Sketch out a Morselize API

Author: alamb

datafusion/datasource/src/file_stream.rs

@@ -42,6 +42,7 @@ use datafusion_common::instant::Instant;
 use futures::future::BoxFuture;
 use futures::stream::BoxStream;
 use futures::{FutureExt as _, Stream, StreamExt as _, ready};
+use crate::morsel::Morselizer;
 
 /// A stream that iterates record batch by record batch, file over file.
 pub struct FileStream {
@@ -50,8 +51,10 @@ pub struct FileStream {
     /// The stream schema (file schema including partition columns and after
     /// projection).
     projected_schema: SchemaRef,
-    /// The remaining number of records to parse, None if no limit
+    /// The remaining number of records to parse until limit is reached, None if no limit
     remain: Option<usize>,
+    /// A type specific [`Morselizer`] that examines the input files and produces a stream of `Morsels`
+    morselizer: Option<Box<dyn Morselizer>>,
     /// A dynamic [`FileOpener`]. Calling `open()` returns a [`FileOpenFuture`],
     /// which can be resolved to a stream of `RecordBatch`.
     file_opener: Arc<dyn FileOpener>,
@@ -81,6 +84,7 @@ impl FileStream {
             file_iter: file_group.into_inner().into_iter().collect(),
             projected_schema,
             remain: config.limit,
+            morselizer: None,
             file_opener,
             state: FileStreamState::Idle,
             file_stream_metrics: FileStreamMetrics::new(metrics, partition),
@@ -292,6 +296,9 @@ impl RecordBatchStream for FileStream {
 }
 
 /// A fallible future that resolves to a stream of [`RecordBatch`]
+/// 
+/// This is typically an `async` function that opens the file, and returns a
+/// stream that reads the file and produces `RecordBatch`es. 
 pub type FileOpenFuture =
     BoxFuture<'static, Result<BoxStream<'static, Result<RecordBatch>>>>;
 
@@ -308,10 +315,15 @@ pub enum OnError {
 /// Generic API for opening a file using an [`ObjectStore`] and resolving to a
 /// stream of [`RecordBatch`]
 ///
+/// TODO: rename this to MorselOpener
+/// Also, add documentation here describing IO expectations
+///
 /// [`ObjectStore`]: object_store::ObjectStore
 pub trait FileOpener: Unpin + Send + Sync {
     /// Asynchronously open the specified file and return a stream
     /// of [`RecordBatch`]
+    ///
+    /// TODO: describe prefetching behavior here, and expectations around IO 
     fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture>;
 }
 
@@ -330,7 +342,7 @@ pub enum FileStreamState {
     /// for a given file
     Open {
         /// A [`FileOpenFuture`] returned by [`FileOpener::open`]
-        future: FileOpenFuture,
+        future: FileOpenFuture, 
     },
     /// Scanning the [`BoxStream`] returned by the completion of a [`FileOpenFuture`]
     /// returned by [`FileOpener::open`]

datafusion/datasource/src/mod.rs

@@ -44,6 +44,7 @@ pub mod sink;
 pub mod source;
 mod statistics;
 pub mod table_schema;
+mod morsel;
 
 #[cfg(test)]
 pub mod test_util;
@@ -70,7 +71,10 @@ use std::ops::Range;
 use std::pin::Pin;
 use std::sync::Arc;
 
+pub use morsel::Morsel;
+
 /// Stream of files get listed from object store
+#[deprecated(since="54.0.0", note = "This type is unused and will be removed in a future release")]
 pub type PartitionedFileStream =
     Pin<Box<dyn Stream<Item = Result<PartitionedFile>> + Send + Sync + 'static>>;
 
@@ -122,7 +126,8 @@ pub struct PartitionedFile {
     /// [`wrap_partition_value_in_dict`]: crate::file_scan_config::wrap_partition_value_in_dict
     /// [`table_partition_cols`]: https://github.com/apache/datafusion/blob/main/datafusion/core/src/datasource/file_format/options.rs#L87
     pub partition_values: Vec<ScalarValue>,
-    /// An optional file range for a more fine-grained parallel execution
+    /// An optional file range for this file. This is used to statically
+    /// schedule non-overlapping sections of a file to be read in parallel.
     pub range: Option<FileRange>,
     /// Optional statistics that describe the data in this file if known.
     ///

datafusion/datasource/src/morsel.rs

@@ -0,0 +1,57 @@
+// 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
+//!
+//! Morsel Driven IO is a technique for parallelizing the reading of large files
+//! by dividing them into smaller "morsels" that can be 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 std::any::Any;
+use std::sync::Arc;
+use futures::future::BoxFuture;
+use datafusion_common::error::Result;
+use futures::stream::BoxStream;
+use crate::PartitionedFile;
+
+/// A Morsel represents a portion of a file that can be processed independently.
+#[derive(Debug)]
+pub struct Morsel {
+    /// The original [`PartitionedFile`] that this morsel belongs to
+    file: Arc<PartitionedFile>,
+    /// File format specific information that describes the morsel, such as byte range, row group, etc.
+    info: Box<dyn Any>,
+}
+
+/// A fallible future that resolves to a stream of [`Morsel`]
+///
+/// This is typically an `async` function that opens the file, and returns a
+/// stream of Morsels that can be processed independently. The stream may yield
+/// an error if the file cannot be opened or read.
+pub type MorselOpenFuture = BoxFuture<'static, Result<BoxStream<'static, Morsel>>>;
+
+pub trait  Morselizer: Send + Sync {
+    /// Given a [`PartitionedFile`], return a Vec of [`Morsel`]s, potentially doing IO
+    ///
+    /// Notes
+    ///
+    /// 1. The API is `async` for file formats that require I/O operations to
+    ///    determine the morsels, such as reading metadata from the file.
+    /// 2.
+   fn morselize(&self, file: Arc<PartitionedFile>) -> MorselOpenFuture;
+}