feat: add LazyPartition trait for LazyMemoryExec, migrate generate_series to partitions by ethan-tyler · Pull Request #20369 · apache/datafusion

ethan-tyler · 2026-02-15T15:15:52Z

Which issue does this PR close?

Closes Simplify LazyMemoryExec with SendableRecordBatchStream #13614

Rationale for this change

LazyMemoryExec currently uses generator closures (LazyBatchGenerator) as its partition interface. This replaces with a LazyPartition trait that gives partitions a proper abstraction while keeping the old generator path working through a legacy adapter.

generate_series/range are migrated to native LazyPartition as the first consumer. Proto roundtrip is updated to handle both native and legacy forms.

What changes are included in this PR?

LazyPartition trait as the native partition interface for LazyMemoryExec
LazyBatchGeneratorPartition adapter allowing existing generator to work
Deprecated compatibility APIs preserved: LazyMemoryExec::try_new, LazyMemoryExec::generators, GenerateSeriesTable::as_generator.
generate_series/range migrated to native GenerateSeriesPartition
Shared helper for generate series state construction (dedupes native and deprecated paths)
Proto roundtrip updated for native GenerateSeriesPartition with legacy adapter fallback
Core coop test helper updated to use LazyPartition
Proto regression test coverage added

Are these changes tested?

yes

cargo test -p datafusion-physical-plan lazy_memory_tests --lib
cargo test -p datafusion-functions-table generate_series_tests --lib
cargo test -p datafusion-proto roundtrip_generate_series --test proto_integration
cargo test -p datafusion execution::coop --test core_integration
cargo check -p datafusion-physical-plan -p datafusion-functions-table -p datafusion-proto -p datafusion

Are there any user-facing changes?

No. LazyBatchGenerator and existing public APIs remain available (deprecated)

ethan-tyler · 2026-02-17T11:53:55Z

@2010YOUY01 @Jefffrey - this might be of interest if you.

kosiew

@ethan-tyler

Thanks for working on this.

kosiew · 2026-03-05T10:02:01Z

datafusion/functions-table/src/generate_series.rs

        Self { schema, args }
    }

