Rewrite FileStream in terms of Morsel API by alamb · Pull Request #21342 · apache/datafusion

alamb · 2026-04-03T13:32:24Z

Stacked on

Which issue does this PR close?

part of Introduce morsel-driven Parquet scan #20529
Broken out of Sketch out a Morselize API #20820

Rationale for this change

The Morsel API allows for finer grain parallelism (and IO). It is important to have the FileStream work in terms of the Morsel API to allow future features (like workstealing, etc)

What changes are included in this PR?

I apologize for the large diff; Note about 1/2 of this PR is tests and a test framework to test the calling sequence of FileStream.

Rewrite FileStream in terms of the MorselAPI
Add snapshot driven test to document the I/O and CPU patterns in FileStream
Add snapshot based tests that show the ordering of files

Are these changes tested?

Yes by existing functional and benchmark tests, as well as new functional snapshot based tests

Are there any user-facing changes?

No (not yet)

alamb · 2026-04-03T16:15:13Z

datafusion/datasource/src/file_stream/scan_state.rs

+/// This groups together ready planners, ready morsels, the active reader,
+/// pending planner I/O, the remaining files and limit, and the metrics
+/// associated with processing that work.
+pub(super) struct ScanState {


This is the new inner state machine for FileStream

I think some more diagrams in the docstring of the struct and/or fields could help. I'm trying to wrap my head around how the IO queue and such work.

I have added a diagram - let me know if that helps or if there is something else I can do

alamb · 2026-04-03T16:15:37Z

datafusion/datasource/src/morsel/adapters.rs

+use std::sync::Arc;
+use std::sync::mpsc::{self, Receiver, TryRecvError};
+
+/// Adapt a legacy [`FileOpener`] to the morsel API.


This is an adapter so that existing FileOpeners continue to have the same behavior

alamb · 2026-04-03T16:15:57Z

datafusion/datasource/src/morsel/mocks.rs

@@ -0,0 +1,556 @@
+// Licensed to the Apache Software Foundation (ASF) under one


This is testing infrastructure to write the snapshot tests

Basically it makes a mock morselizer that records its steps so that the control flow of FileStream can be tested / verified

alamb · 2026-04-03T16:16:52Z

datafusion/datasource/src/file_stream/mod.rs

-                                    return Poll::Ready(Some(Err(err)));
-                                }
-                            }
+                FileStreamState::Scan { scan_state: queue } => {


moved the inner state machine into a separate module/struct to try and keep indenting under control and encapsualte the complexity somewhat

alamb · 2026-04-03T16:18:30Z

datafusion/datasource/src/file_stream/mod.rs

        assert!(err.contains("FileStreamBuilder invalid partition index: 1"));
    }
+
+    /// Verifies the simplest morsel-driven flow: one planner produces one


Here are tests showing the sequence of calls to the various morsel APIs. I intend to use this framework to show how work can migrate from one stream to the other

alamb · 2026-04-03T17:36:36Z

datafusion/datasource/Cargo.toml

 all-features = true

 [features]
+backtrace = ["datafusion-common/backtrace"]


I added this while debugging why the tests failed on CI and not locally (it was when this feature flag was on the Error messages got mangled).

I added a crate level feature to enable the feature in datafusion-common so I could reproduce the error locally

adriangb

Ran out of time for the last couple of files. A lot of the comments are just tracking my thought process, I plan to go over them again to clarify my own understanding but maybe they're helpful as input on how the code reads top to bottom for a first time reader.

adriangb · 2026-04-03T16:26:53Z

datafusion/datasource/src/file.rs

+    /// Creates a `dyn Morselizer` based on given parameters.
+    ///
+    /// The default implementation preserves existing behavior by adapting the
+    /// legacy [`FileOpener`] API into a [`Morselizer`].
+    ///
+    /// It is preferred to implement the [`Morselizer`] API directly by
+    /// implementing this method.
+    fn create_morselizer(
+        &self,
+        object_store: Arc<dyn ObjectStore>,
+        base_config: &FileScanConfig,
+        partition: usize,
+    ) -> Result<Box<dyn Morselizer>> {
+        let opener = self.create_file_opener(object_store, base_config, partition)?;
+        Ok(Box::new(FileOpenerMorselizer::new(opener)))
+    }


