perf: optimize object store requests when reading CSV (apache#22962)

## Which issue does this PR close?

- Closes apache#21419

## Rationale for this change

The CSV scanner currently uses `calculate_range` which issues two extra
`get_opts` requests per byte range to find newline boundaries (one for
the start boundary, one for the end boundary), plus one GET for the
actual data. For a file split into 3 partitions, this results in 8 total
object store requests.

apache#20823 solved this same problem for the JSON scanner by introducing
`AlignedBoundaryStream`, which wraps the raw byte stream and lazily
aligns to newline boundaries as data is read, eliminating the extra
boundary-seeking requests entirely. This PR applies the same approach to
CSV.

## What changes are included in this PR?

Based on the approach from apache#20823:

Moved `AlignedBoundaryStream` from `datasource-json` to the shared
`datasource` crate so it can be reused by both JSON and CSV scanners.
Updated `CsvOpener` to use instead of `calculate_range`, and removed the
`calculate_range` & `find_first_newline` as they no longer had any
callers. Updated tests to reflect.

Public API changes include:
- Removal of the `RangeCalculation` enum
- Removal of the `calculate_change` function
- `boundary_stream` (containing `AlignedBoundaryStream`) moved from
`datafusion_datasource_json` to `datafusion_datasource`

## Are these changes tested?

Yes. The existing `AlignedBoundaryStream` unit tests (16 tests covering
boundary alignment edge cases) were moved along with the implementation
and continue to pass. The `query_csv_file_with_byte_range_partitions`
snapshot test in `object_store_access.rs` has been updated to verify the
new request pattern (4 requests instead of 8).


## Are there any user-facing changes?

No.

Author: saadtajwar

datafusion/core/tests/datasource/object_store_access.rs

@@ -231,23 +231,13 @@ async fn query_multi_csv_file() {
     );
 }
 
-/// Test that a CSV file split into byte ranges via repartitioning exercises
-/// range-based object store access.
+/// Test that a CSV file split into byte ranges via repartitioning produces
+/// exactly one GET request per byte range — no extra requests for boundary seeking.
 ///
 /// With a single file and `target_partitions=3`, the repartitioner produces
-/// exactly 3 ranges.  For each range, `calculate_range` calls
-/// `find_first_newline` via a GET for every non-file boundary it touches
-/// (the start boundary if `start > 0`, the end boundary if `end < file_size`),
-/// plus one GET for the actual data — so 2 GETs for the first range (end scan
-/// + data), 3 for the middle range (start scan + end scan + data), and 2 for
-/// the last range (start scan + data) = 7 data GETs total.  Additionally,
-/// adjacent ranges share a boundary position, so each shared boundary is scanned
-/// twice — once as the left range's end and again as the right range's start —
-/// producing the duplicate GETs visible in the snapshot.  Add the 1 HEAD for
-/// file-size metadata = **8 total**.
-///
-/// This differs from the JSON reader which uses [`AlignedBoundaryStream`] and
-/// needs only 1 GET per range.
+/// exactly 3 ranges.  Each range is served by a single [`AlignedBoundaryStream`]
+/// which issues exactly one bounded `get_opts` call, so there are 3 data GETs
+/// plus 1 HEAD (to determine file size) = **4 total**.
 ///
 /// This test documents the current request pattern to catch regressions.
 #[tokio::test]
@@ -275,15 +265,11 @@ async fn query_csv_file_with_byte_range_partitions() {
     +---------+-------+-------+
     ------- Object Store Request Summary -------
     RequestCountingObjectStore()
-    Total Requests: 8
+    Total Requests: 4
     - GET  (opts) path=csv_range_table.csv head=true
+    - GET  (opts) path=csv_range_table.csv range=0-129
     - GET  (opts) path=csv_range_table.csv range=42-129
-    - GET  (opts) path=csv_range_table.csv range=0-49
-    - GET  (opts) path=csv_range_table.csv range=42-129
-    - GET  (opts) path=csv_range_table.csv range=85-129
-    - GET  (opts) path=csv_range_table.csv range=49-89
     - GET  (opts) path=csv_range_table.csv range=85-129
-    - GET  (opts) path=csv_range_table.csv range=89-129
     "
     );
 }

datafusion/datasource-csv/src/source.rs

@@ -17,24 +17,23 @@
 
 //! Execution plan for reading CSV files
 
