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/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..b6f1a37e3d03f --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q01 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q1: 1 sort key (type: INTEGER, cardinality: 7) + 1 payload column +SELECT l_linenumber, l_partkey +FROM lineitem +ORDER BY l_linenumber +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q01.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..1238beb00583a --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q02 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q2: 1 sort key (type: BIGINT, cardinality: 1.5M) + 1 payload column +SELECT l_orderkey, l_partkey +FROM lineitem +ORDER BY l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q02.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..aadbe86c61602 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q03 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q3: 1 sort key (type: VARCHAR, cardinality: 4.5M) + 1 payload column +SELECT l_comment, l_partkey +FROM lineitem +ORDER BY l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q03.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..8119a6c51be33 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q04 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q4: 2 sort keys {(BIGINT, 1.5M), (INTEGER, 7)} + 1 payload column +SELECT l_orderkey, l_linenumber, l_partkey +FROM lineitem +ORDER BY l_orderkey, l_linenumber +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q04.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..5ee9e610cc3bf --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q05 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q5: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + no payload column +SELECT l_linenumber, l_suppkey, l_orderkey +FROM lineitem +ORDER BY l_linenumber, l_suppkey, l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q05.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..54ce6fa44341d --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q06 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q6: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + 1 payload column +SELECT l_linenumber, l_suppkey, l_orderkey, l_partkey +FROM lineitem +ORDER BY l_linenumber, l_suppkey, l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q06.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..8932810cc1f97 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark @@ -0,0 +1,58 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q07 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q7: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + 12 all other columns +SELECT l_linenumber, + l_suppkey, + l_orderkey, + l_partkey, + l_quantity, + l_extendedprice, + l_discount, + l_tax, + l_returnflag, + l_linestatus, + l_shipdate, + l_commitdate, + l_receiptdate, + l_shipinstruct, + l_shipmode +FROM lineitem +ORDER BY l_linenumber, l_suppkey, l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q07.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..f09e6e9f72f21 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q08 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q8: 4 sort keys {(BIGINT, 1.5M), (BIGINT, 10k), (INTEGER, 7), (VARCHAR, 4.5M)} + no payload column +SELECT l_orderkey, l_suppkey, l_linenumber, l_comment +FROM lineitem +ORDER BY l_orderkey, l_suppkey, l_linenumber, l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q08.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..5e7a2ea63747a --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q09 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q9: 4 sort keys {(BIGINT, 1.5M), (BIGINT, 10k), (INTEGER, 7), (VARCHAR, 4.5M)} + 1 payload column +SELECT l_orderkey, l_suppkey, l_linenumber, l_comment, l_partkey +FROM lineitem +ORDER BY l_orderkey, l_suppkey, l_linenumber, l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q09.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..535393526147e --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark @@ -0,0 +1,59 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q10 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q10: 4 sort keys {(BIGINT, 1.5M), (BIGINT, 10k), (INTEGER, 7), (VARCHAR, 4.5M)} + 12 all other columns +SELECT l_orderkey, + l_suppkey, + l_linenumber, + l_comment, + l_partkey, + l_quantity, + l_extendedprice, + l_discount, + l_tax, + l_returnflag, + l_linestatus, + l_shipdate, + l_commitdate, + l_receiptdate, + l_shipinstruct, + l_shipmode +FROM lineitem +ORDER BY l_orderkey, l_suppkey, l_linenumber, l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q10.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..efce2005f3beb --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q11 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q11: 1 sort key (type: VARCHAR, cardinality: 4.5M) + 1 payload column +SELECT l_shipmode, l_comment, l_partkey +FROM lineitem +ORDER BY l_shipmode +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q11.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/init/load.sql b/benchmarks/sql_benchmarks/sort_tpch/init/load.sql new file mode 100644 index 0000000000000..395d8da009d21 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/init/load.sql @@ -0,0 +1,3 @@ +CREATE EXTERNAL TABLE lineitem_raw STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/lineitem/lineitem.1.parquet'; + +CREATE TABLE lineitem as (SELECT * FROM lineitem_raw${BENCH_SORTED:-false| order by l_orderkey asc| }); \ No newline at end of file diff --git a/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] 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/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/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/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/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/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 { 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; 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/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/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; 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} diff --git a/datafusion/sqllogictest/test_files/pipe_operator.slt b/datafusion/sqllogictest/test_files/pipe_operator.slt index 406ddafc7bdea..4e2b867fd744d 100644 --- a/datafusion/sqllogictest/test_files/pipe_operator.slt +++ b/datafusion/sqllogictest/test_files/pipe_operator.slt @@ -15,11 +15,6 @@ # specific language governing permissions and limitations # under the License. -# BigQuery supports the pipe operator syntax -# TODO: Make the Generic dialect support the pipe operator syntax -statement ok -set datafusion.sql_parser.dialect = 'BigQuery'; - statement ok CREATE TABLE test( a INT, @@ -188,14 +183,10 @@ query TII |> AS produce_sales |> LEFT JOIN ( - SELECT "apples" AS item, 123 AS id + SELECT 'apples' AS item, 123 AS id ) AS produce_data ON produce_sales.item = produce_data.item |> SELECT produce_sales.item, sales, id; ---- apples 2 123 bananas 5 NULL - -# Config reset -statement ok -RESET datafusion.sql_parser.dialect; diff --git a/datafusion/sqllogictest/test_files/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; diff --git a/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt b/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt index b4f5444e8a2da..efa6b898c2a6a 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt @@ -23,5 +23,109 @@ ## Original Query: SELECT weekday('2009-07-30'); ## PySpark 3.5.5 Result: {'weekday(2009-07-30)': 3, 'typeof(weekday(2009-07-30))': 'int', 'typeof(2009-07-30)': 'string'} -#query -#SELECT weekday('2009-07-30'::string); +# Spark `weekday` is 0-indexed with Monday = 0 .. Sunday = 6. +# 2009-07-30 is a Thursday -> 3. +query I +SELECT weekday('2009-07-30'::DATE); +---- +3 + +# All seven days of one week (2024-01-01 is a Monday). +query I +SELECT weekday('2024-01-01'::DATE); +---- +0 + +query I +SELECT weekday('2024-01-02'::DATE); +---- +1 + +query I +SELECT weekday('2024-01-03'::DATE); +---- +2 + +query I +SELECT weekday('2024-01-04'::DATE); +---- +3 + +query I +SELECT weekday('2024-01-05'::DATE); +---- +4 + +query I +SELECT weekday('2024-01-06'::DATE); +---- +5 + +query I +SELECT weekday('2024-01-07'::DATE); +---- +6 + +# NULL handling +query I +SELECT weekday(NULL::DATE); +---- +NULL + +# Array input (mix of weekdays and NULL) +query I +SELECT weekday(d) FROM (VALUES ('2024-01-01'::DATE), ('2024-01-06'::DATE), ('2024-01-07'::DATE), (NULL::DATE)) AS t(d); +---- +0 +5 +6 +NULL + +# Timestamp input: Spark coerces TIMESTAMP/TIMESTAMP_NTZ to DATE before evaluation +query I +SELECT weekday('2009-07-30 12:34:56'::TIMESTAMP); +---- +3 + +query I +SELECT weekday(NULL::TIMESTAMP); +---- +NULL + +# Timestamp array input +query I +SELECT weekday(ts) FROM (VALUES + ('2024-01-01 01:02:03'::TIMESTAMP), + ('2024-01-06 10:20:30'::TIMESTAMP), + ('2024-01-07 23:59:59'::TIMESTAMP), + (NULL::TIMESTAMP) +) AS t(ts); +---- +0 +5 +6 +NULL + +# TIMESTAMP_NTZ (Timestamp without timezone) — explicit Microsecond precision +query I +SELECT weekday(arrow_cast('2009-07-30 09:15:00', 'Timestamp(Microsecond, None)')); +---- +3 + +# TIMESTAMP with timezone (Spark TIMESTAMP / LTZ) — coerces to Date32 +query I +SELECT weekday(arrow_cast('2024-01-07 03:00:00', 'Timestamp(Nanosecond, Some("UTC"))')); +---- +6 + +# Error: wrong argument type (string without cast) +statement error Function 'weekday' requires Date, but received String +SELECT weekday('not-a-date'); + +# Error: wrong argument type (integer) +statement error Function 'weekday' requires Date, but received Int64 +SELECT weekday(123); + +# Error: no arguments +statement error 'weekday' does not support zero arguments +SELECT weekday(); diff --git a/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 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)