You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/04/09 15:51:28 UTC
(datafusion) branch gh-readonly-queue/main/pr-21327-6a770aa49de6ef3746c72838d2a044998e9535e0 created (now 070d1a78b5)
This is an automated email from the ASF dual-hosted git repository.
github-bot pushed a change to branch gh-readonly-queue/main/pr-21327-6a770aa49de6ef3746c72838d2a044998e9535e0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 070d1a78b5 Introduce Morselizer API, rewrite `ParquetOpener` to `ParquetMorselizer` (#21327)
This branch includes the following new commits:
new 070d1a78b5 Introduce Morselizer API, rewrite `ParquetOpener` to `ParquetMorselizer` (#21327)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: Introduce Morselizer API, rewrite `ParquetOpener` to `ParquetMorselizer` (#21327)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-bot pushed a commit to branch gh-readonly-queue/main/pr-21327-6a770aa49de6ef3746c72838d2a044998e9535e0
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 070d1a78b5fc1e3d20a7c698c18999cecfb67d8f
Author: Andrew Lamb <an...@nerdnetworks.org>
AuthorDate: Thu Apr 9 11:51:07 2026 -0400
Introduce Morselizer API, rewrite `ParquetOpener` to `ParquetMorselizer` (#21327)
~(Draft until I am sure I can use this API to make FileStream behave
better)~
## Which issue does this PR close?
- part of https://github.com/apache/datafusion/issues/20529
- Needed for https://github.com/apache/datafusion/pull/21351
- Broken out of https://github.com/apache/datafusion/pull/20820
- Closes https://github.com/apache/datafusion/pull/21427
## Rationale for this change
I can get 10% faster on many ClickBench queries by reordeirng files at
runtime. You can see it all working together here:
https://github.com/apache/datafusion/pull/21351
To do do, I need to rework the FileStream so that it can reorder
operations at runtime. Eventually that will include both CPU and IO.
This PR is a step in the direction by introducing the main Morsel API
and implementing it for Parquet. The next PR
(https://github.com/apache/datafusion/pull/21342) rewrites FileStream in
terms of the Morsel API
## What changes are included in this PR?
1. Add proposed `Morsel` API
2. Rewrite Parquet opener in terms of that API
3. Add an adapter layer (back to FileOpener, so I don't have to rewrite
FileStream in the same PR)
My next PR will rewrite the FileStream to use the Morsel API
## Are these changes tested?
Yes by existing CI.
I will work on adding additional tests for just Parquet opener in a
follow on PR
## Are there any user-facing changes?
No
---
datafusion/datasource-parquet/src/opener.rs | 357 ++++++++++++++++++++--------
datafusion/datasource-parquet/src/source.rs | 56 ++---
datafusion/datasource/src/mod.rs | 1 +
datafusion/datasource/src/morsel/mod.rs | 229 ++++++++++++++++++
4 files changed, 512 insertions(+), 131 deletions(-)
diff --git a/datafusion/datasource-parquet/src/opener.rs b/datafusion/datasource-parquet/src/opener.rs
index 6621706c35..35900e16c1 100644
--- a/datafusion/datasource-parquet/src/opener.rs
+++ b/datafusion/datasource-parquet/src/opener.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-//! [`ParquetOpener`] state machine for opening Parquet files
+//! [`ParquetOpener`] and [`ParquetMorselizer`] state machines for opening Parquet files
use crate::page_filter::PagePruningAccessPlanFilter;
use crate::row_filter::build_projection_read_plan;
@@ -26,11 +26,16 @@ use crate::{
};
use arrow::array::{RecordBatch, RecordBatchOptions};
use arrow::datatypes::DataType;
+use datafusion_common::internal_err;
use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
+use datafusion_datasource::morsel::{
+ Morsel, MorselPlan, MorselPlanner, Morselizer, PendingMorselPlanner,
+};
use datafusion_physical_expr::projection::{ProjectionExprs, Projector};
use datafusion_physical_expr::utils::reassign_expr_columns;
use datafusion_physical_expr_adapter::replace_columns_with_literals;
-use std::collections::HashMap;
+use std::collections::{HashMap, VecDeque};
+use std::fmt;
use std::future::Future;
use std::mem;
use std::pin::Pin;
@@ -77,12 +82,26 @@ use parquet::bloom_filter::Sbbf;
use parquet::errors::ParquetError;
use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader};
-/// Entry point for opening a Parquet file
+/// Implements [`FileOpener`] for Parquet
+#[derive(Clone)]
+pub(super) struct ParquetOpener {
+ pub(super) morselizer: ParquetMorselizer,
+}
+
+impl FileOpener for ParquetOpener {
+ fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
+ let future = ParquetOpenFuture::new(&self.morselizer, partitioned_file)?;
+ Ok(Box::pin(future))
+ }
+}
+
+/// Stateless Parquet morselizer implementation.
///
/// Reading a Parquet file is a multi-stage process, with multiple CPU-intensive
/// steps interspersed with I/O steps. The code in this module implements the steps
/// as an explicit state machine -- see [`ParquetOpenState`] for details.
-pub(super) struct ParquetOpener {
+#[derive(Clone)]
+pub(super) struct ParquetMorselizer {
/// Execution partition index
pub(crate) partition_index: usize,
/// Projection to apply on top of the table schema (i.e. can reference partition columns).
@@ -137,6 +156,23 @@ pub(super) struct ParquetOpener {
pub reverse_row_groups: bool,
}
+impl fmt::Debug for ParquetMorselizer {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("ParquetMorselizer")
+ .field("partition_index", &self.partition_index)
+ .field("preserve_order", &self.preserve_order)
+ .field("enable_page_index", &self.enable_page_index)
+ .field("enable_bloom_filter", &self.enable_bloom_filter)
+ .finish()
+ }
+}
+
+impl Morselizer for ParquetMorselizer {
+ fn plan_file(&self, file: PartitionedFile) -> Result<Box<dyn MorselPlanner>> {
+ Ok(Box::new(ParquetMorselPlanner::try_new(self, file)?))
+ }
+}
+
/// States for [`ParquetOpenFuture`]
///
/// These states correspond to the steps required to read and apply various
@@ -216,6 +252,27 @@ enum ParquetOpenState {
Done,
}
+impl fmt::Debug for ParquetOpenState {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ let state = match self {
+ ParquetOpenState::Start { .. } => "Start",
+ #[cfg(feature = "parquet_encryption")]
+ ParquetOpenState::LoadEncryption(_) => "LoadEncryption",
+ ParquetOpenState::PruneFile(_) => "PruneFile",
+ ParquetOpenState::LoadMetadata(_) => "LoadMetadata",
+ ParquetOpenState::PrepareFilters(_) => "PrepareFilters",
+ ParquetOpenState::LoadPageIndex(_) => "LoadPageIndex",
+ ParquetOpenState::PruneWithStatistics(_) => "PruneWithStatistics",
+ ParquetOpenState::LoadBloomFilters(_) => "LoadBloomFilters",
+ ParquetOpenState::PruneWithBloomFilters(_) => "PruneWithBloomFilters",
+ ParquetOpenState::BuildStream(_) => "BuildStream",
+ ParquetOpenState::Ready(_) => "Ready",
+ ParquetOpenState::Done => "Done",
+ };
+ f.write_str(state)
+ }
+}
+
struct PreparedParquetOpen {
partition_index: usize,
partitioned_file: PartitionedFile,
@@ -290,37 +347,13 @@ struct BloomFiltersLoadedParquetOpen {
row_group_bloom_filters: Vec<BloomFilterStatistics>,
}
-/// Implements state machine described in [`ParquetOpenState`]
-struct ParquetOpenFuture {
- state: ParquetOpenState,
-}
-
-impl ParquetOpenFuture {
- #[cfg(feature = "parquet_encryption")]
- fn new(prepared: PreparedParquetOpen, encryption_context: EncryptionContext) -> Self {
- Self {
- state: ParquetOpenState::Start {
- prepared: Box::new(prepared),
- encryption_context: Arc::new(encryption_context),
- },
- }
- }
-
- #[cfg(not(feature = "parquet_encryption"))]
- fn new(prepared: PreparedParquetOpen) -> Self {
- Self {
- state: ParquetOpenState::Start {
- prepared: Box::new(prepared),
- },
- }
- }
-}
-
impl ParquetOpenState {
/// Applies one CPU-only state transition.
///
/// `Load*` states do not transition here and are returned unchanged so the
/// driver loop can poll their inner futures separately.
+ ///
+ /// Implements state machine described in [`ParquetOpenState`]
fn transition(self) -> Result<ParquetOpenState> {
match self {
ParquetOpenState::Start {
@@ -392,93 +425,208 @@ impl ParquetOpenState {
}
}
+/// Adapter for a [`MorselPlanner`] to the [`FileOpener`] API
+///
+/// Compatibility adapter that drives a morsel planner through the
+/// [`FileOpener`] API.
+struct ParquetOpenFuture {
+ planner: Option<Box<dyn MorselPlanner>>,
+ pending_io: Option<PendingMorselPlanner>,
+ ready_morsels: VecDeque<Box<dyn Morsel>>,
+}
+
+impl ParquetOpenFuture {
+ fn new(
+ morselizer: &ParquetMorselizer,
+ partitioned_file: PartitionedFile,
+ ) -> Result<Self> {
+ Ok(Self {
+ planner: Some(morselizer.plan_file(partitioned_file)?),
+ pending_io: None,
+ ready_morsels: VecDeque::new(),
+ })
+ }
+}
+
impl Future for ParquetOpenFuture {
type Output = Result<BoxStream<'static, Result<RecordBatch>>>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
loop {
- let state = mem::replace(&mut self.state, ParquetOpenState::Done);
- let mut state = state.transition()?;
-
- match state {
- #[cfg(feature = "parquet_encryption")]
- ParquetOpenState::LoadEncryption(mut future) => {
- state = match future.poll_unpin(cx) {
- Poll::Ready(result) => ParquetOpenState::PruneFile(result?),
- Poll::Pending => {
- self.state = ParquetOpenState::LoadEncryption(future);
- return Poll::Pending;
- }
- };
- }
- ParquetOpenState::LoadMetadata(mut future) => {
- state = match future.poll_unpin(cx) {
- Poll::Ready(result) => {
- ParquetOpenState::PrepareFilters(Box::new(result?))
- }
- Poll::Pending => {
- self.state = ParquetOpenState::LoadMetadata(future);
- return Poll::Pending;
- }
- };
- }
- ParquetOpenState::LoadPageIndex(mut future) => {
- state = match future.poll_unpin(cx) {
- Poll::Ready(result) => {
- ParquetOpenState::PruneWithStatistics(Box::new(result?))
- }
- Poll::Pending => {
- self.state = ParquetOpenState::LoadPageIndex(future);
- return Poll::Pending;
- }
- };
- }
- ParquetOpenState::LoadBloomFilters(mut future) => {
- state = match future.poll_unpin(cx) {
- Poll::Ready(result) => {
- ParquetOpenState::PruneWithBloomFilters(Box::new(result?))
- }
- Poll::Pending => {
- self.state = ParquetOpenState::LoadBloomFilters(future);
- return Poll::Pending;
- }
- };
- }
- ParquetOpenState::Ready(stream) => {
- return Poll::Ready(Ok(stream));
- }
- ParquetOpenState::Done => {
- return Poll::Ready(Ok(futures::stream::empty().boxed()));
+ // If planner I/O completed, resume with the returned planner.
+ if let Some(io_future) = self.pending_io.as_mut() {
+ let maybe_planner = ready!(io_future.poll_unpin(cx));
+ // Clear `pending_io` before handling the result so an error
+ // cannot leave both continuation paths populated.
+ self.pending_io = None;
+ if self.planner.is_some() {
+ return Poll::Ready(internal_err!(
+ "ParquetOpenFuture does not support concurrent planners"
+ ));
}
+ self.planner = Some(maybe_planner?);
+ }
+
+ // If a stream morsel is ready, return it.
+ if let Some(morsel) = self.ready_morsels.pop_front() {
+ return Poll::Ready(Ok(morsel.into_stream()));
+ }
- // For all other states, loop again and try to transition
- // immediately. All states are explicitly listed here to ensure any
- // new states are handled correctly
- ParquetOpenState::Start { .. } => {}
- ParquetOpenState::PruneFile(_) => {}
- ParquetOpenState::PrepareFilters(_) => {}
- ParquetOpenState::PruneWithStatistics(_) => {}
- ParquetOpenState::PruneWithBloomFilters(_) => {}
- ParquetOpenState::BuildStream(_) => {}
+ // This shim must always own either a planner, a pending planner
+ // future, or a ready morsel. Reaching this branch means the
+ // continuation was lost.
+ let Some(planner) = self.planner.take() else {
+ return Poll::Ready(internal_err!(
+ "ParquetOpenFuture polled after completion"
+ ));
};
- self.state = state;
+ // Planner completed without producing a stream morsel.
+ // (e.g. all row groups were pruned)
+ let Some(mut plan) = planner.plan()? else {
+ return Poll::Ready(Ok(futures::stream::empty().boxed()));
+ };
+
+ let mut child_planners = plan.take_ready_planners();
+ if child_planners.len() > 1 {
+ return Poll::Ready(internal_err!(
+ "Parquet FileOpener adapter does not support child morsel planners"
+ ));
+ }
+ self.planner = child_planners.pop();
+
+ self.ready_morsels = plan.take_morsels().into();
+
+ if let Some(io_future) = plan.take_pending_planner() {
+ self.pending_io = Some(io_future);
+ }
}
}
}
-impl FileOpener for ParquetOpener {
- fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> {
- let prepared = self.prepare_open_file(partitioned_file)?;
+/// Implements the Morsel API
+struct ParquetStreamMorsel {
+ stream: BoxStream<'static, Result<RecordBatch>>,
+}
+
+impl ParquetStreamMorsel {
+ fn new(stream: BoxStream<'static, Result<RecordBatch>>) -> Self {
+ Self { stream }
+ }
+}
+
+impl fmt::Debug for ParquetStreamMorsel {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("ParquetStreamMorsel")
+ .finish_non_exhaustive()
+ }
+}
+
+impl Morsel for ParquetStreamMorsel {
+ fn into_stream(self: Box<Self>) -> BoxStream<'static, Result<RecordBatch>> {
+ self.stream
+ }
+}
+
+/// Per-file planner that owns the current [`ParquetOpenState`].
+struct ParquetMorselPlanner {
+ /// Ready to perform CPU-only planning work.
+ state: ParquetOpenState,
+}
+
+impl fmt::Debug for ParquetMorselPlanner {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_tuple("ParquetMorselPlanner::Ready")
+ .field(&self.state)
+ .finish()
+ }
+}
+
+impl ParquetMorselPlanner {
+ fn try_new(morselizer: &ParquetMorselizer, file: PartitionedFile) -> Result<Self> {
+ let prepared = morselizer.prepare_open_file(file)?;
#[cfg(feature = "parquet_encryption")]
- let future = ParquetOpenFuture::new(prepared, self.get_encryption_context());
+ let state = ParquetOpenState::Start {
+ prepared: Box::new(prepared),
+ encryption_context: Arc::new(morselizer.get_encryption_context()),
+ };
#[cfg(not(feature = "parquet_encryption"))]
- let future = ParquetOpenFuture::new(prepared);
- Ok(Box::pin(future))
+ let state = ParquetOpenState::Start {
+ prepared: Box::new(prepared),
+ };
+ Ok(Self { state })
+ }
+
+ /// Schedule an I/O future that resolves to the next planner to run.
+ ///
+ /// This helper
+ ///
+ /// 1. drives one I/O phase to completion
+ /// 2. wraps the resulting state in a new [`ParquetMorselPlanner`]
+ /// 3. returns a [`MorselPlan`] containing the boxed future for the caller
+ /// to poll
+ ///
+ fn schedule_io<F>(future: F) -> MorselPlan
+ where
+ F: Future<Output = Result<ParquetOpenState>> + Send + 'static,
+ {
+ let io_future = async move {
+ let next_state = future.await?;
+ Ok(Box::new(ParquetMorselPlanner { state: next_state }) as _)
+ };
+ MorselPlan::new().with_pending_planner(io_future)
+ }
+}
+
+impl MorselPlanner for ParquetMorselPlanner {
+ fn plan(self: Box<Self>) -> Result<Option<MorselPlan>> {
+ if let ParquetOpenState::Done = self.state {
+ return Ok(None);
+ }
+
+ let state = self.state.transition()?;
+
+ match state {
+ #[cfg(feature = "parquet_encryption")]
+ ParquetOpenState::LoadEncryption(future) => {
+ Ok(Some(Self::schedule_io(async move {
+ Ok(ParquetOpenState::PruneFile(future.await?))
+ })))
+ }
+ ParquetOpenState::LoadMetadata(future) => {
+ Ok(Some(Self::schedule_io(async move {
+ Ok(ParquetOpenState::PrepareFilters(Box::new(future.await?)))
+ })))
+ }
+ ParquetOpenState::LoadPageIndex(future) => {
+ Ok(Some(Self::schedule_io(async move {
+ Ok(ParquetOpenState::PruneWithStatistics(Box::new(
+ future.await?,
+ )))
+ })))
+ }
+ ParquetOpenState::LoadBloomFilters(future) => {
+ Ok(Some(Self::schedule_io(async move {
+ Ok(ParquetOpenState::PruneWithBloomFilters(Box::new(
+ future.await?,
+ )))
+ })))
+ }
+ ParquetOpenState::Ready(stream) => {
+ let morsels: Vec<Box<dyn Morsel>> =
+ vec![Box::new(ParquetStreamMorsel::new(stream))];
+ Ok(Some(MorselPlan::new().with_morsels(morsels)))
+ }
+ ParquetOpenState::Done => Ok(None),
+ cpu_state => Ok(Some(
+ MorselPlan::new()
+ .with_planners(vec![Box::new(Self { state: cpu_state })]),
+ )),
+ }
}
}
-impl ParquetOpener {
+impl ParquetMorselizer {
/// Perform the CPU-only setup for opening a parquet file.
fn prepare_open_file(
&self,
@@ -1447,7 +1595,7 @@ impl EncryptionContext {
}
}
-impl ParquetOpener {
+impl ParquetMorselizer {
#[cfg(feature = "parquet_encryption")]
fn get_encryption_context(&self) -> EncryptionContext {
EncryptionContext::new(
@@ -1576,7 +1724,7 @@ fn should_enable_page_index(
mod test {
use std::sync::Arc;
- use super::{ConstantColumns, constant_columns_from_stats};
+ use super::{ConstantColumns, ParquetMorselizer, constant_columns_from_stats};
use crate::{DefaultParquetFileReaderFactory, RowGroupAccess, opener::ParquetOpener};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use bytes::{BufMut, BytesMut};
@@ -1731,11 +1879,12 @@ mod test {
ProjectionExprs::from_indices(&all_indices, &file_schema)
};
- ParquetOpener {
+ let morselizer = ParquetMorselizer {
partition_index: self.partition_index,
projection,
batch_size: self.batch_size,
limit: self.limit,
+ preserve_order: self.preserve_order,
predicate: self.predicate,
table_schema,
metadata_size_hint: self.metadata_size_hint,
@@ -1757,8 +1906,8 @@ mod test {
encryption_factory: None,
max_predicate_cache_size: self.max_predicate_cache_size,
reverse_row_groups: self.reverse_row_groups,
- preserve_order: self.preserve_order,
- }
+ };
+ ParquetOpener { morselizer }
}
}
diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs
index 3a64137a2a..1e54e98dfd 100644
--- a/datafusion/datasource-parquet/src/source.rs
+++ b/datafusion/datasource-parquet/src/source.rs
@@ -23,8 +23,8 @@ use std::sync::Arc;
use crate::DefaultParquetFileReaderFactory;
use crate::ParquetFileReaderFactory;
-use crate::opener::ParquetOpener;
use crate::opener::build_pruning_predicates;
+use crate::opener::{ParquetMorselizer, ParquetOpener};
use crate::row_filter::can_expr_be_pushed_down_with_schemas;
use datafusion_common::config::ConfigOptions;
#[cfg(feature = "parquet_encryption")]
@@ -543,32 +543,34 @@ impl FileSource for ParquetSource {
.map(|time_unit| parse_coerce_int96_string(time_unit.as_str()).unwrap());
let opener = Arc::new(ParquetOpener {
- partition_index: partition,
- projection: self.projection.clone(),
- batch_size: self
- .batch_size
- .expect("Batch size must set before creating ParquetOpener"),
- limit: base_config.limit,
- preserve_order: base_config.preserve_order,
- predicate: self.predicate.clone(),
- table_schema: self.table_schema.clone(),
- metadata_size_hint: self.metadata_size_hint,
- metrics: self.metrics().clone(),
- parquet_file_reader_factory,
- pushdown_filters: self.pushdown_filters(),
- reorder_filters: self.reorder_filters(),
- force_filter_selections: self.force_filter_selections(),
- enable_page_index: self.enable_page_index(),
- enable_bloom_filter: self.bloom_filter_on_read(),
- enable_row_group_stats_pruning: self.table_parquet_options.global.pruning,
- coerce_int96,
- #[cfg(feature = "parquet_encryption")]
- file_decryption_properties,
- expr_adapter_factory,
- #[cfg(feature = "parquet_encryption")]
- encryption_factory: self.get_encryption_factory_with_config(),
- max_predicate_cache_size: self.max_predicate_cache_size(),
- reverse_row_groups: self.reverse_row_groups,
+ morselizer: ParquetMorselizer {
+ partition_index: partition,
+ projection: self.projection.clone(),
+ batch_size: self
+ .batch_size
+ .expect("Batch size must set before creating ParquetOpener"),
+ limit: base_config.limit,
+ preserve_order: base_config.preserve_order,
+ predicate: self.predicate.clone(),
+ table_schema: self.table_schema.clone(),
+ metadata_size_hint: self.metadata_size_hint,
+ metrics: self.metrics().clone(),
+ parquet_file_reader_factory,
+ pushdown_filters: self.pushdown_filters(),
+ reorder_filters: self.reorder_filters(),
+ force_filter_selections: self.force_filter_selections(),
+ enable_page_index: self.enable_page_index(),
+ enable_bloom_filter: self.bloom_filter_on_read(),
+ enable_row_group_stats_pruning: self.table_parquet_options.global.pruning,
+ coerce_int96,
+ #[cfg(feature = "parquet_encryption")]
+ file_decryption_properties,
+ expr_adapter_factory,
+ #[cfg(feature = "parquet_encryption")]
+ encryption_factory: self.get_encryption_factory_with_config(),
+ max_predicate_cache_size: self.max_predicate_cache_size(),
+ reverse_row_groups: self.reverse_row_groups,
+ },
});
Ok(opener)
}
diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs
index bcc4627050..a9600271c2 100644
--- a/datafusion/datasource/src/mod.rs
+++ b/datafusion/datasource/src/mod.rs
@@ -38,6 +38,7 @@ pub mod file_scan_config;
pub mod file_sink_config;
pub mod file_stream;
pub mod memory;
+pub mod morsel;
pub mod projection;
pub mod schema_adapter;
pub mod sink;
diff --git a/datafusion/datasource/src/morsel/mod.rs b/datafusion/datasource/src/morsel/mod.rs
new file mode 100644
index 0000000000..5f200d7794
--- /dev/null
+++ b/datafusion/datasource/src/morsel/mod.rs
@@ -0,0 +1,229 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Structures for Morsel Driven IO.
+//!
+//! NOTE: As of DataFusion 54.0.0, these are experimental APIs that may change
+//! substantially.
+//!
+//! Morsel Driven IO is a technique for parallelizing the reading of large files
+//! by dividing them into smaller "morsels" that are processed independently.
+//!
+//! It is inspired by the paper [Morsel-Driven Parallelism: A NUMA-Aware Query
+//! Evaluation Framework for the Many-Core Age](https://db.in.tum.de/~leis/papers/morsels.pdf).
+
+use crate::PartitionedFile;
+use arrow::array::RecordBatch;
+use datafusion_common::Result;
+use futures::FutureExt;
+use futures::future::BoxFuture;
+use futures::stream::BoxStream;
+use std::fmt::Debug;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+/// A Morsel of work ready to resolve to a stream of [`RecordBatch`]es.
+///
+/// This represents a single morsel of work that is ready to be processed. It
+/// has all data necessary (does not need any I/O) and is ready to be turned
+/// into a stream of [`RecordBatch`]es for processing by the execution engine.
+pub trait Morsel: Send + Debug {
+ /// Consume this morsel and produce a stream of [`RecordBatch`]es for processing.
+ ///
+ /// Note: This may do CPU work to decode already-loaded data, but should not
+ /// do any I/O work such as reading from the file.
+ fn into_stream(self: Box<Self>) -> BoxStream<'static, Result<RecordBatch>>;
+}
+
+/// A Morselizer takes a single [`PartitionedFile`] and creates the initial planner
+/// for that file.
+///
+/// This is the entry point for morsel driven I/O.
+pub trait Morselizer: Send + Sync + Debug {
+ /// Return the initial [`MorselPlanner`] for this file.
+ ///
+ /// Morselizing a file may involve CPU work, such as parsing parquet
+ /// metadata and evaluating pruning predicates. It should NOT do any I/O
+ /// work, such as reading from the file. Any needed I/O should be done using
+ /// [`MorselPlan::with_pending_planner`].
+ fn plan_file(&self, file: PartitionedFile) -> Result<Box<dyn MorselPlanner>>;
+}
+
+/// A Morsel Planner is responsible for creating morsels for a given scan.
+///
+/// The [`MorselPlanner`] is the unit of I/O. There is only ever a single I/O
+/// outstanding for a specific planner. DataFusion may run
+/// multiple planners in parallel, which corresponds to multiple parallel
+/// I/O requests.
+///
+/// It is not a Rust `Stream` so that it can explicitly separate CPU bound
+/// work from I/O work.
+///
+/// The design is similar to `ParquetPushDecoder`: when `plan` is called, it
+/// should do CPU work to produce the next morsels or discover the next I/O
+/// phase.
+///
+/// Best practice is to spawn I/O in a Tokio task on a separate runtime to
+/// ensure that CPU work doesn't block or slow down I/O work, but this is not
+/// strictly required by the API.
+pub trait MorselPlanner: Send + Debug {
+ /// Attempt to plan morsels. This may involve CPU work, such as parsing
+ /// parquet metadata and evaluating pruning predicates.
+ ///
+ /// It should NOT do any I/O work, such as reading from the file. If I/O is
+ /// required, the returned [`MorselPlan`] should contain a pending planner
+ /// future that the caller polls to drive the I/O work to completion. Once
+ /// that future resolves, it yields a planner ready for work.
+ ///
+ /// Note this function is **not async** to make it explicitly clear that if
+ /// I/O is required, it should be done in the returned `io_future`.
+ ///
+ /// Returns `None` if the planner has no more work to do.
+ ///
+ /// # Empty Morsel Plans
+ ///
+ /// It may return `None`, which means no batches will be read from the file
+ /// (e.g. due to late-pruning based on statistics).
+ ///
+ /// # Output Ordering
+ ///
+ /// See the comments on [`MorselPlan`] for the logical output order.
+ fn plan(self: Box<Self>) -> Result<Option<MorselPlan>>;
+}
+
+/// Return result of [`MorselPlanner::plan`].
+///
+/// # Logical Ordering
+///
+/// For plans where the output order of rows is maintained, the output order of
+/// a [`MorselPlanner`] is logically defined as follows:
+/// 1. All morsels that are directly produced
+/// 2. Recursively, all morsels produced by the returned `planners`
+#[derive(Default)]
+pub struct MorselPlan {
+ /// Morsels ready for CPU work
+ morsels: Vec<Box<dyn Morsel>>,
+ /// Planners that are ready for CPU work.
+ ready_planners: Vec<Box<dyn MorselPlanner>>,
+ /// A future with planner I/O that resolves to a CPU ready planner.
+ ///
+ /// DataFusion will poll this future occasionally to drive the I/O work to
+ /// completion. Once it resolves, planning continues with the returned
+ /// planner.
+ pending_planner: Option<PendingMorselPlanner>,
+}
+
+impl MorselPlan {
+ /// Create an empty morsel plan.
+ pub fn new() -> Self {
+ Self::default()
+ }
+
+ /// Set the ready morsels.
+ pub fn with_morsels(mut self, morsels: Vec<Box<dyn Morsel>>) -> Self {
+ self.morsels = morsels;
+ self
+ }
+
+ /// Set the ready child planners.
+ pub fn with_planners(mut self, planners: Vec<Box<dyn MorselPlanner>>) -> Self {
+ self.ready_planners = planners;
+ self
+ }
+
+ /// Set the pending planner for an I/O phase.
+ pub fn with_pending_planner<F>(mut self, io_future: F) -> Self
+ where
+ F: Future<Output = Result<Box<dyn MorselPlanner>>> + Send + 'static,
+ {
+ self.pending_planner = Some(PendingMorselPlanner::new(io_future));
+ self
+ }
+
+ /// Set the pending planner for an I/O phase.
+ pub fn set_pending_planner<F>(&mut self, io_future: F)
+ where
+ F: Future<Output = Result<Box<dyn MorselPlanner>>> + Send + 'static,
+ {
+ self.pending_planner = Some(PendingMorselPlanner::new(io_future));
+ }
+
+ /// Take the ready morsels.
+ pub fn take_morsels(&mut self) -> Vec<Box<dyn Morsel>> {
+ std::mem::take(&mut self.morsels)
+ }
+
+ /// Take the ready child planners.
+ pub fn take_ready_planners(&mut self) -> Vec<Box<dyn MorselPlanner>> {
+ std::mem::take(&mut self.ready_planners)
+ }
+
+ /// Take the pending I/O future, if any.
+ pub fn take_pending_planner(&mut self) -> Option<PendingMorselPlanner> {
+ self.pending_planner.take()
+ }
+
+ /// Returns `true` if this plan contains an I/O future.
+ pub fn has_io_future(&self) -> bool {
+ self.pending_planner.is_some()
+ }
+}
+
+/// Wrapper for I/O that must complete before planning can continue.
+pub struct PendingMorselPlanner {
+ future: BoxFuture<'static, Result<Box<dyn MorselPlanner>>>,
+}
+
+impl PendingMorselPlanner {
+ /// Create a new pending planner future.
+ ///
+ /// Example
+ /// ```
+ /// # use datafusion_common::DataFusionError;
+ /// # use datafusion_datasource::morsel::{MorselPlanner, PendingMorselPlanner};
+ /// let work = async move {
+ /// let planner: Box<dyn MorselPlanner> = {
+ /// // Do I/O work here, then return the next planner to run.
+ /// # unimplemented!();
+ /// };
+ /// Ok(planner) as Result<_, DataFusionError>;
+ /// };
+ /// let pending_io = PendingMorselPlanner::new(work);
+ /// ```
+ pub fn new<F>(future: F) -> Self
+ where
+ F: Future<Output = Result<Box<dyn MorselPlanner>>> + Send + 'static,
+ {
+ Self {
+ future: future.boxed(),
+ }
+ }
+
+ /// Consume this wrapper and return the underlying future.
+ pub fn into_future(self) -> BoxFuture<'static, Result<Box<dyn MorselPlanner>>> {
+ self.future
+ }
+}
+
+/// Forwards polling to the underlying future.
+impl Future for PendingMorselPlanner {
+ type Output = Result<Box<dyn MorselPlanner>>;
+ fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+ // forward request to inner
+ self.future.as_mut().poll(cx)
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org