Repository navigation
[EPIC] JIT support for DataFusion #2703
Description
Summary
TLDR: The key focus of this work is to speed up fundamentally row oriented operations like hash table lookup or comparisons (e.g. #2427)
Background
DataFusion, like many Arrow systems, is a classic "vectorized computation engine" which works quite well for many common operations. The following paper, gives a good treatment on the various tradeoffs between vectorized and JIT's compilation of query plans: https://db.in.tum.de/~kersten/vectorization_vs_compilation.pdf?lang=de
As mentioned in the paper, there are some fundamentally "row oriented" operations in a database that are not typically amenable to vectorization. The "classics" are: Hash table updates in Joins and Hash Aggregates, as well as comparing tuples in sort.
Another example can be found in these slides from this presentation
@yjshen added initial support for JIT'ing in #1849 and it currently lives in https://github.com/apache/arrow-datafusion/tree/master/datafusion/jit. He also added partial support for aggregates in #2375
This ticket aims to be a central location for tracking the status of JIT compiling expressions for anyone who wants to contribute to this effort
Describe the solution you'd like
Activity
[-][EPIC] Full JIT support for `DataFusion`[/-][+][EPIC] JIT support for `DataFusion`[/+]alamb commentedon Jun 6, 2022
Actually, this is basically the same as #1861 so closing in favor of that ticket
leoluan2009 commentedon May 16, 2024
Hi @alamb , do we need to support expression JIT for performance like ClickHouse: https://clickhouse.com/blog/clickhouse-just-in-time-compiler-jit
alamb commentedon May 20, 2024
Hi @leoluan2009
In my opinion, I don't think DataFusion needs JIT to get good performance.
In general, I find the paper "Everything You Always Wanted to Know About Compiled and Vectorized Queries But Were Afraid to Ask" to explain the tradeoffs well
DataFusion is a vectorized engine and we haven't found areas where JIt would be compelling compared to vectorized code. The only area I can really think of would be to implement type specialized comparisons for sorting (to avoid the RowFormat) but we would need to have a pretty compelling benchmark showing improvements to justify I think
faucct commentedon May 24, 2024
Though the paper that you have mentioned admits that JIT-compilation is beneficial for OLTP workloads:
If DataFusion would have JIT, then it could be useful for building Online ML Feature Store engines.
alamb commentedon May 24, 2024
FWIW there is no reason you couldn't JIT compile arbitrary expressions and run them as UDFs. The API is now basically complete: https://docs.rs/datafusion/latest/datafusion/logical_expr/trait.ScalarUDFImpl.html
It would also be possible to do the same with aggregates / windows / etc
faucct commentedon May 24, 2024
I think that compiling SQL-expressions to UDFs by hand would kinda kill the whole point of the framework, but it seems like most of the framework would be irrelevant for the in-memory transformation of Online Features, so I guess it would be easier to build the same thing from scratch, though using the same ideas, like for example SQL and Arrow format for intermediate data representation.
alamb commentedon May 24, 2024
FIW there is a lot more to SQL evaluation than just the expression evaluation, so that might be a reason to use DataFusion even if you had to implement your own expressions 🤔
zero-element commentedon Mar 2, 2025
I agree with @faucct that this approach could be beneficial for building Online ML Feature Store engines. This is because such systems often involve numerous element-wise calculations, which can be memory bandwith bound (considering low flops/bytes ratio with result of each operation needs to be written to memory).
JIT is not inherently in conflict with vectorized code. In fact, the paper referenced earlier also mentioned:
This concept may be similar to Whole-Stage Code Generation in Apache Spark. However, I am not entirely certain about this similarity without inspecting the actual implementation details.
UPDATE: I looked into some discussions further, and I’m wondering if the JIT introduced in https://github.com/apache/datafusion/pull/2587 has been removed, as the generated JIT code does not make good use of vectorized instructions. However, it seems that for the scenario mentioned above, what we need is not full JIT but rather the concept described in the PR above: "glue all the low-level compute functions together."