+use datafusion_datasource::boundary_stream::AlignedBoundaryStream;
 use datafusion_datasource::projection::{ProjectionOpener, SplitProjection};
 use datafusion_physical_plan::projection::ProjectionExprs;
 use std::fmt;
-use std::io::{Read, Seek, SeekFrom};
+use std::io::Read;
 use std::sync::Arc;
-use std::task::Poll;
 
 use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream};
 use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};
 use datafusion_datasource::{
-    FileRange, ListingTableUrl, PartitionedFile, RangeCalculation, TableSchema,
-    as_file_source, calculate_range,
+    FileRange, ListingTableUrl, PartitionedFile, TableSchema, as_file_source,
 };
 
 use arrow::csv;
 use datafusion_common::config::CsvOptions;
-use datafusion_common::{DataFusionError, Result};
+use datafusion_common::{DataFusionError, Result, exec_datafusion_err};
 use datafusion_common_runtime::JoinSet;
 use datafusion_datasource::file::FileSource;
 use datafusion_datasource::file_scan_config::FileScanConfig;
@@ -367,43 +366,53 @@ impl FileOpener for CsvOpener {
 
         Ok(Box::pin(async move {
             // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries)
+            let file_size = partitioned_file.object_meta.size;
+            let location = partitioned_file.object_meta.location;
+
+            if let Some(file_range) = partitioned_file.range.as_ref() {
+                let raw_start: u64 = file_range.start.try_into().map_err(|_| {
+                    exec_datafusion_err!(
+                        "Expected start range to fit in u64, got {}",
+                        file_range.start
+                    )
+                })?;
+                let raw_end: u64 = file_range.end.try_into().map_err(|_| {
+                    exec_datafusion_err!(
+                        "Expected end range to fit in u64, got {}",
+                        file_range.end
+                    )
+                })?;
+
+                let aligned_stream = AlignedBoundaryStream::new(
+                    Arc::clone(&store),
+                    location.clone(),
+                    raw_start,
+                    raw_end,
+                    file_size,
+                    terminator.unwrap_or(b'\n'),
+                )
+                .await?
+                .map_err(DataFusionError::from);
+
+                let decoder = config.builder().build_decoder();
+                let input = file_compression_type
+                    .convert_stream(aligned_stream.boxed())?
+                    .fuse();
+                let stream = deserialize_stream(
+                    input,
+                    DecoderDeserializer::new(CsvDecoder::new(decoder)),
+                );
+                return Ok(stream.map_err(Into::into).boxed());
+            }
 
-            let calculated_range =
-                calculate_range(&partitioned_file, &store, terminator).await?;
-
-            let range = match calculated_range {
-                RangeCalculation::Range(None) => None,
-                RangeCalculation::Range(Some(range)) => Some(range.into()),
-                RangeCalculation::TerminateEarly => {
-                    return Ok(
-                        futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed()
-                    );
-                }
-            };
-
-            let options = GetOptions {
-                range,
-                ..Default::default()
-            };
-
-            let result = store
-                .get_opts(&partitioned_file.object_meta.location, options)
-                .await?;
+            // No range specified — read the entire file
+            let options = GetOptions::default();
+            let result = store.get_opts(&location, options).await?;
 
             match result.payload {
                 #[cfg(not(target_arch = "wasm32"))]
-                GetResultPayload::File(mut file, _) => {
-                    let is_whole_file_scanned = partitioned_file.range.is_none();
-                    let decoder = if is_whole_file_scanned {
-                        // Don't seek if no range as breaks FIFO files
-                        file_compression_type.convert_read(file)?
-                    } else {
-                        file.seek(SeekFrom::Start(result.range.start as _))?;
-                        file_compression_type.convert_read(
-                            file.take((result.range.end - result.range.start) as u64),
-                        )?
-                    };
-
+                GetResultPayload::File(file, _) => {
+                    let decoder = file_compression_type.convert_read(file)?;
                     let mut reader = config.open(decoder)?;
 
                     // Use std::iter::from_fn to wrap execution of iterator's next() method.

datafusion/datasource-json/src/mod.rs

@@ -20,7 +20,6 @@
 // https://github.com/apache/datafusion/issues/11143
 #![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))]
 
-pub mod boundary_stream;
 pub mod file_format;
 pub mod source;
 pub mod utils;

datafusion/datasource-json/src/source.rs

@@ -25,11 +25,10 @@ use std::task::{Context, Poll};
 use crate::file_format::JsonDecoder;
 use crate::utils::{ChannelReader, JsonArrayToNdjsonReader};
 
-use crate::boundary_stream::AlignedBoundaryStream;
-
 use datafusion_common::error::{DataFusionError, Result};
 use datafusion_common::exec_datafusion_err;
 use datafusion_common_runtime::{JoinSet, SpawnedTask};
+use datafusion_datasource::boundary_stream::AlignedBoundaryStream;
 use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream};
 use datafusion_datasource::file_compression_type::FileCompressionType;
 use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener};

…n/datasource-json/src/boundary_stream.rs …fusion/datasource/src/boundary_stream.rsdatafusion/datasource-json/src/boundary_stream.rs renamed to datafusion/datasource/src/boundary_stream.rs datafusion/datasource-json/src/boundary_stream.rs renamed to datafusion/datasource/src/boundary_stream.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-//! Streaming boundary-aligned wrapper for newline-delimited JSON range reads.
+//! Streaming boundary-aligned wrapper for newline-delimited JSON and CSV range reads.
 //!
 //! [`AlignedBoundaryStream`] wraps a raw byte stream and lazily aligns to
 //! record (newline) boundaries, avoiding the need for separate `get_opts`
@@ -398,7 +398,7 @@ impl Stream for AlignedBoundaryStream {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::test_utils::{CHUNK_SIZES, make_chunked_store};
+    use crate::test_util::{CHUNK_SIZES, make_chunked_store};
     use futures::TryStreamExt;
 
     async fn collect_stream(stream: AlignedBoundaryStream) -> Vec<u8> {

datafusion/datasource/src/mod.rs

@@ -28,6 +28,7 @@
 //! A table that uses the `ObjectStore` listing capability
 //! to get the list of files to process.
 
+pub mod boundary_stream;
 pub mod decoder;
 pub mod display;
 pub mod file;
@@ -56,11 +57,10 @@ pub use self::url::ListingTableUrl;
 use crate::file_groups::FileGroup;
 use chrono::TimeZone;
 use datafusion_common::stats::Precision;
-use datafusion_common::{ColumnStatistics, Result, TableReference, exec_datafusion_err};
+use datafusion_common::{ColumnStatistics, Result, TableReference};
 use datafusion_common::{ScalarValue, Statistics};
 use datafusion_physical_expr::LexOrdering;
-use futures::{Stream, StreamExt};
-use object_store::{GetOptions, GetRange, ObjectStore};
+use futures::Stream;
 use object_store::{ObjectMeta, path::Path};
 pub use table_schema::{TableSchema, TableSchemaBuilder};
 // Remove when add_row_stats is remove
@@ -69,7 +69,6 @@ use arrow::datatypes::SchemaRef;
 pub use statistics::add_row_stats;
 pub use statistics::compute_all_files_statistics;
 use std::any::Any;
-use std::ops::Range;
 use std::pin::Pin;
 use std::sync::Arc;
 
@@ -411,119 +410,6 @@ impl From<ObjectMeta> for PartitionedFile {
     }
 }
 
-/// Represents the possible outcomes of a range calculation.
-///
-/// This enum is used to encapsulate the result of calculating the range of
-/// bytes to read from an object (like a file) in an object store.
-///
-/// Variants:
-/// - `Range(Option<Range<usize>>)`:
-///   Represents a range of bytes to be read. It contains an `Option` wrapping a
-///   `Range<usize>`. `None` signifies that the entire object should be read,
-///   while `Some(range)` specifies the exact byte range to read.
-/// - `TerminateEarly`:
-///   Indicates that the range calculation determined no further action is
-///   necessary, possibly because the calculated range is empty or invalid.
-pub enum RangeCalculation {
-    Range(Option<Range<u64>>),
-    TerminateEarly,
-}
-
-/// Calculates an appropriate byte range for reading from an object based on the
-/// provided metadata.
-///
-/// This asynchronous function examines the [`PartitionedFile`] of an object in an object store
-/// and determines the range of bytes to be read. The range calculation may adjust
-/// the start and end points to align with meaningful data boundaries (like newlines).
-///
-/// Returns a `Result` wrapping a [`RangeCalculation`], which is either a calculated byte range or an indication to terminate early.
-///
-/// Returns an `Error` if any part of the range calculation fails, such as issues in reading from the object store or invalid range boundaries.
-pub async fn calculate_range(
-    file: &PartitionedFile,
-    store: &Arc<dyn ObjectStore>,
-    terminator: Option<u8>,
-) -> Result<RangeCalculation> {
-    let location = &file.object_meta.location;
-    let file_size = file.object_meta.size;
-    let newline = terminator.unwrap_or(b'\n');
-
-    match file.range {
-        None => Ok(RangeCalculation::Range(None)),
-        Some(FileRange { start, end }) => {
-            let start: u64 = start.try_into().map_err(|_| {
-                exec_datafusion_err!("Expect start range to fit in u64, got {start}")
-            })?;
-            let end: u64 = end.try_into().map_err(|_| {
-                exec_datafusion_err!("Expect end range to fit in u64, got {end}")
-            })?;
-
-            let start_delta = if start != 0 {
-                find_first_newline(store, location, start - 1, file_size, newline).await?
-            } else {
-                0
-            };
-
-            if start + start_delta > end {
-                return Ok(RangeCalculation::TerminateEarly);
-            }
-
-            let end_delta = if end != file_size {
-                find_first_newline(store, location, end - 1, file_size, newline).await?
-            } else {
-                0
-            };
-
-            let range = start + start_delta..end + end_delta;
-
-            if range.start >= range.end {
-                return Ok(RangeCalculation::TerminateEarly);
-            }
-
-            Ok(RangeCalculation::Range(Some(range)))
-        }
-    }
-}
-
-/// Asynchronously finds the position of the first newline character in a specified byte range
-/// within an object, such as a file, in an object store.
-///
-/// This function scans the contents of the object starting from the specified `start` position
-/// up to the `end` position, looking for the first occurrence of a newline character.
-/// It returns the position of the first newline relative to the start of the range.
-///
-/// Returns a `Result` wrapping a `usize` that represents the position of the first newline character found within the specified range. If no newline is found, it returns the length of the scanned data, effectively indicating the end of the range.
-///
-/// The function returns an `Error` if any issues arise while reading from the object store or processing the data stream.
-async fn find_first_newline(
-    object_store: &Arc<dyn ObjectStore>,
-    location: &Path,
-    start: u64,
-    end: u64,
-    newline: u8,
-) -> Result<u64> {
-    let options = GetOptions {
-        range: Some(GetRange::Bounded(start..end)),
-        ..Default::default()
-    };
-
-    let result = object_store.get_opts(location, options).await?;
-    let mut result_stream = result.into_stream();
-
-    let mut index = 0;
-
-    while let Some(chunk) = result_stream.next().await.transpose()? {
-        if let Some(position) = chunk.iter().position(|&byte| byte == newline) {
-            let position = position as u64;
-            return Ok(index + position);
-        }
-
-        index += chunk.len() as u64;
-    }
-
-    Ok(index)
-}
-
 /// Generates test files with min-max statistics in different overlap patterns.
 ///
 /// Used by tests and benchmarks.
@@ -647,7 +533,7 @@ mod tests {
     use datafusion_execution::object_store::{
         DefaultObjectStoreRegistry, ObjectStoreRegistry,
     };
-    use object_store::{ObjectStoreExt, local::LocalFileSystem, path::Path};
+    use object_store::{local::LocalFileSystem, path::Path};
     use std::{collections::HashMap, ops::Not, sync::Arc};
     use url::Url;
 
@@ -844,31 +730,4 @@ mod tests {
         // testing an empty path with `ignore_subdirectory` set to false
         assert!(url.contains(&Path::parse("/var/data/mytable/").unwrap(), false));
     }
-
-    /// Regression test for <https://github.com/apache/datafusion/issues/19605>
-    #[tokio::test]
-    async fn test_calculate_range_single_line_file() {
-        use super::{PartitionedFile, RangeCalculation, calculate_range};
-        use object_store::ObjectStore;
-        use object_store::memory::InMemory;
-
-        let content = r#"{"id":1,"data":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}"#;
-        let file_size = content.len() as u64;
-
-        let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
-        let path = Path::from("test.json");
-        store.put(&path, content.into()).await.unwrap();
-
-        let mid = file_size / 2;
-        let partitioned_file = PartitionedFile::new_with_range(
-            path.to_string(),
-            file_size,
-            mid as i64,
-            file_size as i64,
-        );
-
-        let result = calculate_range(&partitioned_file, &store, None).await;
-
-        assert!(matches!(result, Ok(RangeCalculation::TerminateEarly)));
-    }
 }