Skip to content

Optimize split calculation in DF source #9167

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AdamGS wants to merge 2 commits into from
Open

Optimize split calculation in DF source #9167

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2404723
Precompute natural split assignment bytes
AdamGS Aug 4, 2026
762c28f
better cache plan
AdamGS Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
247 changes: 161 additions & 86 deletions vortex-datafusion/src/persistent/opener.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,18 +37,15 @@ use futures::FutureExt;
use futures::StreamExt;
use futures::TryStreamExt;
use futures::stream;
use itertools::Itertools;
use object_store::path::Path;
use tracing::Instrument;
use vortex::array::VortexSessionExecute;
use vortex::dtype::FieldMask;
use vortex::error::VortexError;
use vortex::error::VortexExpect;
use vortex::file::OpenOptionsSessionExt;
use vortex::io::InstrumentedReadAt;
use vortex::layout::LayoutReader;
use vortex::layout::scan::scan_builder::ScanBuilder;
use vortex::layout::scan::split_by::SplitBy;
use vortex::metrics::Label;
use vortex::metrics::MetricsRegistry;
use vortex::session::VortexSession;
Expand Down Expand Up @@ -98,8 +95,8 @@ pub(crate) struct VortexOpener {
/// To save on the overhead of reparsing FlatBuffers and rebuilding the layout tree, we cache
/// a file reader the first time we read a file.
pub layout_readers: Arc<DashMap<Path, Weak<dyn LayoutReader>>>,
/// Shared full-file natural split ranges keyed by file path.
pub natural_split_ranges: Arc<DashMap<Path, Arc<[Range<u64>]>>>,
/// Shared full-file natural splits keyed by file path.
pub natural_splits: Arc<DashMap<Path, Arc<NaturalSplits>>>,
/// Whether the query has output ordering specified
pub has_output_ordering: bool,

Expand Down Expand Up @@ -140,7 +137,7 @@ impl FileOpener for VortexOpener {
let unified_file_schema = Arc::clone(self.table_schema.file_schema());
let limit = self.limit;
let layout_readers = Arc::clone(&self.layout_readers);
let natural_split_ranges = Arc::clone(&self.natural_split_ranges);
let natural_splits = Arc::clone(&self.natural_splits);
let has_output_ordering = self.has_output_ordering;
let scan_concurrency = self.scan_concurrency;

Expand Down Expand Up @@ -360,34 +357,6 @@ impl FileOpener for VortexOpener {
scan_builder = vortex_plan.apply_to_builder(scan_builder);
}

if let Some(file_range) = file.range {
let byte_range = Range {
start: u64::try_from(file_range.start)
.map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?,
end: u64::try_from(file_range.end)
.map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?,
};
if byte_range.start != 0 || byte_range.end != file.object_meta.size {
// Full-file scans already cover every natural split. Only translate the
// byte range back into row boundaries when DataFusion has trimmed the file.
let natural_split_ranges = natural_split_ranges_for_file(
natural_split_ranges.as_ref(),
&file.object_meta.location,
&layout_reader,
)?;

let Some(row_range) = split_aligned_row_range(
byte_range,
file.object_meta.size,
natural_split_ranges.as_ref(),
) else {
return Ok(stream::empty().boxed());
};

scan_builder = scan_builder.with_row_range(row_range);
}
}

let filter = filter
.and_then(|f| {
// Verify that all filters we've accepted from DataFusion get pushed down.
Expand Down Expand Up @@ -430,11 +399,46 @@ impl FileOpener for VortexOpener {
scan_builder = scan_builder.with_concurrency(concurrency);
}

// Set before the byte-range translation below, which computes natural splits for
// the fields the scan's projection and filter reference.
scan_builder = scan_builder
.with_projection(scan_projection)
.with_some_filter(filter);

if let Some(file_range) = file.range {
let byte_range = Range {
start: u64::try_from(file_range.start)
.map_err(|_| exec_datafusion_err!("Vortex file range start is negative"))?,
end: u64::try_from(file_range.end)
.map_err(|_| exec_datafusion_err!("Vortex file range end is negative"))?,
};
if byte_range.start != 0 || byte_range.end != file.object_meta.size {
// Full-file scans already cover every natural split. Only translate the
// byte range back into row boundaries when DataFusion has trimmed the file.
let natural_splits = natural_splits_for_file(
natural_splits.as_ref(),
&file.object_meta.location,
&scan_builder,
file.object_meta.size,
)?;

let Some(row_range) =
split_aligned_row_range(byte_range, natural_splits.as_ref())
else {
return Ok(stream::empty().boxed());
};

scan_builder = scan_builder
.with_row_range(row_range)
// Hand the shared full-file boundaries back to the scan so prepare()
// skips its own layout walk.
.with_natural_splits(Arc::clone(&natural_splits.row_boundaries));
}
}

let stream_target_field = Field::new_struct("", stream_schema.fields().clone(), false);
let stream = scan_builder
.with_metrics_registry(metrics_registry)
.with_projection(scan_projection)
.with_some_filter(filter)
.with_ordered(has_output_ordering)
.map(move |chunk| {
let mut ctx = session.create_execution_ctx();
Expand Down Expand Up @@ -482,72 +486,124 @@ impl FileOpener for VortexOpener {
}
}

fn natural_split_ranges_for_file(
natural_split_ranges: &DashMap<Path, Arc<[Range<u64>]>>,
path: &Path,
layout_reader: &Arc<dyn LayoutReader>,
) -> DFResult<Arc<[Range<u64>]>> {
if let Some(split_ranges) = natural_split_ranges.get(path) {
return Ok(Arc::clone(split_ranges.value()));
/// A file's natural split boundaries plus the precomputed byte each split is assigned to,
/// enabling [`split_aligned_row_range`] to translate a DataFusion byte range into row
/// boundaries with a binary search instead of re-projecting every split per partition.
///
/// The boundaries are computed for the fields referenced by the scan's projection and filter.
/// All partitions translate through the first opener's cached entry (the cache lives on the
/// source, so projection and filter are fixed for its lifetime), which keeps the byte ranges
/// tiling the file's rows exactly once.
#[derive(Debug)]
pub(crate) struct NaturalSplits {
/// Sorted row boundaries of the natural splits; split `i` covers
/// `row_boundaries[i]..row_boundaries[i + 1]`. Shared so partitions can hand the
/// boundaries back to the scan via [`ScanBuilder::with_natural_splits`], skipping the
/// per-partition layout walk in `prepare`.
row_boundaries: Arc<[u64]>,
/// For each split, the byte a DataFusion byte range must contain to own it (see
/// [`split_assignment_byte`]); one entry per split, sorted because split midpoints
/// increase monotonically under the row-to-byte projection.
assignment_bytes: Box<[u64]>,
}

impl NaturalSplits {
fn new(row_boundaries: Arc<[u64]>, total_size: u64) -> Self {
let row_count = row_boundaries.last().copied().unwrap_or_default();
let assignment_bytes = if row_count == 0 {
Box::default()
} else {
row_boundaries
.windows(2)
.enumerate()
.map(|(idx, boundaries)| {
split_assignment_byte(
idx,
&(boundaries[0]..boundaries[1]),
row_count,
total_size,
)
})
.collect()
};

debug_assert!(assignment_bytes.is_sorted());
debug_assert_eq!(
assignment_bytes.len() + usize::from(!row_boundaries.is_empty()),
row_boundaries.len()
);
Comment on lines +531 to +534

@robert3005 robert3005 Aug 5, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is incorrect if row_count == 0 but there's no test that hits this

@AdamGS AdamGS Aug 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

if there are 0 rows, do we report a single split at 0?


Self {
row_boundaries,
assignment_bytes,
}
}
}

let split_ranges = compute_natural_split_ranges(layout_reader.as_ref())?;
/// Return the cached [`NaturalSplits`] for `path`, computing and caching them on first use.
fn natural_splits_for_file<A: 'static + Send>(
natural_splits: &DashMap<Path, Arc<NaturalSplits>>,
path: &Path,
scan_builder: &ScanBuilder<A>,
total_size: u64,
) -> DFResult<Arc<NaturalSplits>> {
if let Some(splits) = natural_splits.get(path) {
return Ok(Arc::clone(splits.value()));
}

match natural_split_ranges.entry(path.clone()) {
// Compute while holding the entry so concurrent partitions opening the same file wait
// for the winner instead of all walking the layout tree; the redundant walks contend on
// the lazily-initialized layout children and dominate the cost of the computation itself.
match natural_splits.entry(path.clone()) {
Entry::Occupied(entry) => Ok(Arc::clone(entry.get())),
Entry::Vacant(entry) => {
entry.insert(Arc::clone(&split_ranges));
Ok(split_ranges)
let splits = compute_natural_splits(scan_builder, total_size)?;
Comment thread
robert3005 marked this conversation as resolved.
entry.insert(Arc::clone(&splits));
Ok(splits)
}
}
}

fn compute_natural_split_ranges(layout_reader: &dyn LayoutReader) -> DFResult<Arc<[Range<u64>]>> {
let row_count = layout_reader.row_count();
let row_range = 0..row_count;
let split_points: Vec<_> = SplitBy::Layout
.splits(layout_reader, &row_range, &[FieldMask::All])
.map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?
.into_iter()
.tuple_windows()
.map(|(s, e)| s..e)
.collect::<Vec<_>>();

Ok(split_points.into())
/// Walk the layout tree to compute the file's full natural split boundaries for the fields
/// referenced by the scan's projection and filter.
fn compute_natural_splits<A: 'static + Send>(
scan_builder: &ScanBuilder<A>,
total_size: u64,
) -> DFResult<Arc<NaturalSplits>> {
let row_boundaries = scan_builder
.full_file_splits()
.map_err(|e| exec_datafusion_err!("Failed to compute Vortex natural splits: {e}"))?;

Ok(Arc::new(NaturalSplits::new(
row_boundaries.into(),
total_size,
)))
}

/// Translate a DataFusion byte range to the contiguous natural split ranges it owns.
/// Most splits are assigned by midpoint, but the leading split stays with the range that owns
/// byte 0 so a tiny first byte range still claims the first rows.
fn split_aligned_row_range(
byte_range: Range<u64>,
total_size: u64,
split_ranges: &[Range<u64>],
natural_splits: &NaturalSplits,
) -> Option<Range<u64>> {
if byte_range.start >= byte_range.end {
return None;
}

let row_count = split_ranges.last().map(|split| split.end)?;
if row_count == 0 {
let first_split = natural_splits
.assignment_bytes
.partition_point(|&assignment_byte| assignment_byte < byte_range.start);
let after_last_split = natural_splits
.assignment_bytes
.partition_point(|&assignment_byte| assignment_byte < byte_range.end);
if first_split == after_last_split {
return None;
}

let mut owned_splits = split_ranges
.iter()
.enumerate()
.filter_map(|(idx, split_range)| {
let assignment_byte = split_assignment_byte(idx, split_range, row_count, total_size);
byte_range.contains(&assignment_byte).then_some(split_range)
});

let first_split = owned_splits.next()?;
let mut row_range = first_split.start..first_split.end;
for split_range in owned_splits {
row_range.end = split_range.end;
}

Some(row_range)
Some(
natural_splits.row_boundaries[first_split]..natural_splits.row_boundaries[after_last_split],
)
}

fn split_assignment_byte(
Expand Down Expand Up @@ -675,6 +731,15 @@ mod tests {
}
}

fn natural_splits(total_size: u64, split_ranges: &[Range<u64>]) -> NaturalSplits {
let mut row_boundaries = Vec::with_capacity(split_ranges.len() + 1);
if let Some(first) = split_ranges.first() {
row_boundaries.push(first.start);
row_boundaries.extend(split_ranges.iter().map(|range| range.end));
}
NaturalSplits::new(row_boundaries.into(), total_size)
}

#[rstest]
#[case(0..3, 10, vec![0..2, 2..5, 5..10], Some(0..2))]
#[case(3..7, 10, vec![0..2, 2..5, 5..10], Some(2..5))]
Expand All @@ -688,7 +753,7 @@ mod tests {
#[case] expected: Option<Range<u64>>,
) {
assert_eq!(
split_aligned_row_range(byte_range, total_size, &split_ranges),
split_aligned_row_range(byte_range, &natural_splits(total_size, &split_ranges)),
expected
);
}
Expand All @@ -697,10 +762,11 @@ mod tests {
fn test_split_aligned_ranges_cover_splits_exactly_once() {
let split_ranges = vec![0..1, 1..4, 4..10, 10..13];
let byte_ranges = [0..4, 4..8, 8..12, 12..16];
let natural_splits = natural_splits(16, &split_ranges);

let assigned = byte_ranges
.into_iter()
.filter_map(|byte_range| split_aligned_row_range(byte_range, 16, &split_ranges))
.filter_map(|byte_range| split_aligned_row_range(byte_range, &natural_splits))
.collect::<Vec<_>>();

assert_eq!(assigned, vec![0..4, 4..10, 10..13]);
Expand Down Expand Up @@ -731,6 +797,15 @@ mod tests {
}
}

#[test]
fn test_split_aligned_row_range_keeps_colliding_assignments_together() {
let natural_splits = natural_splits(2, &[0..1, 1..2, 2..3, 3..4]);

assert_eq!(natural_splits.assignment_bytes.as_ref(), [0, 0, 1, 1]);
assert_eq!(split_aligned_row_range(0..1, &natural_splits), Some(0..2));
assert_eq!(split_aligned_row_range(1..2, &natural_splits), Some(2..4));
}

async fn write_arrow_to_vortex(
object_store: Arc<dyn ObjectStore>,
path: &str,
Expand Down Expand Up @@ -767,7 +842,7 @@ mod tests {
metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
df_metrics: ExecutionPlanMetricsSet::new(),
layout_readers: Default::default(),
natural_split_ranges: Default::default(),
natural_splits: Default::default(),
has_output_ordering: false,
expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
file_metadata_cache: None,
Expand Down Expand Up @@ -1098,7 +1173,7 @@ mod tests {
metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
df_metrics: ExecutionPlanMetricsSet::new(),
layout_readers: Default::default(),
natural_split_ranges: Default::default(),
natural_splits: Default::default(),
has_output_ordering: false,
expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
file_metadata_cache: None,
Expand Down Expand Up @@ -1185,7 +1260,7 @@ mod tests {
metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
df_metrics: ExecutionPlanMetricsSet::new(),
layout_readers: Default::default(),
natural_split_ranges: Default::default(),
natural_splits: Default::default(),
has_output_ordering: false,
expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
file_metadata_cache: None,
Expand Down Expand Up @@ -1342,7 +1417,7 @@ mod tests {
metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
df_metrics: ExecutionPlanMetricsSet::new(),
layout_readers: Default::default(),
natural_split_ranges: Default::default(),
natural_splits: Default::default(),
has_output_ordering: false,
expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
file_metadata_cache: None,
Expand Down Expand Up @@ -1402,7 +1477,7 @@ mod tests {
metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
df_metrics: ExecutionPlanMetricsSet::new(),
layout_readers: Default::default(),
natural_split_ranges: Default::default(),
natural_splits: Default::default(),
has_output_ordering: false,
expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
file_metadata_cache: None,
Expand Down Expand Up @@ -1611,7 +1686,7 @@ mod tests {
metrics_registry: Arc::new(DefaultMetricsRegistry::default()),
df_metrics: ExecutionPlanMetricsSet::new(),
layout_readers: Default::default(),
natural_split_ranges: Default::default(),
natural_splits: Default::default(),
has_output_ordering: false,
expression_convertor: Arc::new(DefaultExpressionConvertor::default()),
file_metadata_cache: None,
Expand Down
Loading
Loading