Add generate_series() udtf (and introduce 'lazy' `MemoryExec`) #13540

2010YOUY01 · 2024-11-23T16:15:42Z

Which issue does this PR close?

Rationale for this change

To make it easier to test memory-limited queries, i was looking for a datasource which can generate long streams, and itself doesn't consume lots of memory.
There are several existing approaches are close, but can't fully satisfy this requirement:

MemoryExec have to buffer all output data up front
select * from unnest(generate_series(1,10000)) also have to buffer all output during construction
It's possible to do such tests on large files, however it would be hard to setup and takes longer to run

So in this PR, a new StreamingMemoryExec is introduced, it can be used to generate batches lazily by defining a 'iterator' on RecordBatch in its interface. And a new user-defined table function is implemented with this new execution plan.

I also think this UDTF can be useful to tests/micro-benchmarks in general

> select * from generate_series(1,3);
+-------+
| value |
+-------+
| 1     |
| 2     |
| 3     |
+-------+

What changes are included in this PR?

Move TableFunctionImpl from datafusion crate -> datafusion-catalog crate, so new UDTFs defined in datafusion-functions-table crate won't have circular dependency with datafusion crate (TableFunctionImpl depends on TableProvider in datafusion-catalog crate, so I think it's the best place to move it to)
Introduced StreamingMemoryExec
Implemented a new UDTF generate_series()

Are these changes tested?

Unit tests for StreamingMemoryExec, and sqllogictests for generate_series() UDTF

Are there any user-facing changes?

No

2010YOUY01 · 2024-11-23T16:23:27Z

datafusion/physical-plan/src/memory.rs