datafusion/datasource-parquet/src/source.rs

adriangb · 2026-04-04T13:53:09Z

datafusion/datasource/src/file_stream/builder.rs

    /// Configure the [`FileOpener`] used to open files.
+    ///
+    /// This will overwrite any setting from [`Self::with_morselizer`]
    pub fn with_file_opener(mut self, file_opener: Arc<dyn FileOpener>) -> Self {


While I think it could make sense to keep FileOpener as a public API for building data sources (if we consider it simpler, for folks who don't care about perf), this method in particular seems like a mostly internal method (even if it is pub) on we might as well deprecate / remove.

This method is the way we could keep using FileOpener (as it is simpler)

I am not sure how we could still allow using FileOpener but not keep this method

datafusion/datasource/src/file_stream/metrics.rs

datafusion/datasource/src/file_stream/mod.rs

adriangb · 2026-04-04T14:37:24Z

datafusion/datasource/src/file_stream/scan_state.rs

+    /// The active reader, if any.
+    reader: Option<BoxStream<'static, Result<RecordBatch>>>,


Is there one ScanState across all partitions or one per partition? I'm guessing the latter: file_iter: VecDeque<PartitionedFile> is the files for this partition, we pump all of the files into one output stream of RecordBatch (reader). But we can have multiple planners / morsels ready and merge those all into a single stream of RecordBatch on the way out.

One per partition

we pump all of the files into one output stream of RecordBatch (reader). But we can have multiple planners / morsels ready and merge those all into a single stream of RecordBatch on the way out.

My initial proposal (following @Dandandan 's original design" is that when possible the files are put into a shared queue so that when a FileStream is ready it gets the next file

I think once we get that structure in place, we can contemplate more sophisticated designs (like one filestream preparing a parquet file, and then divying up the record batches between other cores)

Is there one ScanState across all partitions or one per partition? I'm guessing the latter: file_iter: VecDeque is the files for this partition

yes, it is one ScanState per partition

we pump all of the files into one output stream of RecordBatch (reader). But we can have multiple planners / morsels ready and merge those all into a single stream of RecordBatch on the way out.

Yes this is right

My initial proposal (following @Dandandan 's original design" is that when possible the files are put into a shared queue so that when a FileStream is ready it gets the next file

yes, it is one ScanState per partition

I'm a bit confused then: if there is one ScanState per partition then there is one VecDeque<PartitionedFile>, which means it's not shared between partitions. But that would contradict
"files are put into a shared queue so that when a FileStream is ready it gets the next file"?

You can see how cross stream sharing works in the next stacked PR:

Dynamic work scheduling in FileStream #21351

the ScanState is not shared across partitions, but it has a new work_queue that is (potentially) shared. The relevant change is to replace the file_iter with this work_source thing and then handle setting up the work_source in the DataSource exec

pub(super) struct ScanState { /// Files that still need to be planned. file_iter: VecDeque<PartitionedFile>, ...

Wth

pub(super) struct ScanState { /// Unopened files that still need to be planned for this stream. work_source: WorkSource, ...

datafusion/datasource/src/file_stream/scan_state.rs

alamb · 2026-04-06T11:47:57Z

Thanks @adriangb -- I think I am now convinced this will make things faster (see #21351)

Once I finalze that I will then go back and ping / respond to each of these PRs in turn

alamb · 2026-04-06T12:56:54Z

Ok the first PR in the chain is ready for review:

Reduce SortExec memory usage by void constructing single huge batch #2132

(that is basically 50% of this PR)

alamb

Thanks for the comments @adriangb

Once I get this PR merged

#21327

I will come back and update this PR / address your comments

alamb · 2026-04-06T12:58:17Z

datafusion/datasource/src/file_stream/builder.rs

    /// Configure the [`FileOpener`] used to open files.
+    ///
+    /// This will overwrite any setting from [`Self::with_morselizer`]
    pub fn with_file_opener(mut self, file_opener: Arc<dyn FileOpener>) -> Self {


This method is the way we could keep using FileOpener (as it is simpler)

I am not sure how we could still allow using FileOpener but not keep this method

datafusion/datasource/src/file_stream/scan_state.rs

alamb · 2026-04-06T13:05:16Z

datafusion/datasource/src/file_stream/scan_state.rs

+    /// The active reader, if any.
+    reader: Option<BoxStream<'static, Result<RecordBatch>>>,


One per partition

we pump all of the files into one output stream of RecordBatch (reader). But we can have multiple planners / morsels ready and merge those all into a single stream of RecordBatch on the way out.

My initial proposal (following @Dandandan 's original design" is that when possible the files are put into a shared queue so that when a FileStream is ready it gets the next file

I think once we get that structure in place, we can contemplate more sophisticated designs (like one filestream preparing a parquet file, and then divying up the record batches between other cores)

datafusion/datasource/src/file_stream/scan_state.rs

alamb · 2026-04-10T16:33:54Z

My summary of the benchmark results is that this PR does not change performance characteristics (as expected)

The changes come in #21351

alamb · 2026-04-11T09:04:08Z

datafusion/datasource-parquet/src/opener.rs

+    /// This mirrors how `FileStream` consumes the morsel APIs: it repeatedly
+    /// plans CPU work, awaits any discovered I/O futures, and feeds the planner
+    /// back into the ready queue until a stream morsel is ready.
+    async fn open_file(


this is needed so that all the tests in this file don't need to be rewritten to the morselizer API

alamb · 2026-04-11T09:04:31Z

datafusion/datasource-parquet/src/opener.rs

-                    + Send,
-            >,
-        >,
+        mut stream: BoxStream<'static, Result<RecordBatch>>,


just a simpler way to write this

alamb · 2026-04-11T09:05:11Z

datafusion/datasource-parquet/src/source.rs

-            },
-        });
-        Ok(opener)
+        Ok(Box::new(ParquetMorselizer {


this passes the new morselizer directly rather than wrapping it first

alamb · 2026-04-11T09:05:50Z

datafusion/datasource/src/file_stream/scan_state.rs

+/// This groups together ready planners, ready morsels, the active reader,
+/// pending planner I/O, the remaining files and limit, and the metrics
+/// associated with processing that work.
+pub(super) struct ScanState {


I have added a diagram - let me know if that helps or if there is something else I can do

alamb · 2026-04-11T09:06:15Z

datafusion/datasource/src/file_stream/scan_state.rs

+    ///
+    /// The return [`ScanAndReturn`] tells `poll_inner` how to update the
+    /// outer `FileStreamState`.
+    pub(super) fn poll_scan(&mut self, cx: &mut Context<'_>) -> ScanAndReturn {


This is the main state machine loop that shows how the file stream transitions between states

alamb · 2026-04-11T09:07:02Z

datafusion/datasource/src/morsel/mocks.rs

@@ -0,0 +1,556 @@
+// Licensed to the Apache Software Foundation (ASF) under one


Basically it makes a mock morselizer that records its steps so that the control flow of FileStream can be tested / verified

datafusion/datasource/src/file_stream/mod.rs

adriangb · 2026-04-12T21:17:19Z

datafusion/datasource/src/file_stream/scan_state.rs

+                // Morsels should ideally only expose ready-to-decode streams,
+                // but tolerate pending readers here.


Is this another case where we could better encode the behavior constraints into the type system?

That is a good question... I think we would have to change the inner API to use something other than Stream (perhaps just an iterator). I'll see what I can come up with

Ok, an update here. I tried this out and hit some snags -- see here for deatils

Support Morsel output for Parquet known to be non blocking #21598

adriangb · 2026-04-12T21:19:20Z

datafusion/datasource/src/file_stream/scan_state.rs

+                    self.ready_morsels.extend(plan.take_morsels());
+                    self.ready_planners.extend(plan.take_ready_planners());
+                    if let Some(pending_planner) = plan.take_pending_planner() {
+                        self.pending_planner = Some(pending_planner);


Should we at least assert the invariant assert!(self.pending_planner.is_none())? Or should we have a queue (VecDequeu) of pending planners?

THis is a good call -- I think a queue of pending planners is best. Will do

I played around with this and it turns out to be non trivial (mostly due to testing) and the current parquet morselizer doesn't actually return multiple child planners (yet)

So for this PR I would like to proceed with just an assert and I will leave some comments / a note about supporting multiple returned planners in the future

I played around with this more and changed the core loop here to make things better in 42d8750

The current design of FileStream only allows for a single outstanding IO (to avoid having to deal with backpressure, etc)

adriangb · 2026-04-12T21:21:30Z

datafusion/datasource/src/file_stream/scan_state.rs

+    /// The active reader, if any.
+    reader: Option<BoxStream<'static, Result<RecordBatch>>>,


My initial proposal (following @Dandandan 's original design" is that when possible the files are put into a shared queue so that when a FileStream is ready it gets the next file

yes, it is one ScanState per partition

I'm a bit confused then: if there is one ScanState per partition then there is one VecDeque<PartitionedFile>, which means it's not shared between partitions. But that would contradict
"files are put into a shared queue so that when a FileStream is ready it gets the next file"?

Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>

alamb

Thank you @adriangb -- I am working on your comments

alamb · 2026-04-13T12:42:15Z

datafusion/datasource/src/file_stream/scan_state.rs

+    /// The active reader, if any.
+    reader: Option<BoxStream<'static, Result<RecordBatch>>>,


You can see how cross stream sharing works in the next stacked PR:

Dynamic work scheduling in FileStream #21351

the ScanState is not shared across partitions, but it has a new work_queue that is (potentially) shared. The relevant change is to replace the file_iter with this work_source thing and then handle setting up the work_source in the DataSource exec

pub(super) struct ScanState { /// Files that still need to be planned. file_iter: VecDeque<PartitionedFile>, ...

Wth

pub(super) struct ScanState { /// Unopened files that still need to be planned for this stream. work_source: WorkSource, ...

alamb · 2026-04-13T12:43:15Z

datafusion/datasource/src/file_stream/scan_state.rs

+                // Morsels should ideally only expose ready-to-decode streams,
+                // but tolerate pending readers here.


That is a good question... I think we would have to change the inner API to use something other than Stream (perhaps just an iterator). I'll see what I can come up with

alamb · 2026-04-13T12:44:40Z

datafusion/datasource/src/file_stream/scan_state.rs

+                    self.ready_morsels.extend(plan.take_morsels());
+                    self.ready_planners.extend(plan.take_ready_planners());
+                    if let Some(pending_planner) = plan.take_pending_planner() {
+                        self.pending_planner = Some(pending_planner);


THis is a good call -- I think a queue of pending planners is best. Will do

alamb · 2026-04-13T15:10:58Z

I have also updated the testing mocks to more closely follow the morsel API (so that I can test the suggestions from @adriangb )

I ran out of time now, but hopefully soon I'll move on to trying to explore:

Is this another case where we could better encode the behavior constraints into the type system?

adriangb · 2026-04-13T15:18:06Z

Let me know when you want me to do the next (which feels like it might be the final) round of review.

alamb · 2026-04-13T15:35:59Z

Let me know when you want me to do the next (which feels like it might be the final) round of review.

Thanks -- I haven't explored this more yet:

Is this another case where we could better encode the behavior constraints into the type system?

But I can also do it as part of a follow on PR (I think it would need to change the Morsel API)

I don't think it really changes this PR per se

adriangb

Makes sense to me. We can iterate on the APIs as a followup but I think we should keep that in mind, it feels like there's some improvements we can make.

alamb · 2026-04-13T16:18:53Z

Makes sense to me. We can iterate on the APIs as a followup but I think we should keep that in mind, it feels like there's some improvements we can make.

Yes it is an excellent point actually -- and one I think we can resolve

alamb · 2026-04-13T16:19:12Z

FYI @Dandandan @zhuliquan and @xudong963 in case you would like to review

alamb · 2026-04-14T11:22:36Z

Makes sense to me. We can iterate on the APIs as a followup but I think we should keep that in mind, it feels like there's some improvements we can make.

Yes it is an excellent point actually -- and one I think we can resolve

Update is I filed a ticket to explain what I found here

Support Morsel output for Parquet known to be non blocking #21598

alamb · 2026-04-14T11:23:03Z

I'll merge this one in and get the final one ready for review

alamb · 2026-04-14T11:23:11Z

Thank you for helping this along @adriangb

Dandandan · 2026-04-14T15:15:07Z

Woah merged!!! 🥳 🥳 🥳

github-actions bot added the datasource Changes to the datasource crate label Apr 3, 2026

alamb force-pushed the alamb/file_stream_split branch 3 times, most recently from 816d243 to 3346af7 Compare April 3, 2026 16:14

alamb commented Apr 3, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

alamb mentioned this pull request Apr 3, 2026

Dynamic work scheduling in FileStream #21351

Open

alamb force-pushed the alamb/file_stream_split branch from b5c452a to d5a1f74 Compare April 3, 2026 17:34

This comment has been minimized.

Sign in to view

alamb commented Apr 3, 2026

View reviewed changes

This comment has been minimized.

Sign in to view

alamb force-pushed the alamb/file_stream_split branch from d5a1f74 to b2c9bd6 Compare April 3, 2026 17:38

This comment has been minimized.

Sign in to view

adriangb reviewed Apr 4, 2026

View reviewed changes

alamb mentioned this pull request Apr 6, 2026

Introduce Morselizer API, rewrite ParquetOpener to ParquetMorselizer #21327

Merged

This was referenced Apr 6, 2026

Sketch out a Morselize API #20820

Closed

Introduce morsel-driven Parquet scan #20529

Open

alamb commented Apr 6, 2026

View reviewed changes

alamb requested a review from adriangb April 10, 2026 15:31

alamb commented Apr 11, 2026

View reviewed changes

adriangb reviewed Apr 12, 2026

View reviewed changes

alamb and others added 2 commits April 13, 2026 08:35

Apply suggestions from code review from @adriangb

90fe8e1

Co-authored-by: Adrian Garcia Badaracco <1755071+adriangb@users.noreply.github.com>

Merge remote-tracking branch 'apache/main' into alamb/file_stream_split

40b027c

alamb commented Apr 13, 2026

View reviewed changes

alamb added 7 commits April 13, 2026 08:49

fix compilation

671938b

Make it clearer that there is only ever a single outstanding IO

42d8750

Add test for pending IO

8b591ac

Update mocks to better match API

e7913da

more updates

7c0ca4f

Simplify testing

7d33eb8

Merge remote-tracking branch 'apache/main' into alamb/file_stream_split

ed9873c

new test

04dbbbf

adriangb approved these changes Apr 13, 2026

View reviewed changes

alamb mentioned this pull request Apr 13, 2026

Support Morsel output for Parquet known to be non blocking #21598

Open

alamb added this pull request to the merge queue Apr 14, 2026

Merged via the queue into apache:main with commit 776b723 Apr 14, 2026
34 checks passed

alamb deleted the alamb/file_stream_split branch April 14, 2026 11:25

		@@ -0,0 +1,556 @@
		// Licensed to the Apache Software Foundation (ASF) under one

		/// The active reader, if any.
		reader: Option<BoxStream<'static, Result<RecordBatch>>>,

		// Morsels should ideally only expose ready-to-decode streams,
		// but tolerate pending readers here.

Conversation

alamb commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

adriangb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alamb commented Apr 6, 2026

Uh oh!

alamb commented Apr 6, 2026

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alamb commented Apr 3, 2026 •

edited

Loading

alamb Apr 13, 2026 •

edited

Loading