IN LIST: add branchless filter for small primitive lists - #23014
Conversation
428e3cd to
eae4046
Compare
eae4046 to
3e3651b
Compare
3e3651b to
6a1869f
Compare
bf69ed2 to
e040466
Compare
04ca744 to
f89d3b6
Compare
f89d3b6 to
409b0cb
Compare
409b0cb to
78e3631
Compare
6795d52 to
3abc50d
Compare
in fact DataFuson does this rewrite explicitly today: DataFusion CLI v54.0.0
> create table t (v int) as values (1),(2), (3), (4);
0 row(s) fetched.
Elapsed 0.034 seconds.
> explain select * from t where v IN (10, 20, 30);
+---------------+-------------------------------+
| plan_type | plan |
+---------------+-------------------------------+
| physical_plan | ┌───────────────────────────┐ |
| | │ FilterExec │ |
| | │ -------------------- │ |
| | │ predicate: │ |
| | │ v = 10 OR v = 20 OR v = 30│ |
| | └─────────────┬─────────────┘ |
| | ┌─────────────┴─────────────┐ |
| | │ DataSourceExec │ |
| | │ -------------------- │ |
| | │ bytes: 112 │ |
| | │ format: memory │ |
| | │ rows: 1 │ |
| | └───────────────────────────┘ |
| | |
+---------------+-------------------------------+
1 row(s) fetched.
Elapsed 0.014 seconds. |
alamb
left a comment
There was a problem hiding this comment.
I am pretty sure this code is unreachable via SQL / the external APIs as small IN LISTS get rewritten to an OR chain (see my other comment)
Thus if we are going to add this specialization, we should also turn off the IN LIST rewrite (and benchmark performance against the OR)
The OR rewrite threshold is currently set to 3 by THRESHOLD_INLINE_INLIST, so this code is reachable starting at size 4. The OR rewrite also doesn't trigger for dynamic filters, which directly create a standard IN filter. create table t (v int) as values (1),(2),(3),(4);
explain select * from t where v IN (10, 20, 30, 40);
Output:
+---------------+-------------------------------+
| plan_type | plan |
+---------------+-------------------------------+
| physical_plan | ┌───────────────────────────┐ |
| | │ FilterExec │ |
| | │ -------------------- │ |
| | │ predicate: │ |
| | │ v IN (10, 20, 30, 40) │ |
| | └─────────────┬─────────────┘ |
| | ┌─────────────┴─────────────┐ |
| | │ DataSourceExec │ |
| | │ -------------------- │ |
| | │ bytes: 112 │ |
| | │ format: memory │ |
| | │ rows: 1 │ |
| | └───────────────────────────┘ |
| | |
+---------------+-------------------------------+ -- Dynamic filters also reach IN-list evaluation directly.
set datafusion.explain.analyze_level = 'dev';
set datafusion.execution.target_partitions = 1;
explain analyze
with
build(a,b,c) as (values (1,10,100),(2,20,200),(3,30,300)),
probe(x,y,z) as (values (5,10,1000),(15,20,2000),(25,30,3000),(35,40,4000))
select *
from build
join probe on build.b = probe.y
order by x;
-- Relevant output:
-- FilterExec: DynamicFilter [ y@1 >= 10 AND y@1 <= 30 AND y@1 IN (SET) ([10, 20, 30]) ] |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #23014 +/- ##
=======================================
Coverage ? 80.74%
=======================================
Files ? 1089
Lines ? 369136
Branches ? 369136
=======================================
Hits ? 298058
Misses ? 53292
Partials ? 17786 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
I hope to review this very soon (like tomorrow) |
@alamb I looked into this a bit more (especially the OR rewrite micro-benchmarks), and I think this should actually be done as a dedicated new PR once the full IN LIST optimization stack has landed. This PR is still valuable by itself starting at lists of size 4 (and for dynamic filters), and removing the OR rewrite is not trivial to do (as the gains depend on the data type, eg utf-8 probably still benefits from OR rewrite...) |
|
I will review this today |
|
run benchmark in_list_strategy |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing perf/in_list_branchless_filter (06f2b6b) to f33dcec (merge-base) diff using: in_list_strategy File an issue against this benchmark runner |
alamb
left a comment
There was a problem hiding this comment.
Thank you @geoffreyclaude and @kosiew
My main concern about this Pr is that it generates a LOT of new code (by instantiating several hundred BranchlessFilter) -- by Claude Code's estiamte it is "~507 distinct BranchlessFilter<T, N> types, e" and by my measurement increases the code size by 4-6.7 percent
I think if we templated on the physical type (e.g. i32) rather than all the logical types we could reduce the number of distinct types substantially, but fundamentally creating many specialized code paths for the different sizes seems like a very expensive way to increase performance
Release binary size is 5MB (6.7%) larger and debug binary is 17MB (4%) larger
One way we could potentially reduce the code size would be to minimize the parts of the code that really need specializations. For example, I think the actual storage values could be a runtime value, and then just optimize the inner compare loop:
Compare the values 8 at a time or something using chunks_exact and then just generate a special comparison for 7,6,5,4,3,2,1 values to catch the last value.
That could also make this optimization more applicable as well, to any size where we wanted implement branchless (brute force comparison)
Details
Release
# release
cargo build --release --bin datafusion-cli
# debug
cargo build --bin datafusion-cliSize of datafusion-cli on main
74M /Users/andrewlamb/Downloads/datafusion-main-base-in-list
Size on this branch
79M /Users/andrewlamb/Downloads/datafusion-in-list
Debug size
main
412M target/debug/datafusion-cli
this branch
429M target/debug/datafusion-cli
| branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); | ||
| branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); | ||
|
|
||
| branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); |
There was a problem hiding this comment.
const BRANCHLESS_MAX_4B: usize = 32;
do I read the code above that this is creating 32 specializations (just) for UInt32 IN lists
There was a problem hiding this comment.
Yes, that was correct for the previous version: UInt32 alone generated 32 copies of the full BranchlessFilter 😬.
The follow-up removes the list length from BranchlessFilter, so there is now only one full filter per type. Only the small comparison function still varies by length. I put the detailed size and performance results in my reply to the overall review.
| }; | ||
| } | ||
|
|
||
| match in_array.data_type() { |
There was a problem hiding this comment.
i found it somewhat hard to follow by reading the code to understand what cases would use the branchless filter and what would use the bitmap filter for small integer types -- because the cases that are covered by branchless filters are behind some macros and then the bitmap filtering happens only if there isn't a matchess branchless filter
I am not sure there is a concrete action to take from this observation
There was a problem hiding this comment.
Good point. The latest refactor removed the per-length dispatch from this module, but it's still not super clear, I admit.
For the small types, Int8/UInt8 use direct comparisons for up to 16 non-null values, and Int16/UInt16/Float16 for up to 8. Longer lists fall back to the bitmap filter. I added a short comment to make that explicit...
| }; | ||
| } | ||
|
|
||
| branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); |
There was a problem hiding this comment.
Given we now have fast bitmap filter implementations for 8 and 16 bit integer types, what extra value do branchless filters give over those bitmaps?
UNless the performance win is compelling, I suggest we simply use the bitmap implementation for those types
There was a problem hiding this comment.
Branchless is still faster (1.5x) for small lists (eg, 4 values). Since the latest commits resolves the compilation size issue, I think the branchless path is worth keeping for short 8 and 16 bit type lists, with the bitmap filter remaining the fallback for larger ones.
| // The branchless filter reads the same Arrow value buffer as the | ||
| // comparison type. That is only valid when both native types have the | ||
| // same width, so catch any bad mapping here at compile time. | ||
| const _: () = assert!( |
There was a problem hiding this comment.
my reading of this loop is that it instantiates max_len copies of the branchless filter
There was a problem hiding this comment.
Yes, exactly: this loop was the main source of the code-size increase.
After the follow-up it still creates one small comparison function per supported length, but it no longer creates max_len copies of the whole filter. The rest of the filter is shared now; I included the resulting size and benchmark numbers in my reply to the overall review.
|
🤖 Benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagein_list_strategy — base (merge-base)
in_list_strategy — branch
File an issue against this benchmark runner |
@alamb I agree, and I pushed a new remediation commit along those lines. The I didn't use I also rebuilt release
That brings the release binary back to roughly the 74M For performance, I compared the two versions across 28 focused Criterion cases. The worst median regression was 2.45%; the next two were 1.58% and 1.57%. None reached 3%, and many cases were flat or faster. Thanks for pushing on this! I think this follow-up addresses the main concern without giving up the performance gains. WDYT? |
5af3108 to
08277d2
Compare
alamb
left a comment
There was a problem hiding this comment.
Thank you @geoffreyclaude -- I went through this PR again and I think it looks great.
I double checked this the code size and it seems like it is much smaller (a few kb I think)
I personally think it would help to have some more slt coverage, but I will propose that as a follow on PR
Code size measurements
(venv) andrewlamb@Andrews-MacBook-Pro-3:/tmp/branchless_in$ du -s -h target/release/datafusion-cli target/debug/datafusion-cli
74M target/release/datafusion-cli
414M target/debug/datafusion-cli
andrewlamb@Andrews-MacBook-Pro-3:/tmp/branchless_in_merge_base$ du -s -h target/release/datafusion-cli target/debug/datafusion-cli
74M target/release/datafusion-cli
413M target/debug/datafusion-cli| let non_null_count = in_array.len() - in_array.null_count(); | ||
| // `try_new` can be called on its own, so check the limit here too. | ||
| if non_null_count > T::MAX_LIST_LEN { | ||
| return Err(exec_datafusion_err!( |
There was a problem hiding this comment.
This is probably more correctly an internal error rather than an execution error as it isn't really based on user input (it would signal a bug in the code I think)
| 4 => choose!(0, 1, 2, 3, 4), | ||
| 8 => choose!(0, 1, 2, 3, 4, 5, 6, 7, 8), | ||
| 16 => choose!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16), | ||
| 32 => choose!( |
There was a problem hiding this comment.
this is cool -- it makes much more sense to isolate the code repetition to just the comparison function (which is where it is most helpful)
|
🚀 |
|
Thanks again @geoffreyclaude and @kosiew for the iteration -- |
Which issue does this PR close?
INperformance with specialized implementations #19390.Rationale for this change
For very small
INlists, building or probing a hash table can be more work than just comparing the input value with each constant.For example, for
x IN (10, 20, 30), the fast path can behave like:Because the list is tiny, those comparisons are cheap. The implementation stores the constants in a fixed-size array and checks them with a compact comparison chain.
“Branchless” here means the comparisons are combined without stopping at the first match. That can be faster for these small fixed-width lists because the CPU gets a predictable sequence of simple operations instead of hash-table setup and probe logic.
For primitive values that are not already plain unsigned integers, this PR keeps the logical Arrow type explicit and uses a matching same-width comparison representation only inside the branchless filter. For example,
Float16usesUInt16storage,Float32usesUInt32storage, andTimestampNanosecondusesUInt64storage.Decimal128andIntervalMonthDayNanouse their own 16-byte native representation. This preserves bit-pattern equality while relying on Arrow's native primitive compatibility rules: timestamp timezone metadata and Decimal128 precision/scale metadata may differ, while incompatible primitive representations remain rejected.What changes are included in this PR?
BranchlessFilterfor small primitiveINlists.strategy.rs.Int8->UInt8Int16,Float16->UInt16Int32,Float32,Date32,Time32->UInt32Int64,Float64,Date64,Time64,Timestamp,Duration->UInt64Decimal128,IntervalMonthDayNano-> their native 16-byte representationDecimal256and unsupported complex types on the generic path.IN/NOT INnull behavior as the rest of the stack.Are these changes tested?
Yes.
cargo fmt --all -- --checkcargo test -p datafusion-physical-expr expressions::in_list --libcargo test -p datafusion-physical-expr --bench in_list_strategy --no-runcargo clippy --all-targets --all-features -- -D warningsAre there any user-facing changes?
No. This is an internal performance optimization only.
Local benchmark snapshot
Built and run with
release-nonlto, filtered to the relevant small primitive-list rows:Filters used:
narrow_integer,primitive/i32/small_list,primitive/i64/small_list,f32/small_list,timestamp_ns/small_list, andinterval_month_day_nano/small_list.Method: directly compared Criterion's raw sample minima (
min(time / iterations)) fromsample.json. Lower is better; changes within +/-5% are treated as noise.Compared baselines: #23311 -> #23014
Relevant scope: small primitive-list rows.
Summary: 39 relevant rows, 28 faster, 0 slower, 11 within +/-5%.
Largest relevant deltas:
timestamp_ns/small_list/list=4/match=50%f32/small_list/list=4/match=50%primitive/i32/small_list/list=4/match=50%primitive/i64/small_list/list=4/match=50%timestamp_ns/small_list/list=4/match=0%f32/small_list/list=4/match=0%primitive/i32/small_list/list=4/match=0%primitive/i64/small_list/list=4/match=0%primitive/i32/small_list/list=16/match=50%/NOT_INnulls/primitive/i32/small_list/list=16/match=50%/nulls=20%timestamp_ns/small_list/list=16/match=50%nulls/primitive/i32/small_list/list=16/match=50%/nulls=50%nulls/primitive/i32/small_list/list=16/match=50%/nulls=20%/NOT_INinterval_month_day_nano/small_list/list=4/match=50%f32/small_list/list=32/match=50%primitive/i64/small_list/list=16/match=50%Full relevant table (39 rows)
narrow_integer/u8/list=4/match=0%narrow_integer/u8/list=4/match=50%narrow_integer/u8/list=16/match=0%narrow_integer/u8/list=16/match=50%narrow_integer/i16/list=4/match=0%narrow_integer/i16/list=4/match=50%narrow_integer/i16/list=64/match=0%narrow_integer/i16/list=64/match=50%narrow_integer/i16/list=256/match=0%narrow_integer/i16/list=256/match=50%narrow_integer/f16/list=4/match=0%narrow_integer/f16/list=4/match=50%narrow_integer/f16/list=64/match=0%narrow_integer/f16/list=64/match=50%narrow_integer/f16/list=256/match=0%narrow_integer/f16/list=256/match=50%nulls/narrow_integer/u8/list=16/match=50%/nulls=20%primitive/i32/small_list/list=4/match=0%primitive/i32/small_list/list=4/match=50%primitive/i32/small_list/list=32/match=0%primitive/i32/small_list/list=32/match=50%primitive/i32/small_list/list=16/match=50%/NOT_INnulls/primitive/i32/small_list/list=16/match=50%/nulls=20%nulls/primitive/i32/small_list/list=16/match=50%/nulls=20%/NOT_INnulls/primitive/i32/small_list/list=16/match=50%/nulls=50%primitive/i64/small_list/list=4/match=0%primitive/i64/small_list/list=4/match=50%primitive/i64/small_list/list=16/match=0%primitive/i64/small_list/list=16/match=50%f32/small_list/list=4/match=0%f32/small_list/list=4/match=50%f32/small_list/list=32/match=0%f32/small_list/list=32/match=50%timestamp_ns/small_list/list=4/match=0%timestamp_ns/small_list/list=4/match=50%timestamp_ns/small_list/list=16/match=0%timestamp_ns/small_list/list=16/match=50%interval_month_day_nano/small_list/list=4/match=0%interval_month_day_nano/small_list/list=4/match=50%