+    ) -> Result<Self> {
+        let cache = PlanProperties::new(
+            EquivalenceProperties::new(Arc::clone(&schema)),
+            Partitioning::RoundRobinBatch(generators.len()),


Is this field mean: Let's say a exec only have 1 generator, the optimizer will try to insert a RepartitionExec to the output of StreamingMemoryExec, and use round robin to repartition output batches to target_partitions number?
I found it works on some aggregate queries, but not for a sort query

> explain select * from generate_series(1, 10000) as t1(v1) order by v1; +---------------+----------------------------------------------------------------------------------------------------------------+ | plan_type | plan | +---------------+----------------------------------------------------------------------------------------------------------------+ | logical_plan | Sort: t1.v1 ASC NULLS LAST | | | SubqueryAlias: t1 | | | Projection: tmp_table.value AS v1 | | | TableScan: tmp_table projection=[value] | | physical_plan | SortExec: expr=[v1@0 ASC NULLS LAST], preserve_partitioning=[false] | | | ProjectionExec: expr=[value@0 as v1] | | | StreamingMemoryExec: partitions=1, batch_generators=[generate_series: start=1, end=10000, batch_size=8192] | | | | +---------------+----------------------------------------------------------------------------------------------------------------+

Parallelized sort should look like

> explain select * from lineitem order by l_orderkey; +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | plan_type | plan | +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+ | logical_plan | Sort: lineitem.l_orderkey ASC NULLS LAST | | | TableScan: lineitem projection=[l_orderkey, l_partkey, l_suppkey, l_linenumber, l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate, l_commitdate, l_receiptdate, l_shipinstruct, l_shipmode, l_comment] | | physical_plan | SortPreservingMergeExec: [l_orderkey@0 ASC NULLS LAST] | | | SortExec: expr=[l_orderkey@0 ASC NULLS LAST], preserve_partitioning=[true] | | | ParquetExec: file_groups={14 groups: [[Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-0.parquet:0..11525426], [Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-0.parquet:11525426..20311205, Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-1.parquet:0..2739647], [Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-1.parquet:2739647..14265073], [Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-1.parquet:14265073..20193593, Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-2.parquet:0..5596906], [Users/yongting/Code/datafusion/benchmarks/data/tpch_sf1/lineitem/part-2.parquet:5596906..17122332], ...]}, projection=[l_orderkey, l_partkey, l_suppkey, l_linenumber, l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate, l_commitdate, l_receiptdate, l_shipinstruct, l_shipmode, l_comment] | | | | +---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

I want to know is it correct here, and this sort query in theory can be made parallelized by changing somewhere in optimization phase? 🤔

jayzhan211 · 2024-11-26T02:29:14Z

datafusion/physical-plan/src/memory.rs

+/// Stream that generates record batches on demand
+pub struct StreamingMemoryStream {
+    schema: SchemaRef,
+    generator: Box<dyn StreamingBatchGenerator>,


Should we use Arc given it makes sense to have generator generate across threads

jayzhan211 · 2024-11-26T02:29:46Z

datafusion/physical-plan/src/memory.rs

+    /// Schema representing the data
+    schema: SchemaRef,
+    /// Functions to generate batches for each partition
+    batch_generators: Vec<Box<dyn StreamingBatchGenerator>>,


Suggested change

batch_generators: Vec<Box<dyn StreamingBatchGenerator>>,

batch_generators: Vec<Arc<dyn StreamingBatchGenerator>>,

jayzhan211 · 2024-11-26T02:31:46Z

datafusion/catalog/src/table.rs

@@ -297,3 +299,40 @@ pub trait TableProviderFactory: Debug + Sync + Send {
        cmd: &CreateExternalTable,
    ) -> Result<Arc<dyn TableProvider>>;
 }
+
+/// A trait for table function implementations
+pub trait TableFunctionImpl: Debug + Sync + Send {


Suggested change

pub trait TableFunctionImpl: Debug + Sync + Send {

/// Returns this function's name

fn name(&self) -> &str;

This is much more consistent with other UDFs

jayzhan211 · 2024-11-26T02:40:00Z

datafusion/sqllogictest/test_files/table_functions.slt

+SELECT * FROM generate_series(5, 1)
+
+statement error DataFusion error: Error during planning: First argument must be an integer literal
+SELECT * FROM generate_series(NULL, 5)


We might not return error for this query

D SELECT * FROM generate_series(null, 5); ┌─────────────────┐ │ generate_series │ │ int64 │ ├─────────────────┤ │ 0 rows │ └─────────────────┘

jayzhan211 · 2024-11-26T02:40:24Z

datafusion/sqllogictest/test_files/table_functions.slt

+SELECT * FROM generate_series(NULL, 5)
+
+statement error DataFusion error: Error during planning: Second argument must be an integer literal
+SELECT * FROM generate_series(1, NULL)


D SELECT * FROM generate_series(1, null); ┌─────────────────┐ │ generate_series │ │ int64 │ ├─────────────────┤ │ 0 rows │ └─────────────────┘

jayzhan211 · 2024-11-26T02:40:42Z

datafusion/sqllogictest/test_files/table_functions.slt

+statement error DataFusion error: Error during planning: Second argument must be an integer literal
+SELECT * FROM generate_series(1, NULL)
+
+statement error DataFusion error: Error during planning: generate_series expects 2 arguments


D SELECT * FROM generate_series(1, 5, null); ┌─────────────────┐ │ generate_series │ │ int64 │ ├─────────────────┤ │ 0 rows │ └─────────────────┘

Add generate_series() udtf

ac50592

github-actions bot added physical-expr Physical Expressions core Core DataFusion crate sqllogictest SQL Logic Tests (.slt) catalog Related to the catalog crate labels Nov 23, 2024

liscence

92e94e6

2010YOUY01 commented Nov 23, 2024

View reviewed changes

2010YOUY01 added 2 commits November 24, 2024 00:26

fix examples

95f95f2

clippy

5df8964

jayzhan211 reviewed Nov 26, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add generate_series() udtf (and introduce 'lazy' `MemoryExec`) #13540

Add generate_series() udtf (and introduce 'lazy' `MemoryExec`) #13540

2010YOUY01 commented Nov 23, 2024 •

edited

Loading

2010YOUY01 Nov 23, 2024 •

edited

Loading

jayzhan211 Nov 26, 2024

jayzhan211 Nov 26, 2024

jayzhan211 Nov 26, 2024

jayzhan211 Nov 26, 2024

jayzhan211 Nov 26, 2024

jayzhan211 Nov 26, 2024

	batch_generators: Vec<Box<dyn StreamingBatchGenerator>>,
	batch_generators: Vec<Arc<dyn StreamingBatchGenerator>>,

	pub trait TableFunctionImpl: Debug + Sync + Send {
	/// Returns this function's name
	fn name(&self) -> &str;

Add generate_series() udtf (and introduce 'lazy' MemoryExec) #13540

Are you sure you want to change the base?

Add generate_series() udtf (and introduce 'lazy' MemoryExec) #13540

Conversation

2010YOUY01 commented Nov 23, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

2010YOUY01 Nov 23, 2024 • edited Loading

Choose a reason for hiding this comment

jayzhan211 Nov 26, 2024

Choose a reason for hiding this comment

jayzhan211 Nov 26, 2024

Choose a reason for hiding this comment

jayzhan211 Nov 26, 2024

Choose a reason for hiding this comment

jayzhan211 Nov 26, 2024

Choose a reason for hiding this comment

jayzhan211 Nov 26, 2024

Choose a reason for hiding this comment

jayzhan211 Nov 26, 2024

Choose a reason for hiding this comment

Add generate_series() udtf (and introduce 'lazy' `MemoryExec`) #13540

Add generate_series() udtf (and introduce 'lazy' `MemoryExec`) #13540

2010YOUY01 commented Nov 23, 2024 •

edited

Loading

2010YOUY01 Nov 23, 2024 •

edited

Loading