+    pub fn as_partition(&self, batch_size: usize) -> Result<Arc<dyn LazyPartition>> {


GenerateSeriesTable::as_partition currently constructs GenerateSeriesPartition without validating timestamp timezone input.
Previously as_generator parsed timezone eagerly and could fail during planning. Consider validating in GenerateSeriesPartition::new (or a try_new) to keep error timing consistent and fail earlier.

kosiew · 2026-03-05T10:07:32Z

datafusion/physical-plan/src/memory.rs

@@ -268,10 +408,10 @@ impl DisplayAs for LazyMemoryExec {
                write!(
                    f,
                    "LazyMemoryExec: partitions={}, batch_generators=[{}]",


DisplayAs still prints batch_generators= even though the primary abstraction is now partitions?

kosiew · 2026-03-05T10:11:47Z

datafusion/physical-plan/src/memory.rs

    /// Get the batch generators
+    #[deprecated(note = "Use LazyMemoryExec::partitions instead")]
+    #[expect(deprecated)]
    pub fn generators(&self) -> &Vec<Arc<RwLock<dyn LazyBatchGenerator>>> {


Not the fault of this PR but
generators() returns &Vec<_>. Since this is a compatibility accessor, returning a slice (&[_]) would avoid exposing container type and align with other accessor style.

kosiew · 2026-03-05T10:16:50Z

datafusion/proto/src/physical_plan/mod.rs

+    }
+
+    #[allow(deprecated)]
    fn try_from_lazy_memory_exec(exec: &LazyMemoryExec) -> Result<Option<Self>> {


The native GenerateSeriesPartition path uses serialize_generate_series_args, but legacy generator downcasts still manually reconstruct protobuf args.

Could we extract a helper that converts legacy generator variants into GenSeriesArgs first, then reuse serialize_generate_series_args to reduce duplicated conversion logic?

kosiew · 2026-03-05T10:19:15Z

datafusion/physical-plan/src/memory.rs

    pub fn try_set_partitioning(&mut self, partitioning: Partitioning) -> Result<()> {
        let partition_count = partitioning.partition_count();
-        let generator_count = self.batch_generators.len();
+        let generator_count = self.partitions.len();


generator_count although this now counts partitions?

…tions Signed-off-by: Ethan Urbanski <ethanurbanski@gmail.com>

Signed-off-by: Ethan Urbanski <ethanurbanski@gmail.com>

ethan-tyler · 2026-03-17T21:13:18Z

@ethan-tyler

Thanks for working on this.

@kosiew I must have forgotten to reply! Thanks for the review, I addressed your feedback. Let me know if there's anymore, thanks!

kosiew · 2026-03-18T07:24:49Z

@ethan-tyler

Just want to confirmt that you have pushed the changes already.

I noted a few outstanding points:

GenerateSeriesPartition::try_new does not eagerly validate timestamp timezones

The fail-fast behavior called out in the review remains unresolved.
- generate_series.rs (line 246)
- generate_series.rs (line 274)
These locations only validate:
```
batch_size == 0
```
Timezone parsing still happens later in:
```
build_generate_series_state
```
- generate_series.rs (line 362)

Naming cleanup in try_set_partitioning was not completed

The local variable is still named:
```
generator_count
```
even though it now counts partitions.
- memory.rs (line 372)

DisplayAs still emits batch_generators=
- memory.rs (line 418)

Proto code still duplicates the legacy generator-to-protobuf conversion

Instead of normalizing legacy variants into GenSeriesArgs and reusing:
```
serialize_generate_series_args
```
The native partition path uses the helper, but the legacy adapter path still manually reconstructs protobuf args.

github-actions bot added core Core DataFusion crate proto Related to proto crate functions Changes to functions implementation physical-plan Changes to the physical-plan crate labels Feb 15, 2026

kosiew reviewed Mar 5, 2026

View reviewed changes

ethan-tyler and others added 5 commits March 11, 2026 12:18

Feat physical plan add LazyPartition and wire LazyMemoryExec to parti…

b58ce18

…tions Signed-off-by: Ethan Urbanski <ethanurbanski@gmail.com>

Refactor functions table migrate generate_series range to LazyPartition

2eebfbf

Signed-off-by: Ethan Urbanski <ethanurbanski@gmail.com>

Test core migrate coop lazy source test to LazyPartition

5019968

Signed-off-by: Ethan Urbanski <ethanurbanski@gmail.com>

Feat proto support generate_series native partition roundtrip

3d5241d

Signed-off-by: Ethan Urbanski <ethanurbanski@gmail.com>

Address reviewer feedback on native generate_series

3253f6c

ethan-tyler force-pushed the lazy-partition-migration branch from 5bbb386 to 3253f6c Compare March 11, 2026 16:20

Merge branch 'main' into lazy-partition-migration

52d293e

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: add LazyPartition trait for LazyMemoryExec, migrate generate_series to partitions#20369

feat: add LazyPartition trait for LazyMemoryExec, migrate generate_series to partitions#20369
ethan-tyler wants to merge 6 commits intoapache:mainfrom
ethan-tyler:lazy-partition-migration

ethan-tyler commented Feb 15, 2026

Uh oh!

ethan-tyler commented Feb 17, 2026

Uh oh!

kosiew left a comment

Uh oh!

kosiew Mar 5, 2026

Uh oh!

kosiew Mar 5, 2026

Uh oh!

kosiew Mar 5, 2026

Uh oh!

kosiew Mar 5, 2026

Uh oh!

kosiew Mar 5, 2026

Uh oh!

ethan-tyler commented Mar 17, 2026

Uh oh!

kosiew commented Mar 18, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

ethan-tyler commented Feb 15, 2026

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?

Uh oh!

ethan-tyler commented Feb 17, 2026

Uh oh!

kosiew left a comment

Choose a reason for hiding this comment

Uh oh!

kosiew Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

kosiew Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

kosiew Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

kosiew Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

kosiew Mar 5, 2026

Choose a reason for hiding this comment

Uh oh!

ethan-tyler commented Mar 17, 2026

Uh oh!

kosiew commented Mar 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

kosiew commented Mar 18, 2026 •

edited

Loading