From 666f86210e0e03f1207cf94a8718d5b0fdb2270d Mon Sep 17 00:00:00 2001 From: Brent Gardner Date: Wed, 10 Jun 2026 06:23:49 -0600 Subject: [PATCH 01/10] Revert custom allocator auditing of MemoryPool tracking in SLTs (#22860) ## Which issue does this PR close? No issue, just responding to discussion in #22723 ## Rationale for this change Agreed this was an overly strict approach. ## What changes are included in this PR? Reverting previous work ## Are these changes tested? By definition ## Are there any user-facing changes? Contributors won't be held to higher standards in SLTs than in the contributor guide. --- .github/workflows/rust.yml | 2 +- datafusion/execution/src/memory_pool/mod.rs | 12 +- datafusion/execution/src/memory_pool/pool.rs | 68 +-- datafusion/execution/src/runtime_env.rs | 60 +-- datafusion/sqllogictest/Cargo.toml | 4 - datafusion/sqllogictest/README.md | 29 -- datafusion/sqllogictest/bin/sqllogictests.rs | 71 +-- datafusion/sqllogictest/src/accounting.rs | 434 ------------------ .../sqllogictest/src/accounting_pool.rs | 174 ------- .../src/engines/datafusion_engine/runner.rs | 46 +- datafusion/sqllogictest/src/lib.rs | 16 +- datafusion/sqllogictest/src/test_context.rs | 39 +- docs/source/contributor-guide/testing.md | 12 - 13 files changed, 18 insertions(+), 949 deletions(-) delete mode 100644 datafusion/sqllogictest/src/accounting.rs delete mode 100644 datafusion/sqllogictest/src/accounting_pool.rs diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 6c74a9539ad5f..770e705dddc90 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -471,7 +471,7 @@ jobs: export RUST_MIN_STACK=20971520 export TPCH_DATA=`realpath datafusion/sqllogictest/test_files/tpch/data` cargo test plan_q --package datafusion-benchmarks --profile ci --features=ci -- --test-threads=1 - INCLUDE_TPCH=true cargo test --features backtrace,parquet_encryption,substrait,memory-accounting --profile ci --package datafusion-sqllogictest --test sqllogictests -- --default-pool-size-mb 16384 + INCLUDE_TPCH=true cargo test --features backtrace,parquet_encryption,substrait --profile ci --package datafusion-sqllogictest --test sqllogictests - name: Verify Working Directory Clean run: git diff --exit-code diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs index e50f72632b3f2..2b36ee7f40add 100644 --- a/datafusion/execution/src/memory_pool/mod.rs +++ b/datafusion/execution/src/memory_pool/mod.rs @@ -18,7 +18,7 @@ //! [`MemoryPool`] for memory management during query execution, [`proxy`] for //! help with allocation accounting. -use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{Result, internal_datafusion_err}; use std::any::Any; use std::fmt::Display; use std::hash::{Hash, Hasher}; @@ -223,16 +223,6 @@ pub trait MemoryPool: Any + Send + Sync + std::fmt::Debug + Display { fn memory_limit(&self) -> MemoryLimit { MemoryLimit::Unknown } - - /// Attempt to update this pool's limit in place to `new_limit` bytes. - /// - /// Default impl returns `Err`. Callers that route through - /// [`crate::runtime_env::RuntimeEnvBuilder::with_memory_limit`] fall - /// back to replacing the pool wholesale on `Err`, preserving historical - /// behavior for pools that can't be resized in place. - fn try_resize(&self, _new_limit: usize) -> Result<()> { - not_impl_err!("{} does not support resize", self.name()) - } } impl dyn MemoryPool { diff --git a/datafusion/execution/src/memory_pool/pool.rs b/datafusion/execution/src/memory_pool/pool.rs index ecbc2bd5c6f82..52b601d5cd78b 100644 --- a/datafusion/execution/src/memory_pool/pool.rs +++ b/datafusion/execution/src/memory_pool/pool.rs @@ -73,15 +73,9 @@ impl Display for UnboundedMemoryPool { /// This pool works well for queries that do not need to spill or have /// a single spillable operator. See [`FairSpillPool`] if there are /// multiple spillable operators that all will spill. -/// -/// Supports [`MemoryPool::try_resize`] for in-place limit adjustment, so -/// callers routing through -/// [`RuntimeEnvBuilder::with_memory_limit`](crate::runtime_env::RuntimeEnvBuilder::with_memory_limit) -/// can keep the existing pool (and any wrappers around it) rather than -/// replacing it on every change. #[derive(Debug)] pub struct GreedyMemoryPool { - pool_size: AtomicUsize, + pool_size: usize, used: AtomicUsize, } @@ -90,7 +84,7 @@ impl GreedyMemoryPool { pub fn new(pool_size: usize) -> Self { debug!("Created new GreedyMemoryPool(pool_size={pool_size})"); Self { - pool_size: AtomicUsize::new(pool_size), + pool_size, used: AtomicUsize::new(0), } } @@ -110,17 +104,16 @@ impl MemoryPool for GreedyMemoryPool { } fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { - let pool_size = self.pool_size.load(Ordering::Relaxed); self.used .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |used| { let new_used = used + additional; - (new_used <= pool_size).then_some(new_used) + (new_used <= self.pool_size).then_some(new_used) }) .map_err(|used| { insufficient_capacity_err( reservation, additional, - pool_size.saturating_sub(used), + self.pool_size.saturating_sub(used), self, ) })?; @@ -132,25 +125,19 @@ impl MemoryPool for GreedyMemoryPool { } fn memory_limit(&self) -> MemoryLimit { - MemoryLimit::Finite(self.pool_size.load(Ordering::Relaxed)) - } - - fn try_resize(&self, new_limit: usize) -> Result<()> { - self.pool_size.store(new_limit, Ordering::Relaxed); - Ok(()) + MemoryLimit::Finite(self.pool_size) } } impl Display for GreedyMemoryPool { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { let used = self.used.load(Ordering::Relaxed); - let pool_size = self.pool_size.load(Ordering::Relaxed); write!( f, "{}(used: {}, pool_size: {})", &self.name(), human_readable_size(used), - human_readable_size(pool_size) + human_readable_size(self.pool_size) ) } } @@ -613,10 +600,6 @@ impl MemoryPool for TrackConsumersPool { fn memory_limit(&self) -> MemoryLimit { self.inner.memory_limit() } - - fn try_resize(&self, new_limit: usize) -> Result<()> { - self.inner.try_resize(new_limit) - } } fn provide_top_memory_consumers_to_error_msg( @@ -1063,43 +1046,4 @@ mod tests { "TrackConsumersPool Display" ); } - - #[test] - fn test_greedy_try_resize_in_place() { - let pool: Arc = Arc::new(GreedyMemoryPool::new(100)); - let r = MemoryConsumer::new("r").register(&pool); - - // Fill the pool, then verify it rejects further growth. - r.try_grow(100).unwrap(); - r.try_grow(1).unwrap_err(); - - // Resize *up*: previously-rejected growth now succeeds. - pool.try_resize(200).unwrap(); - assert!(matches!(pool.memory_limit(), MemoryLimit::Finite(200))); - r.try_grow(50).unwrap(); - assert_eq!(pool.reserved(), 150); - - // Resize *down* below current usage: subsequent grows fail because - // reserved (150) already exceeds the new limit (120). Already-issued - // reservations are not retroactively shrunk. - pool.try_resize(120).unwrap(); - assert!(matches!(pool.memory_limit(), MemoryLimit::Finite(120))); - r.try_grow(1).unwrap_err(); - } - - #[test] - fn test_track_consumers_try_resize_forwards() { - let pool: Arc = Arc::new(TrackConsumersPool::new( - GreedyMemoryPool::new(100), - NonZeroUsize::new(3).unwrap(), - )); - pool.try_resize(500).unwrap(); - assert!(matches!(pool.memory_limit(), MemoryLimit::Finite(500))); - } - - #[test] - fn test_unbounded_try_resize_returns_err() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - assert!(pool.try_resize(100).is_err()); - } } diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index 31f663e19557b..5b90f28a141ef 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -409,23 +409,12 @@ impl RuntimeEnvBuilder { /// Specify the total memory to use while running the DataFusion /// plan to `max_memory * memory_fraction` in bytes. /// - /// If a memory pool is already configured on this builder, this first - /// attempts to resize it in place via [`MemoryPool::try_resize`]. Pools - /// that support resize (e.g. [`GreedyMemoryPool`]) keep their identity - /// — useful for any wrapper that needs to observe limit changes (e.g. - /// to retune external accounting). Pools whose [`MemoryPool::try_resize`] - /// returns `Err` (the default) fall back to wholesale replacement - /// with a [`TrackConsumersPool`]-wrapped [`GreedyMemoryPool`] (top 5 - /// consumers), preserving the historical behavior. + /// This defaults to using [`GreedyMemoryPool`] wrapped in the + /// [`TrackConsumersPool`] with a maximum of 5 consumers. /// /// Note DataFusion does not yet respect this limit in all cases. pub fn with_memory_limit(self, max_memory: usize, memory_fraction: f64) -> Self { let pool_size = (max_memory as f64 * memory_fraction) as usize; - if let Some(existing) = &self.memory_pool - && existing.try_resize(pool_size).is_ok() - { - return self; - } self.with_memory_pool(Arc::new(TrackConsumersPool::new( GreedyMemoryPool::new(pool_size), NonZeroUsize::new(5).unwrap(), @@ -573,48 +562,3 @@ impl RuntimeEnvBuilder { docs } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::memory_pool::{GreedyMemoryPool, MemoryLimit, UnboundedMemoryPool}; - - #[test] - fn with_memory_limit_resizes_in_place_when_pool_supports_it() { - let pool: Arc = Arc::new(GreedyMemoryPool::new(100)); - let pool_ptr = Arc::as_ptr(&pool); - - let env = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&pool)) - .with_memory_limit(500, 1.0) - .build() - .unwrap(); - - // Same Arc as before — wrapper-or-other-resize-capable pools survive. - assert!(std::ptr::eq(Arc::as_ptr(&env.memory_pool), pool_ptr)); - assert!(matches!( - env.memory_pool.memory_limit(), - MemoryLimit::Finite(500) - )); - } - - #[test] - fn with_memory_limit_falls_back_to_replace_when_resize_unsupported() { - let pool: Arc = Arc::new(UnboundedMemoryPool::default()); - let pool_ptr = Arc::as_ptr(&pool); - - let env = RuntimeEnvBuilder::new() - .with_memory_pool(Arc::clone(&pool)) - .with_memory_limit(500, 1.0) - .build() - .unwrap(); - - // Different Arc — wholesale replacement happened because Unbounded's - // default `try_resize` returns Err. - assert!(!std::ptr::eq(Arc::as_ptr(&env.memory_pool), pool_ptr)); - assert!(matches!( - env.memory_pool.memory_limit(), - MemoryLimit::Finite(500) - )); - } -} diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index cda73ba4e8766..a642fbe22a6e3 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -70,10 +70,6 @@ tokio-postgres = { version = "0.7.17", optional = true } [features] avro = ["datafusion/avro"] backtrace = ["datafusion/backtrace"] -# Enable the `AccountingAllocator` `GlobalAlloc` wrapper and its thread-local -# byte counter. The binary still has to declare `#[global_allocator]` for it -# to actually take effect — building with this feature on alone is harmless. -memory-accounting = [] postgres = [ "bytes", "chrono", diff --git a/datafusion/sqllogictest/README.md b/datafusion/sqllogictest/README.md index 57aabca361553..f0a54cf978fbf 100644 --- a/datafusion/sqllogictest/README.md +++ b/datafusion/sqllogictest/README.md @@ -360,35 +360,6 @@ For focusing on one specific failing test, a file:line filter can be used: cargo test --test sqllogictests -- --substrait-round-trip binary.slt:23 ``` -## Running tests: allocator-level memory accounting - -Build with `--features memory-accounting` to install a global allocator -wrapper that tracks actual bytes allocated per SLT file and reconciles them -against DataFusion's voluntary `MemoryPool` tracking. The point isn't to -enforce a process-wide budget — it's to catch DataFusion lying about how -much memory it's using. If `MemoryPool` reports 1 MB while the allocator -sees 100 MB go by, _that gap is the bug_. - -```shell -cargo test --features memory-accounting --test sqllogictests -- \ - --default-pool-size-mb 16384 -``` - -`--default-pool-size-mb` seeds each per-file SLT context's MemoryPool with -the given size in MB and arms the bank as a no-op until a test opts in. - -**Opting an individual test in.** Add `SET datafusion.runtime.memory_limit = 'N'` at the top of the `.slt`. The wrapping `AccountingMemoryPool` then -tightens its allocator-level bank to `N * 1.10` (10% headroom). If the test -allocates more than that — including bytes DataFusion's tracker didn't see -— the test panics with an `OverdraftPanic` reporting the actual balance at -panic time. SLTs without a `SET` of `memory_limit` see no change in -behavior; the bank stays loose and `SHOW ALL` continues to render the limit -as `unlimited`. - -Inside the runner each file gets its own multi-thread Tokio runtime so -context-ids stamped onto worker threads stay stable for the allocator -hook, and per-file accounts in the bank are isolated from each other. - ## `.slt` file format [`sqllogictest`] was originally written for SQLite to verify the diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs index 2b08769bf5208..e43f03fcf46a7 100644 --- a/datafusion/sqllogictest/bin/sqllogictests.rs +++ b/datafusion/sqllogictest/bin/sqllogictests.rs @@ -15,11 +15,6 @@ // specific language governing permissions and limitations // under the License. -#[cfg(feature = "memory-accounting")] -#[global_allocator] -static GLOBAL: datafusion_sqllogictest::AccountingAllocator = - datafusion_sqllogictest::AccountingAllocator::system(); - use clap::{ColorChoice, Parser}; use datafusion::common::instant::Instant; use datafusion::common::utils::get_available_parallelism; @@ -143,19 +138,6 @@ async fn run_tests() -> Result<()> { options.warn_on_ignored(); - #[cfg(feature = "memory-accounting")] - if let Some(pool_mb) = options.default_pool_size_mb { - let pool_bytes = pool_mb.saturating_mul(1024 * 1024); - // Same value drives the inner MemoryPool's size and the bank's - // default budget. The wrapper renders this value as `unlimited` in - // `SHOW ALL` (sentinel for "no SET has happened"); once a test - // calls `SET datafusion.runtime.memory_limit`, the wrapper retunes - // the bank to that limit + 10% headroom. - datafusion_sqllogictest::set_memory_tracker_limit(pool_bytes); - datafusion_sqllogictest::set_default_budget(pool_bytes as isize); - log::info!("memory-accounting on: default pool size = {pool_mb} MB"); - } - // Print parallelism info for debugging CI performance eprintln!( "Running with {} test threads (available parallelism: {})", @@ -228,7 +210,7 @@ async fn run_tests() -> Result<()> { let currently_running_sql_tracker_clone = currently_running_sql_tracker.clone(); let file_start = Instant::now(); - let body = async move { + SpawnedTask::spawn(async move { let result = match ( options.postgres_runner, options.complete, @@ -301,41 +283,9 @@ async fn run_tests() -> Result<()> { } (result, elapsed) - }; - // Each file gets its own multi-thread runtime so a stable per-file - // context-id (stamped via `on_thread_start`) is readable from the - // global allocator hook. Bank accounting and SET-driven limit - // retuning will key off this id in later steps. The outer - // orchestration runtime hosts this via `spawn_blocking` so its - // worker threads aren't blocked by the per-file `block_on`. - // - // Worker count matches `SLT_TARGET_PARTITIONS` so a query's - // partition streams each get a worker rather than contending. - #[cfg(feature = "memory-accounting")] - let spawned = { - let context_id = datafusion_sqllogictest::next_context_id(); - SpawnedTask::spawn_blocking(move || { - // Stamp this thread too — `block_on` polls `body` here, so - // statements that don't suspend (e.g. `SET memory_limit`, - // pool construction) run on this thread, not a worker. - datafusion_sqllogictest::set_thread_context_id(context_id); - let runtime = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .worker_threads(datafusion_sqllogictest::SLT_TARGET_PARTITIONS) - .thread_name(format!("slt-file-{context_id}")) - .on_thread_start(move || { - datafusion_sqllogictest::set_thread_context_id(context_id); - }) - .build() - .expect("build per-file Tokio runtime"); - let out = runtime.block_on(body); - runtime.shutdown_background(); - out - }) - }; - #[cfg(not(feature = "memory-accounting"))] - let spawned = SpawnedTask::spawn(body); - spawned.join().map(move |result| { + }) + .join() + .map(move |result| { let elapsed = match &result { Ok((_, elapsed)) => *elapsed, Err(_) => Duration::ZERO, @@ -965,19 +915,6 @@ struct Options { default_value_t = ColorChoice::Auto )] color: ColorChoice, - - #[clap( - long, - help = "Default MemoryPool size in MB for each per-file SLT context. \ - The pool is wrapped in AccountingMemoryPool, which doubles \ - this value as the 'no SET has happened yet' sentinel — until \ - an SLT calls `SET datafusion.runtime.memory_limit`, SHOW ALL \ - renders the limit as 'unlimited' and the allocator bank \ - stays loose. Once a test SETs a limit, the bank tightens to \ - that limit + 10% headroom. Requires the memory-accounting \ - feature; ignored without it." - )] - default_pool_size_mb: Option, } impl Options { diff --git a/datafusion/sqllogictest/src/accounting.rs b/datafusion/sqllogictest/src/accounting.rs deleted file mode 100644 index 46b6120c24d28..0000000000000 --- a/datafusion/sqllogictest/src/accounting.rs +++ /dev/null @@ -1,434 +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. - -//! Allocator-driven memory accounting with per-context budgets. -//! -//! The bank ([`ACCOUNTS`]) holds one [`AtomicIsize`] account per stamped -//! `CONTEXT_ID`, each tracking its own remaining budget. Allocations debit -//! the current thread's account, deallocations credit it; below zero is an -//! overdraft. Threads with `CONTEXT_ID == 0` (main, the outer orchestration -//! runtime, blocking-pool hosts) are untracked and skip the hot path. -//! -//! Per-alloc bookkeeping accumulates in a thread-local `LOCAL_BALANCE` -//! drift counter; it settles into the account once `|drift|` crosses -//! [`SETTLE_THRESHOLD`] (64 KB), amortizing the `RwLock` read + atomic -//! op across thousands of allocations. -//! -//! [`account_balance`] reads the current thread's account; it lags reality -//! by up to one threshold's worth of un-settled drift per thread. -//! -//! # Enforcement -//! -//! An allocation that drives the bank negative on a stamped thread -//! (`CONTEXT_ID != 0`) panics with [`OverdraftPanic`] on the polling thread. -//! Drop-chain credits during unwind never re-panic — `track` only fires on -//! debits (`delta < 0`). Unstamped threads are silently skipped. -//! -//! Compiled in only when the `memory-accounting` feature is on. - -use std::alloc::{GlobalAlloc, Layout, System}; -use std::cell::Cell; -use std::collections::HashMap; -use std::sync::atomic::{AtomicIsize, AtomicUsize, Ordering}; -use std::sync::{OnceLock, RwLock}; - -/// Net byte change at which a thread flushes its local count into the bank. -/// 64 KB chosen to keep per-thread drift tight (≤1 MB on a 16-core box) while -/// still settling rarely enough to make the bank's atomic op amortized-free. -const SETTLE_THRESHOLD: isize = 64 * 1024; - -/// The bank: every account, keyed by context-id, valued by remaining budget. -/// Debits on alloc, credits on free, negative = overdraft. ctx-id 0 never -/// gets an entry — that's the "untracked thread" marker. -static ACCOUNTS: OnceLock>> = OnceLock::new(); - -/// Starting budget for any new account, set by [`set_default_budget`] and -/// inherited by per-file SLT contexts spawned after. -static DEFAULT_BUDGET: AtomicIsize = AtomicIsize::new(0); - -fn accounts() -> &'static RwLock> { - ACCOUNTS.get_or_init(|| RwLock::new(HashMap::new())) -} - -/// Run `f` against the current thread's account balance, or return `None` -/// if there isn't one — silently skipping the update is fine on the alloc -/// hot path. -fn with_current_balance(op: impl FnOnce(&AtomicIsize) -> R) -> Option { - let ctx_id = CONTEXT_ID.with(|ctx| ctx.get()); - if ctx_id == 0 { - return None; - } - // PERF: acquires an `RwLock` read on every settle. If it ever shows up - // hot, stash a `&'static AtomicIsize` in a thread-local (set in - // `set_thread_context_id`, backed by `Box::leak`) and skip the lookup. - let accounts_lock = ACCOUNTS.get()?; - let accounts = accounts_lock.read().ok()?; - accounts.get(&ctx_id).map(op) -} - -thread_local! { - static LOCAL_BALANCE: Cell = const { Cell::new(0) }; - - /// Account-id stamped onto worker threads via [`set_thread_context_id`]. - /// Zero = untracked thread; nothing to track, nothing to enforce. - static CONTEXT_ID: Cell = const { Cell::new(0) }; -} - -/// Monotonic source of fresh context-ids. Starts at 1; the zero value is -/// reserved for "no per-file runtime" so callers can distinguish. -static CONTEXT_ID_COUNTER: AtomicUsize = AtomicUsize::new(0); - -/// Returns a fresh, never-before-used context-id. Call once per file in the -/// SLT binary and pass the result into the per-file runtime's -/// `on_thread_start` callback so every worker thread of that runtime shares -/// the same id. -pub fn next_context_id() -> usize { - CONTEXT_ID_COUNTER.fetch_add(1, Ordering::Relaxed) + 1 -} - -/// Stamp the current thread with `id`. Intended for `on_thread_start`. -/// Creates the account if it doesn't already exist. -pub fn set_thread_context_id(id: usize) { - if id == 0 { - CONTEXT_ID.with(|ctx| ctx.set(0)); - return; - } - // Insert under the write lock *before* stamping the thread. A HashMap - // resize allocates → recurses through `track` → `with_current_account`, - // which sees `CONTEXT_ID == 0` and bails out instead of trying to - // read-lock the map we're holding for writing on the same thread. - { - let accounts_lock = accounts(); - let mut accounts = accounts_lock - .write() - .unwrap_or_else(|poison| poison.into_inner()); - accounts - .entry(id) - .or_insert_with(|| AtomicIsize::new(DEFAULT_BUDGET.load(Ordering::Relaxed))); - } - CONTEXT_ID.with(|ctx| ctx.set(id)); -} - -/// Current thread's context-id, or 0 if none has been set. -pub fn current_context_id() -> usize { - CONTEXT_ID.with(|ctx| ctx.get()) -} - -/// Payload attached to allocator-induced panics. Catch with: -/// -/// ```ignore -/// match std::panic::catch_unwind(|| { /* ... */ }) { -/// Err(e) if e.is::() => { /* it was an overdraft */ } -/// ... -/// } -/// ``` -#[derive(Debug, Clone)] -pub struct OverdraftPanic { - /// Account balance at the moment the panic fired (negative — that's the point). - pub account_balance: isize, -} - -/// Set the default budget new accounts will be created with. Existing -/// accounts are untouched. -pub fn set_default_budget(value: isize) { - DEFAULT_BUDGET.store(value, Ordering::Relaxed); -} - -/// Current default budget — what a fresh account starts at and what -/// [`reset_account_to_default`] restores to. -pub fn default_budget() -> isize { - DEFAULT_BUDGET.load(Ordering::Relaxed) -} - -/// Restore the current thread's account to [`default_budget`]. Used by the -/// SLT runner after catching an [`OverdraftPanic`] so the next statement -/// starts clean — otherwise the bank stays negative and every subsequent -/// allocation refires, which is unsafe (allocator hooks must not panic -/// repeatedly within a single thread). -pub fn reset_account_to_default() { - set_account_balance(default_budget()); -} - -/// Set the current thread's account balance to `value`. No-op on untracked -/// threads (`CONTEXT_ID == 0`). -pub fn set_account_balance(value: isize) { - let _ = with_current_balance(|bal| bal.store(value, Ordering::Relaxed)); -} - -/// Cross-module config for DataFusion's voluntary `MemoryPool` limit, set -/// from the SLT binary's CLI and read by test_context when building each -/// per-file `RuntimeEnv`. Zero means "use the default `UnboundedMemoryPool`". -static MEMORY_TRACKER_LIMIT: AtomicUsize = AtomicUsize::new(0); - -/// Set the size (in bytes) the per-file `MemoryPool` should be built with. -/// Zero (the default) leaves the existing `UnboundedMemoryPool` behavior. -pub fn set_memory_tracker_limit(bytes: usize) { - MEMORY_TRACKER_LIMIT.store(bytes, Ordering::Relaxed); -} - -/// Current `MemoryPool` limit configured via [`set_memory_tracker_limit`]. -pub fn memory_tracker_limit() -> usize { - MEMORY_TRACKER_LIMIT.load(Ordering::Relaxed) -} - -/// Current account balance. Negative = overdraft. `0` if untracked. -pub fn account_balance() -> isize { - with_current_balance(|bal| bal.load(Ordering::Relaxed)).unwrap_or(0) -} - -/// Current thread's local balance — not yet reflected in the global bank. -/// Always in `(-SETTLE_THRESHOLD, +SETTLE_THRESHOLD)`. Sign matches the bank: -/// negative on a thread that's net-allocated, positive on one that's net-freed. -pub fn local_balance() -> isize { - LOCAL_BALANCE.with(|loc_bal| loc_bal.get()) -} - -/// Force the current thread to flush its local count into its context bank. -/// No-op on untracked threads (`CONTEXT_ID == 0`). -pub fn settle_thread_local() { - if CONTEXT_ID.with(|ctx| ctx.get()) == 0 { - return; - } - let _ = LOCAL_BALANCE.try_with(|loc_bal| { - let drift = loc_bal.replace(0); - if drift != 0 { - let _ = with_current_balance(|bal| bal.fetch_add(drift, Ordering::Relaxed)); - } - }); -} - -/// Record a delta into the current thread's account: settle local drift into -/// the bank when it crosses `±SETTLE_THRESHOLD`, fire the kill panic on a -/// debit that leaves the account negative. -#[inline(always)] -fn track(delta: isize) { - if CONTEXT_ID.with(|ctx| ctx.get()) == 0 { - return; - } - let _ = LOCAL_BALANCE.try_with(|loc_bal| { - let drift = loc_bal.get() + delta; - // 99% case: drift fits — accumulate locally and bail. - if -SETTLE_THRESHOLD < drift && drift < SETTLE_THRESHOLD { - loc_bal.set(drift); - return; - } - // Drop the read lock *before* maybe_kill — the panic allocates, - // recurses through track, and would self-deadlock on std::sync::RwLock. - let new_bal = with_current_balance(|bal| { - bal.fetch_add(drift, Ordering::Relaxed).wrapping_add(drift) - }); - loc_bal.set(0); - // Only debits fire the kill — credits run inside Drop chains during - // unwinding, where a panic would double-fault and abort the process. - if delta >= 0 { - return; - } - let Some(new_bal) = new_bal else { return }; - if new_bal >= 0 { - return; - } - // Skip if we're already unwinding — `panic_any` boxes the payload, - // which allocates, which re-enters `track`; without this gate the - // second debit would fire a nested panic and abort the process. - if std::thread::panicking() { - return; - } - std::panic::panic_any(OverdraftPanic { - account_balance: new_bal, - }); - }); -} - -/// `GlobalAlloc` wrapper that counts bytes against a thread-local + global bank. -/// -/// Forwards every operation unchanged to the inner allocator; the bookkeeping -/// is a thread-local update on the fast path plus an amortized atomic settle. -pub struct AccountingAllocator { - inner: A, -} - -impl AccountingAllocator { - pub const fn new(inner: A) -> Self { - Self { inner } - } -} - -impl AccountingAllocator { - /// Convenience constructor for the typical `System`-backed case. - pub const fn system() -> Self { - Self { inner: System } - } -} - -unsafe impl GlobalAlloc for AccountingAllocator { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - // Account BEFORE the inner alloc. If we panicked AFTER `inner.alloc` - // succeeded, the bytes are physically allocated but no caller ever - // sees the pointer → unwind leaks the very bytes that pushed us - // over the budget — the opposite of what the kill panic is for. - let delta = -(layout.size() as isize); - track(delta); - // SAFETY: layout is forwarded unchanged. - let ptr = unsafe { self.inner.alloc(layout) }; - if ptr.is_null() { - // Allocator refused — refund so the bank matches reality. - track(-delta); - } - ptr - } - - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged. - unsafe { self.inner.dealloc(ptr, layout) }; - // Credit only; `track()` short-circuits on `delta >= 0` and never - // panics, so ordering relative to `inner.dealloc` doesn't matter. - track(layout.size() as isize); - } - - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - // Same panic-then-leak hazard as `alloc`; account first. - let delta = -(layout.size() as isize); - track(delta); - // SAFETY: layout is forwarded unchanged. - let ptr = unsafe { self.inner.alloc_zeroed(layout) }; - if ptr.is_null() { - track(-delta); - } - ptr - } - - unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - // Account BEFORE the inner realloc so a kill panic doesn't strand the - // caller with a freed `ptr`. `inner.realloc` frees `ptr` on success; - // if we panicked after that, the caller's `Vec`-or-similar would - // still hold the old pointer and double-free on unwind (glibc - // "double free or corruption (out)" + SIGABRT). - let delta = layout.size() as isize - new_size as isize; - track(delta); - // SAFETY: caller upholds GlobalAlloc invariants; we forward unchanged. - let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) }; - if new_ptr.is_null() { - // Allocator refused — refund so the bank matches reality. - track(-delta); - } - new_ptr - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[global_allocator] - static GLOBAL: AccountingAllocator = AccountingAllocator::system(); - - /// Each test runs on its own thread (cargo-test parallelism) and stamps a - /// fresh context-id, so per-context isolation makes them naturally - /// independent — no shared mutex required. - fn enter_fresh_context() { - set_thread_context_id(next_context_id()); - } - - #[test] - fn alloc_debits_and_free_credits_account() { - enter_fresh_context(); - // Bump budget well above the alloc + this thread's own background - // drift so the test's own activity can't accidentally overdraw. - set_account_balance(10_000_000); - settle_thread_local(); - let before = account_balance(); - - let buf: Vec = vec![0u8; 8192]; - settle_thread_local(); - let mid = account_balance(); - // Alloc debited the account → mid should be at least 8192 below before. - assert!( - before - mid >= 8192, - "alloc didn't debit: before={before} mid={mid}" - ); - - drop(buf); - settle_thread_local(); - let after = account_balance(); - // Free credited the account → after should be at least 8192 above mid. - assert!( - after - mid >= 8192, - "free didn't credit: mid={mid} after={after}" - ); - } - - #[test] - fn set_account_balance_sticks() { - enter_fresh_context(); - set_account_balance(1_000_000); - // Balance drifts a little from this thread's own allocator activity - // between the set and the read, so we expect at-or-below the set value. - let bal = account_balance(); - assert!( - (900_000..=1_000_000).contains(&bal), - "set_account_balance didn't stick: bal={bal}" - ); - } - - #[test] - fn overdraft_on_stamped_thread_panics() { - use std::panic::{AssertUnwindSafe, catch_unwind}; - enter_fresh_context(); - set_account_balance(1024); - - let result = catch_unwind(AssertUnwindSafe(|| { - // Alloc large enough to cross SETTLE_THRESHOLD in one shot — the - // settle drives the bank negative on a stamped thread, which now - // unconditionally panics. - let _buf: Vec = vec![0u8; SETTLE_THRESHOLD as usize + 4096]; - unreachable!("alloc should have panicked"); - })); - - let payload = result.expect_err("alloc should have panicked"); - let overdraft = payload - .downcast_ref::() - .expect("panic payload should be OverdraftPanic"); - assert!( - overdraft.account_balance < 0, - "payload should report negative balance; got {}", - overdraft.account_balance - ); - } - - #[test] - fn threshold_settlement_flushes_to_account() { - enter_fresh_context(); - // Bump budget — the settle on threshold crossing now panics on - // a stamped thread if it goes negative. We just want to observe the - // flush mechanism here, not the kill. - set_account_balance(10_000_000); - settle_thread_local(); - let before = account_balance(); - - let buf: Vec = vec![0u8; SETTLE_THRESHOLD as usize + 1024]; - // Crossing the threshold auto-settles; account balance should have - // dropped by at least SETTLE_THRESHOLD without us calling - // settle_thread_local. - let after_alloc = account_balance(); - assert!( - before - after_alloc >= SETTLE_THRESHOLD, - "balance didn't auto-settle on threshold crossing: \ - before={before} after_alloc={after_alloc}" - ); - drop(buf); - } -} diff --git a/datafusion/sqllogictest/src/accounting_pool.rs b/datafusion/sqllogictest/src/accounting_pool.rs deleted file mode 100644 index a9d2db9f12261..0000000000000 --- a/datafusion/sqllogictest/src/accounting_pool.rs +++ /dev/null @@ -1,174 +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. - -//! [`AccountingMemoryPool`] bridges DataFusion's voluntary memory tracking -//! to the allocator-level bank in [`crate::accounting`]. -//! -//! It wraps any [`MemoryPool`] and re-tunes the current thread's bank -//! account whenever the pool's limit changes (via [`MemoryPool::try_resize`], -//! which `RuntimeEnvBuilder::with_memory_limit` triggers on `SET -//! datafusion.runtime.memory_limit = '…'`). -//! -//! Each retune sets the bank to `new_limit * HEADROOM_FACTOR`. A query -//! that allocates past that envelope panics with an `OverdraftPanic` — -//! the gap between DF's voluntary tracker and the allocator's reality -//! is the bug we're hunting. - -use crate::set_account_balance; -use datafusion::common::Result; -use datafusion::execution::memory_pool::{ - MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation, -}; -use std::fmt::{self, Display, Formatter}; -use std::sync::Arc; - -/// Headroom over the pool's declared limit. Anything past this is an -/// untracked allocation — by definition, since DF's pool didn't see it. -/// -/// 800% high, but that's what it takes to pass the SLT suite right now. Goal should be ~10% -const HEADROOM_FACTOR: f64 = 8.0; - -pub struct AccountingMemoryPool { - inner: Arc, - /// The operator-configured default pool size, used as a "no SET has - /// happened yet" sentinel by [`Self::memory_limit`]. - default_size: usize, -} - -impl AccountingMemoryPool { - pub fn new(inner: Arc, default_size: usize) -> Self { - Self { - inner, - default_size, - } - } -} - -impl fmt::Debug for AccountingMemoryPool { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - f.debug_struct("AccountingMemoryPool") - .field("inner", &self.inner) - .field("default_size", &self.default_size) - .finish() - } -} - -impl Display for AccountingMemoryPool { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "accounting({})", self.inner) - } -} - -impl MemoryPool for AccountingMemoryPool { - fn name(&self) -> &str { - "accounting" - } - - 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) - } - - fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { - self.inner.shrink(reservation, shrink) - } - - fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { - self.inner.try_grow(reservation, additional) - } - - fn reserved(&self) -> usize { - self.inner.reserved() - } - - fn memory_limit(&self) -> MemoryLimit { - // HACK: When the inner pool still reports the operator-configured - // default, no `SET datafusion.runtime.memory_limit` has happened — - // render as `Infinite` so `information_schema.slt`'s `SHOW ALL` - // expectation of `unlimited` for an un-SET context stays satisfied. - // Once a SET fires, `try_resize` mutates the inner pool to some - // other value and we report the real limit. - match self.inner.memory_limit() { - MemoryLimit::Finite(n) if n == self.default_size => MemoryLimit::Infinite, - other => other, - } - } - - fn try_resize(&self, new_limit: usize) -> Result<()> { - self.inner.try_resize(new_limit)?; - set_account_balance((new_limit as f64 * HEADROOM_FACTOR) as isize); - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::{account_balance, next_context_id, set_thread_context_id}; - use datafusion::execution::memory_pool::GreedyMemoryPool; - - #[test] - fn memory_limit_returns_infinite_for_sentinel() { - let default_size = 1_000_000; - let pool = AccountingMemoryPool::new( - Arc::new(GreedyMemoryPool::new(default_size)), - default_size, - ); - assert!(matches!(pool.memory_limit(), MemoryLimit::Infinite)); - } - - #[test] - fn memory_limit_returns_finite_after_resize() { - let default_size = 1_000_000; - let pool = AccountingMemoryPool::new( - Arc::new(GreedyMemoryPool::new(default_size)), - default_size, - ); - pool.try_resize(50_000).unwrap(); - assert!(matches!(pool.memory_limit(), MemoryLimit::Finite(50_000))); - } - - #[test] - fn try_resize_retunes_current_account_balance() { - // Stamp a fresh context so set_account_balance lands somewhere - // visible. Otherwise CONTEXT_ID == 0 means the call is a no-op. - set_thread_context_id(next_context_id()); - - let default_size = 1_000_000; - let pool = AccountingMemoryPool::new( - Arc::new(GreedyMemoryPool::new(default_size)), - default_size, - ); - pool.try_resize(50_000).unwrap(); - - // Balance is reset to limit * HEADROOM_FACTOR, minus a small - // drift from this test thread's own allocs between set and read. - let expected = (50_000.0 * HEADROOM_FACTOR) as isize; - let bal = account_balance(); - assert!( - (50_000..=expected).contains(&bal), - "balance not in expected range: got {bal}, expected ≤ {expected}" - ); - } -} diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs index 0c038fb00fa08..08facc48005dc 100644 --- a/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs +++ b/datafusion/sqllogictest/src/engines/datafusion_engine/runner.rs @@ -83,50 +83,6 @@ impl DataFusion { self } - /// Run a single query through the engine. Under the `memory-accounting` - /// feature, allocator-detected overdrafts panic with `OverdraftPanic`; - /// catch them here and translate to a clean `Err`. - async fn run_one(&self, sql: &str) -> Result { - #[cfg(feature = "memory-accounting")] - { - use crate::OverdraftPanic; - use futures::FutureExt; - - let fut = run_query(&self.ctx, is_spark_path(&self.relative_path), sql); - - return match std::panic::AssertUnwindSafe(fut).catch_unwind().await { - Ok(r) => r, - Err(payload) => { - if let Some(od) = payload.downcast_ref::() { - let df_reserved_mb = - (self.ctx.runtime_env().memory_pool.reserved() as u64) - / (1024 * 1024); - warn!( - "[{}] killed by allocator overdraft: \ - account balance = {} bytes, df-pool reserved = {df_reserved_mb} MB; \ - sql = {sql:?}", - self.relative_path.display(), - od.account_balance, - ); - // Restore the bank so the next statement starts clean - crate::reset_account_to_default(); - Err(DFSqlLogicTestError::Other(format!( - "allocator overdraft: account balance at panic = {} bytes", - od.account_balance, - ))) - } else { - // Not our panic — re-raise so test runner sees it. - std::panic::resume_unwind(payload); - } - } - }; - } - #[cfg(not(feature = "memory-accounting"))] - { - run_query(&self.ctx, is_spark_path(&self.relative_path), sql).await - } - } - fn update_slow_count(&self) { let msg = self.pb.message(); let split: Vec<&str> = msg.split(" ").collect(); @@ -198,7 +154,7 @@ impl sqllogictest::AsyncDB for DataFusion { let tracked_sql = self.currently_executing_sql_tracker.set_sql(sql); let start = Instant::now(); - let result = self.run_one(sql).await; + let result = run_query(&self.ctx, is_spark_path(&self.relative_path), sql).await; let duration = start.elapsed(); self.currently_executing_sql_tracker.remove_sql(tracked_sql); diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs index 54f460958c0ab..6b6c40365f855 100644 --- a/datafusion/sqllogictest/src/lib.rs +++ b/datafusion/sqllogictest/src/lib.rs @@ -26,23 +26,9 @@ //! DataFusion sqllogictest driver -#[cfg(feature = "memory-accounting")] -mod accounting; -#[cfg(feature = "memory-accounting")] -mod accounting_pool; mod engines; mod test_file; -#[cfg(feature = "memory-accounting")] -pub use accounting::{ - AccountingAllocator, OverdraftPanic, account_balance, current_context_id, - default_budget, local_balance, memory_tracker_limit, next_context_id, - reset_account_to_default, set_account_balance, set_default_budget, - set_memory_tracker_limit, set_thread_context_id, settle_thread_local, -}; -#[cfg(feature = "memory-accounting")] -pub use accounting_pool::AccountingMemoryPool; - pub use engines::CurrentlyExecutingSqlTracker; pub use engines::DFColumnType; pub use engines::DFOutput; @@ -61,6 +47,6 @@ mod test_context; mod util; pub use filters::*; -pub use test_context::{SLT_TARGET_PARTITIONS, TestContext}; +pub use test_context::TestContext; pub use test_file::TestFile; pub use util::*; diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 8d437271fee86..e0aaa91ef6369 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -60,20 +60,12 @@ use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; use datafusion::execution::runtime_env::RuntimeEnv; -#[cfg(feature = "memory-accounting")] -use datafusion::execution::runtime_env::RuntimeEnvBuilder; use log::info; use sqlparser::ast; use tempfile::TempDir; mod range_partitioning; -/// Target partition count used for every SLT file's `SessionConfig`. Hardcoded -/// so query plans are deterministic across machines. The SLT binary also -/// sizes each file's per-file Tokio runtime to this value so partition streams -/// each get a worker rather than contending. -pub const SLT_TARGET_PARTITIONS: usize = 4; - /// Context for running tests pub struct TestContext { /// Context for running queries @@ -99,33 +91,6 @@ impl TypePlanner for SqlLogicTestTypePlanner { } } -/// Construct the per-file `RuntimeEnv`. With the `memory-accounting` feature -/// on and a non-zero `memory_tracker_limit()` configured, this wraps the -/// usual `TrackConsumersPool(GreedyMemoryPool)` in an `AccountingMemoryPool` -/// so the allocator-level bank retunes on every `SET datafusion.runtime. -/// memory_limit`. Otherwise falls back to the historical default. -fn build_runtime_env() -> RuntimeEnv { - #[cfg(feature = "memory-accounting")] - { - use datafusion::execution::memory_pool::{GreedyMemoryPool, TrackConsumersPool}; - use std::num::NonZeroUsize; - - let limit = crate::memory_tracker_limit(); - if limit > 0 { - let tracked = TrackConsumersPool::new( - GreedyMemoryPool::new(limit), - NonZeroUsize::new(5).unwrap(), - ); - let wrapped = crate::AccountingMemoryPool::new(Arc::new(tracked), limit); - return RuntimeEnvBuilder::new() - .with_memory_pool(Arc::new(wrapped)) - .build() - .expect("RuntimeEnvBuilder::build with accounting pool"); - } - } - RuntimeEnv::default() -} - impl TestContext { pub fn new(ctx: SessionContext) -> Self { Self { @@ -142,8 +107,8 @@ impl TestContext { pub async fn try_new_for_test_file(relative_path: &Path) -> Option { let config = SessionConfig::new() // hardcode target partitions so plans are deterministic - .with_target_partitions(SLT_TARGET_PARTITIONS); - let runtime = Arc::new(build_runtime_env()); + .with_target_partitions(4); + let runtime = Arc::new(RuntimeEnv::default()); let mut state_builder = SessionStateBuilder::new() .with_config(config) diff --git a/docs/source/contributor-guide/testing.md b/docs/source/contributor-guide/testing.md index 3e44e3aabaeef..3b644f610b90e 100644 --- a/docs/source/contributor-guide/testing.md +++ b/docs/source/contributor-guide/testing.md @@ -113,18 +113,6 @@ Like similar systems such as [DuckDB](https://duckdb.org/dev/testing), DataFusio DataFusion has integrated [sqlite's test suite](https://sqlite.org/sqllogictest/doc/trunk/about.wiki) as a supplemental test suite that is run whenever a PR is merged into DataFusion. To run it manually please refer to the [README](https://github.com/apache/datafusion/blob/main/datafusion/sqllogictest/README.md#running-tests-sqlite) file for instructions. -### Allocator-level memory accounting (`--features memory-accounting`) - -For tests that need to verify DataFusion's voluntary memory tracking -matches actual heap usage, the `sqllogictest` runner ships an optional -`memory-accounting` feature that installs a global allocator wrapper. -Adding `SET datafusion.runtime.memory_limit = 'N'` at the top of an -`.slt` file opts that file into allocator-vs-`MemoryPool` reconciliation -with 10% headroom — any divergence panics the test with an -`OverdraftPanic` reporting the actual allocator balance. See -[the sqllogictest README](https://github.com/apache/datafusion/blob/main/datafusion/sqllogictest/README.md#running-tests-allocator-level-memory-accounting) -for the runner flag and the full mechanism. - ## Snapshot testing (`cargo insta`) [Insta](https://github.com/mitsuhiko/insta) is used for snapshot testing. Snapshots are generated From 3f52debc532331d3a9b07abaf31868c53d77740d Mon Sep 17 00:00:00 2001 From: Raz Luvaton <16746759+rluvaton@users.noreply.github.com> Date: Wed, 10 Jun 2026 17:51:27 +0300 Subject: [PATCH 02/10] fix: add backtrace for `assert_*_or_internal_err` helpers (#18910) ## Which issue does this PR close? N/A ## Rationale for this change No backtrace is added when using the assert macros, so fixing that ## What changes are included in this PR? used `internal_datafusion_err` macro in the `assert_*` helpers ## Are these changes tested? yes ## Are there any user-facing changes? now have backtrace when feature enabled --- datafusion/common/src/error.rs | 412 +++++++++++++++++++++++++++------ 1 file changed, 343 insertions(+), 69 deletions(-) diff --git a/datafusion/common/src/error.rs b/datafusion/common/src/error.rs index 71ae9ec71081d..ce6f8e68aee43 100644 --- a/datafusion/common/src/error.rs +++ b/datafusion/common/src/error.rs @@ -823,10 +823,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) - )) + ) })? }; } @@ -844,19 +844,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)+) - ))); + )); } }; } @@ -876,27 +876,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)+) - ))); + )); } }}; } @@ -916,27 +916,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)+) - ))); + )); } }}; } @@ -1204,7 +1204,6 @@ mod test { use std::sync::Arc; use arrow::error::ArrowError; - use insta::assert_snapshot; fn ok_result() -> Result<()> { Ok(()) @@ -1223,14 +1222,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] @@ -1246,14 +1239,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] @@ -1270,14 +1257,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] @@ -1287,13 +1268,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") ); } @@ -1304,14 +1281,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] @@ -1339,23 +1310,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::>() @@ -1363,15 +1380,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] From 5e3bc99fb0000dfa1e4d0c8e311dbf35080b734b Mon Sep 17 00:00:00 2001 From: EeshanBembi <33062610+EeshanBembi@users.noreply.github.com> Date: Wed, 10 Jun 2026 22:12:56 +0530 Subject: [PATCH 03/10] perf: fast-path inline strings in ByteViewGroupValueBuilder::vectorized_append (#21794) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #21568. ## Rationale for this change `ByteViewGroupValueBuilder::vectorized_append` was doing unnecessary work for short strings (≤12 bytes): for each row it called `array.value(row)` to decode the u128 view into a `&[u8]`, then called `make_view` to re-encode it back into a u128. The input `GenericByteViewArray` already stores inline values in exactly that u128 format, so the round-trip is redundant. This mirrors the existing `HAS_BUFFERS` specialisation in `vectorized_equal_to_inner`, which uses the same `data_buffers().is_empty()` guard to take a direct-view-compare fast path for inline strings. ## What changes are included in this PR? In `vectorized_append_inner`, the `Nulls::None` branch now dispatches on `arr.data_buffers().is_empty()`: - **Fast path** (no data buffers → all values ≤12 bytes inline): copies u128 views directly via `self.views.extend(rows.iter().map(|&row| arr.views()[row]))`. Arrow's validity invariant guarantees inline views are zero-padded, so direct copy is semantically identical to `value() → make_view()`. - **Slow path** (array has non-inline strings): adds `self.views.reserve(rows.len())` before the existing loop to avoid repeated reallocation. ## Are these changes tested? Covered by the existing 6 unit tests in `bytes_view::tests`, all passing unchanged. `test_byte_view_vectorized_operation_special_case` exercises the fast path directly (11-byte strings, no data buffers). ## Are there any user-facing changes? No. Internal performance improvement only. ## Benchmark `inline_null_0.0_size_1000/vectorized_append` (8-byte strings, no nulls, 1 000 rows): | | time | |---|---| | Before | 3.37 µs | | After | 495 ns | | Change | **−85.3% (6.8× faster)** | --- datafusion/physical-plan/Cargo.toml | 1 - .../benches/aggregate_vectorized.rs | 15 ++--- .../group_values/multi_group_by/bytes_view.rs | 56 +++++++++++++++++-- 3 files changed, 59 insertions(+), 13 deletions(-) diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index 0fc75043bf333..4b2b31febef2a 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -123,7 +123,6 @@ required-features = ["test_utils"] [[bench]] harness = false name = "aggregate_vectorized" -required-features = ["test_utils"] [[bench]] harness = false 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/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 e94e4547e1a75..abc3aba88ad48 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 @@ -145,7 +145,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(); @@ -166,8 +170,50 @@ 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 are non-inline (>12 bytes). + // Read views directly to avoid array.value(row) overhead and + // reuse the source view's prefix instead of recomputing it via make_view. + 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 { + let view = arr.views()[row]; + let len = view as u32; + if len <= 12 { + // This row happens to be inline; copy view directly. + self.views.push(view); + } else { + let src = ByteView::from(view); + // ensure_in_progress_big_enough must be called before computing + // new_buffer_index / new_offset — it may flush in_progress to completed. + 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 = &arr.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], + ); + // Reuse prefix from the source view — avoids re-reading first 4 bytes. + 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); + } + } } } @@ -177,6 +223,7 @@ impl ByteViewGroupValueBuilder { self.views.resize(new_len, 0); } } + Ok(()) } fn do_append_val_inner(&mut self, array: &GenericByteViewArray, row: usize) @@ -548,8 +595,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 { From 92820c8047b3dcee5553d057f58381e2942a6852 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 10 Jun 2026 12:48:55 -0400 Subject: [PATCH 04/10] docs: Update/improve `SELECT` reference (#22672) ## Which issue does this PR close? - N/A ## Rationale for this change Various improvements and cleanups to the `SELECT` reference page in the documentation. Update the documentation to cover more of the supported SQL syntax, to include more details when appropriate, and to add more examples. ## What changes are included in this PR? * Improve the `SELECT` reference page in the documentation * Update `pipe_operator.slt` to reflect that the pipe operator syntax is now supported in the default SQL dialect ## Are these changes tested? Manually and LLM-checked the examples. ## Are there any user-facing changes? Documentation improvements, but should be no regressions. No functional changes in DataFusion. --- .../sqllogictest/test_files/pipe_operator.slt | 11 +- docs/source/user-guide/sql/select.md | 556 +++++++++++++++--- 2 files changed, 475 insertions(+), 92 deletions(-) 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/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index 3564884b041ad..ea96f6ae4528d 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -19,76 +19,320 @@ # 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. + +
[ 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 ... ]
+ +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 t.a FROM table AS t +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 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). + ## 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. + +Join conditions can use `ON` or `USING`. + +Examples: + +```sql +SELECT * +FROM orders AS o +JOIN customers AS c ON o.customer_id = c.id; -The following examples are based on this table: +SELECT * +FROM orders +JOIN customers USING (customer_id); +``` + +The join examples below use this table: ```sql select * from x; @@ -104,7 +348,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 +356,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 +383,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 +398,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 +413,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 +427,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 +441,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 +455,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 +466,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 +479,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 +618,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_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; +``` + +`GROUP BY ALL` groups by every non-aggregate expression in the `SELECT` list. ```sql -SELECT a, b, MAX(c) FROM table GROUP BY a, b +SELECT a, b, SUM(c) FROM table_name GROUP BY ALL; ``` -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. +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) 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 a, b, MAX(c) FROM table GROUP BY a, b HAVING MAX(c) > 10 +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. -Example: +`UNION` returns rows from both inputs and removes duplicates by default. +`UNION DISTINCT` is equivalent to `UNION`; `UNION ALL` preserves duplicates. + +Examples: ```sql SELECT @@ -419,55 +758,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) From 6d4cb31d3a844993ec82279d363783780481a53d Mon Sep 17 00:00:00 2001 From: "7. Sun" Date: Wed, 10 Jun 2026 12:49:23 -0400 Subject: [PATCH 05/10] feat: implement Spark-compatible weekday function (#22740) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Which issue does this PR close? - Closes #22599. ## Rationale for this change Implements the Spark `weekday` function as part of the datafusion-spark function library (#15914). Spark's `weekday` returns the day-of-week as a 0-indexed integer with Monday = 0 .. Sunday = 6, which differs from the existing `dayofweek` (1-indexed, Sunday = 1). This fills a gap for Spark-compatibility consumers. ## What changes are included in this PR? - New `SparkWeekDay` scalar UDF in `datafusion/spark/src/function/datetime/weekday.rs`, modeled on the sibling `monthname` function. - Accepts Date and (implicitly coerced) Timestamp input; returns `Int32`. - Uses Arrow's `DatePart::DayOfWeekMonday0` kernel, which is exactly 0=Monday .. 6=Sunday — a direct match for Spark semantics. - Null input propagates to null output for both scalar and array paths. - Registered in the datetime function `mod.rs` (UDF macro, doc export, `functions()` list). ## Are these changes tested? Yes: - Rust unit tests: return-field nullability, scalar evaluation (incl. null), and array evaluation. - sqllogictest coverage in `datafusion/sqllogictest/test_files/spark/datetime/weekday.slt`: scalar dates for all 7 weekdays, array input, TIMESTAMP / TIMESTAMP_NTZ / LTZ coercion, null handling, and argument-type / zero-argument error cases. The original PySpark 3.5.5 reference (`weekday('2009-07-30') = 3`) is preserved and validated. ## Are there any user-facing changes? Adds the new Spark-compatible scalar function `weekday`. No breaking changes. --- Note: I noticed there is an existing draft PR #22601 for the same function (last updated 2026-05-28). It appears to be a stale draft, so I've opened this as a complete, tested implementation — happy to defer or collaborate if the draft author is still active. Signed-off-by: sjhddh <151469562+sjhddh@users.noreply.github.com> Co-authored-by: sjhddh <151469562+sjhddh@users.noreply.github.com> Co-authored-by: Andrew Lamb --- datafusion/spark/src/function/datetime/mod.rs | 8 + .../spark/src/function/datetime/weekday.rs | 191 ++++++++++++++++++ .../test_files/spark/datetime/weekday.slt | 108 +++++++++- 3 files changed, 305 insertions(+), 2 deletions(-) create mode 100644 datafusion/spark/src/function/datetime/weekday.rs diff --git a/datafusion/spark/src/function/datetime/mod.rs b/datafusion/spark/src/function/datetime/mod.rs index 98afa91ddc834..70ab024c329aa 100644 --- a/datafusion/spark/src/function/datetime/mod.rs +++ b/datafusion/spark/src/function/datetime/mod.rs @@ -32,6 +32,7 @@ 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; @@ -59,6 +60,7 @@ 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, @@ -186,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> { @@ -212,5 +219,6 @@ pub fn functions() -> Vec> { unix_micros(), unix_millis(), unix_seconds(), + weekday(), ] } 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/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(); From 5615f251245c60cc72f0f449e299c89a1cb82505 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 10 Jun 2026 12:50:16 -0400 Subject: [PATCH 06/10] fix: Correct output-count stats for partitioned partial aggs (#22780) ## Which issue does this PR close? - Closes #22779 ## Rationale for this change A partial aggregate with no group by expression emits one row per output partition, even for input partitions that did not receive any rows. The stats code gets this incorrect, and claims that the partial agg outputs `Exact(0)` rows in this scenario. This is off by a factor of `partition_count`, which can lead to suboptimal planning decisions downstream. Because `total_byte_size` is scaled from the output row count, this is also improved: the byte-size estimate now scales with the corrected output row count rather than the input row count. ## What changes are included in this PR? * Pass the requested partition into `statistics_inner` and adjust `statistics_inner` as described above * Tighten optimizer metadata for grouping sets queries (minor correctness fix) * Add/extend existing unit tests ## Are these changes tested? Yes; new tests added. ## Are there any user-facing changes? No. --- .../partition_statistics.rs | 37 ++++ .../physical-plan/src/aggregates/mod.rs | 197 ++++++++++++++++-- 2 files changed, 215 insertions(+), 19 deletions(-) diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs index 4fba94ec3a1dc..e1bf22201dbad 100644 --- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs @@ -982,6 +982,43 @@ 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, + *agg_partial.partition_statistics(Some(0))? + ); + assert_eq!( + expect_partial_stat, + *agg_partial.partition_statistics(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, + *agg_partial.partition_statistics(None)? + ); + + // 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( diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index a5f1621812561..940bdd41a88e4 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -873,6 +873,7 @@ impl AggregateExec { &input, Arc::clone(&schema), &group_expr_mapping, + group_by.is_true_no_grouping(), &mode, &input_order_mode, aggr_expr.as_ref(), @@ -1115,6 +1116,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], @@ -1124,9 +1126,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) @@ -1185,6 +1190,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: /// @@ -1222,7 +1232,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 @@ -1246,20 +1260,18 @@ 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, }) } - _ => { + None => { let num_rows = self.estimate_num_rows(child_statistics); let total_byte_size = num_rows @@ -1280,6 +1292,51 @@ 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(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() + } + }) + } + + /// 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 { @@ -1708,7 +1765,9 @@ impl ExecutionPlan for AggregateExec { fn partition_statistics(&self, partition: Option) -> Result> { let child_statistics = self.input().partition_statistics(partition)?; - Ok(Arc::new(self.statistics_inner(&child_statistics)?)) + Ok(Arc::new( + self.statistics_inner(&child_statistics, partition)?, + )) } fn cardinality_effect(&self) -> CardinalityEffect { @@ -4514,6 +4573,26 @@ mod tests { let stats_zero = agg_zero.partition_statistics(None)?; 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 = single_agg_zero.partition_statistics(None)?; + assert_eq!(single_stats_zero.num_rows, Precision::Exact(1)); + Ok(()) } @@ -4522,19 +4601,39 @@ mod tests { 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), @@ -4969,6 +5068,66 @@ 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!( + single_agg.partition_statistics(None)?.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!( + partial_agg.partition_statistics(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!( + partial_agg.partition_statistics(None)?.num_rows, + Precision::Exact(4) + ); + + Ok(()) + } + #[test] fn test_aggregate_stats_non_column_expr_bails_out() -> Result<()> { use datafusion_common::ColumnStatistics; From 68cbe60b1d965085a0613cf444913d53e149aa53 Mon Sep 17 00:00:00 2001 From: Nathan <56370526+nathanb9@users.noreply.github.com> Date: Wed, 10 Jun 2026 12:51:24 -0400 Subject: [PATCH 07/10] fix: map() fails when keys are literals and values are column expressions (#22784) ## Which issue does this PR close? Closes #22781 ## Rationale for this change `map(['a','b'], [col, col * 10])` fails at execution time with "map requires key and value lists to have the same length" when keys are all literals and values contain column references. Root cause: literal keys evaluate to `ColumnarValue::Scalar(FixedSizeList[N])` (length N) while column values evaluate to `ColumnarValue::Array` (length batch_size). The length check compares N != batch_size and errors. ## What changes are included in this PR? In `make_map_batch`: when one argument is scalar and the other is array, expand the scalar to match `batch_size` via `ScalarValue::to_array_of_size(number_rows)` before the length comparison. ## Are these changes tested? Yes. - Unit test: scalar keys + array values in `map.rs` - SLT test: `SELECT map(['a','b'], [column1, column1 * 10]) FROM (VALUES (1), (2), (3))` ## Are there any user-facing changes? No. Previously-failing queries now execute correctly. --------- Co-authored-by: Jeffrey Vo --- datafusion/functions-nested/src/map.rs | 72 +++++++++++++++------- datafusion/sqllogictest/test_files/map.slt | 15 +++++ 2 files changed, 66 insertions(+), 21 deletions(-) diff --git a/datafusion/functions-nested/src/map.rs b/datafusion/functions-nested/src/map.rs index c7418e9021494..868f6dd29009b 100644 --- a/datafusion/functions-nested/src/map.rs +++ b/datafusion/functions-nested/src/map.rs @@ -63,15 +63,33 @@ fn can_evaluate_to_const(args: &[ColumnarValue]) -> bool { .all(|arg| matches!(arg, ColumnarValue::Scalar(_))) } -fn make_map_batch(args: &[ColumnarValue]) -> Result { +fn expand_if_scalar(arg: &ColumnarValue, rows: usize) -> Result { + match arg { + ColumnarValue::Scalar(s) => Ok(ColumnarValue::Array(s.to_array_of_size(rows)?)), + ColumnarValue::Array(a) => Ok(ColumnarValue::Array(Arc::clone(a))), + } +} + +fn make_map_batch(args: &[ColumnarValue], number_rows: usize) -> Result { let [keys_arg, values_arg] = take_function_args("make_map", args)?; let can_evaluate_to_const = can_evaluate_to_const(args); - let keys = get_first_array_ref(keys_arg)?; + // if we can't evaluate to const (inputs are not both scalar) then ensure they + // are expanded to arrays which following logic expects + let (keys_arg, values_arg) = if !can_evaluate_to_const { + ( + expand_if_scalar(keys_arg, number_rows)?, + expand_if_scalar(values_arg, number_rows)?, + ) + } else { + (keys_arg.clone(), values_arg.clone()) + }; + + let keys = get_first_array_ref(&keys_arg)?; let key_array = keys.as_ref(); - match keys_arg { + match &keys_arg { ColumnarValue::Array(_) => match key_array.data_type() { DataType::List(_) => keys .as_list::() @@ -101,7 +119,7 @@ fn make_map_batch(args: &[ColumnarValue]) -> Result { } } - let values = get_first_array_ref(values_arg)?; + let values = get_first_array_ref(&values_arg)?; make_map_batch_internal(&keys, &values, can_evaluate_to_const, &keys_arg.data_type()) } @@ -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( + &[ + 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( + &[ + 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( + &[ + 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( + &[ + ColumnarValue::Array(keys_array), + ColumnarValue::Array(values_array), + ], + 3, + ); assert!( result.is_ok(), diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index 62e70e6080bab..2b390c3748e35 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -901,3 +901,18 @@ 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} From 8f0635438fe605da4181852fd8a1332aea42750b Mon Sep 17 00:00:00 2001 From: linfeng <33561138+lyne7-sc@users.noreply.github.com> Date: Thu, 11 Jun 2026 00:51:46 +0800 Subject: [PATCH 08/10] minor: handle NULL array input in array_remove and array_replace (#22790) ## Which issue does this PR close? - Closes #. ## Rationale for this change `array_remove` errored out when given a `NULL` array argument. `array_replace` handled the `DataType::Null` case, but it always built a NULL array of length 1, regardless of the actual input length. `remove` and `replace` share nearly identical implementations, so they should treat a NULL array argument the same way ## What changes are included in this PR? - `array_remove`: handle `DataType::Null` in both `array_remove_internal` and the scalar fast-path `array_remove_with_scalar_args`, returning a NULL array of the input length. - `array_replace`: fix the existing `DataType::Null` branches to use the input array length (`array.len()` / `list_array.len()`) instead of a hard-coded `1`. ## Are these changes tested? Yes, added slt coverage ## Are there any user-facing changes? `array_remove`/`array_remove_n`/`array_remove_all` now return `NULL` for a `NULL`-typed array argument instead of raising an error. No breaking API changes. --- datafusion/functions-nested/src/remove.rs | 2 + datafusion/functions-nested/src/replace.rs | 4 +- .../test_files/array/array_remove.slt | 40 ++++++++++++++----- .../test_files/array/array_replace.slt | 31 ++++++++++++++ 4 files changed, 64 insertions(+), 13 deletions(-) diff --git a/datafusion/functions-nested/src/remove.rs b/datafusion/functions-nested/src/remove.rs index 44ef56c039b71..9d7dd4f44d91b 100644 --- a/datafusion/functions-nested/src/remove.rs +++ b/datafusion/functions-nested/src/remove.rs @@ -403,6 +403,7 @@ 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}'.") } @@ -425,6 +426,7 @@ fn array_remove_with_scalar_args( 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}'." ), diff --git a/datafusion/functions-nested/src/replace.rs b/datafusion/functions-nested/src/replace.rs index f129972fc7ea8..e8a95cff9670c 100644 --- a/datafusion/functions-nested/src/replace.rs +++ b/datafusion/functions-nested/src/replace.rs @@ -631,7 +631,7 @@ fn array_replace_with_scalar_args( let list = list_array.as_list::(); general_replace_with_scalar::(list, &needle, scalar_to, max_replacements) } - DataType::Null => Ok(new_null_array(list_array.data_type(), 1)), + DataType::Null => Ok(new_null_array(list_array.data_type(), list_array.len())), array_type => exec_err!("array_replace does not support type '{array_type}'."), } } @@ -651,7 +651,7 @@ fn array_replace_internal( let list_array = array.as_list::(); general_replace::(list_array, from, to, arr_n) } - DataType::Null => Ok(new_null_array(array.data_type(), 1)), + DataType::Null => Ok(new_null_array(array.data_type(), array.len())), array_type => exec_err!("array_replace does not support type '{array_type}'."), } } diff --git a/datafusion/sqllogictest/test_files/array/array_remove.slt b/datafusion/sqllogictest/test_files/array/array_remove.slt index 23ebf00239530..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 @@ -406,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); diff --git a/datafusion/sqllogictest/test_files/array/array_replace.slt b/datafusion/sqllogictest/test_files/array/array_replace.slt index 77793228c9ebe..cab84007bcd53 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; From 656dc47ac9563edaa795a2da5b513bce1e5716ed Mon Sep 17 00:00:00 2001 From: Bruce Ritchie Date: Wed, 10 Jun 2026 12:57:53 -0400 Subject: [PATCH 09/10] Add sort tpch SQL benchmark (#22814) ## Which issue does this PR close? Part of #21706 ## Rationale for this change Continue work on sql benchmark migration. ## What changes are included in this PR? sort_tpch sql benchmark ## Are these changes tested? Yes `BENCH_NAME=sort_tpch cargo bench --bench sql` `BENCH_NAME=sort_tpch BENCH_SIZE=10 cargo bench --bench sql` `BENCH_NAME=sort_tpch LIMIT=true cargo bench --bench sql` `BENCH_NAME=sort_tpch BENCH_SIZE=10 LIMIT=true cargo bench --bench sql` ## Are there any user-facing changes? No --------- Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com> --- .../sort_tpch/benchmarks/q01.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q02.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q03.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q04.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q05.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q06.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q07.benchmark | 58 ++++++++++++++++++ .../sort_tpch/benchmarks/q08.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q09.benchmark | 44 ++++++++++++++ .../sort_tpch/benchmarks/q10.benchmark | 59 +++++++++++++++++++ .../sort_tpch/benchmarks/q11.benchmark | 44 ++++++++++++++ .../sql_benchmarks/sort_tpch/init/load.sql | 3 + 12 files changed, 516 insertions(+) create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark create mode 100644 benchmarks/sql_benchmarks/sort_tpch/init/load.sql 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 From 8bbc46004183c22bd6caf0b7492c21540248a084 Mon Sep 17 00:00:00 2001 From: Neil Conway Date: Wed, 10 Jun 2026 13:51:22 -0400 Subject: [PATCH 10/10] fix: Avoid incorrectly rounding large integers in `nanvl` (#22575) ## Which issue does this PR close? - Closes #22567 ## Rationale for this change `nanvl`'s function signature listed `Float16` first, which meant that data types not natively supported by the function (e.g., integers) would be converted to `Float16`. This resulted in unnecessary and incorrect rounding for large integer inputs (e.g., `nanvl(16777217, 1)` returned incorrect results). Instead, we now convert non-floating point inputs to `Float64`. This might still round for inputs larger than 2^53, but it's the best we can do without more fundamental changes. Floating point inputs are coerced to the widest float type they have in common -- e.g., (`Float16`, `Float32`) -> `Float32`. ## What changes are included in this PR? * Change function signature to avoid unnecessary rounding * Update SLT ## Are these changes tested? Yes; new tests added. ## Are there any user-facing changes? Yes; `nanvl` now avoids rounding in more cases. --------- Co-authored-by: Kumar Ujjawal --- datafusion/functions/src/math/nanvl.rs | 59 +++++++++++++------ datafusion/sqllogictest/test_files/scalar.slt | 25 ++++++++ 2 files changed, 66 insertions(+), 18 deletions(-) diff --git a/datafusion/functions/src/math/nanvl.rs b/datafusion/functions/src/math/nanvl.rs index 251e98bb72c03..b1f69032efae6 100644 --- a/datafusion/functions/src/math/nanvl.rs +++ b/datafusion/functions/src/math/nanvl.rs @@ -20,11 +20,11 @@ use std::sync::Arc; use arrow::array::{ArrayRef, AsArray, Float16Array, Float32Array, Float64Array}; use arrow::datatypes::DataType::{Float16, Float32, Float64}; use arrow::datatypes::{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; @@ -63,12 +63,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 +106,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 +117,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,6 +133,15 @@ 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) diff --git a/datafusion/sqllogictest/test_files/scalar.slt b/datafusion/sqllogictest/test_files/scalar.slt index 9dbf8f16d85ab..64874eb316d8b 100644 --- a/datafusion/sqllogictest/test_files/scalar.slt +++ b/datafusion/sqllogictest/test_files/scalar.slt @@ -796,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;