From 5827ae789b52e1bbf7f3139251cd5d563956e4bc Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 20 Sep 2026 22:09:06 +0800 Subject: [PATCH 1/5] feat(python): plan chunk-shuffled reads natively --- bindings/python/src/read.rs | 101 +- crates/paimon/src/table/chunk_shuffle.rs | 959 ++++++++++++++++++ .../paimon/src/table/data_evolution_reader.rs | 243 +++-- crates/paimon/src/table/data_file_reader.rs | 15 +- crates/paimon/src/table/incremental_scan.rs | 12 +- crates/paimon/src/table/mod.rs | 2 + crates/paimon/src/table/source.rs | 150 ++- crates/paimon/src/table/table_scan.rs | 112 +- 8 files changed, 1474 insertions(+), 120 deletions(-) create mode 100644 crates/paimon/src/table/chunk_shuffle.rs diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index 22af9ee61..dd71a1ad1 100644 --- a/bindings/python/src/read.rs +++ b/bindings/python/src/read.rs @@ -23,7 +23,9 @@ use arrow::pyarrow::ToPyArrow; use arrow::record_batch::RecordBatch; use futures::TryStreamExt; use paimon::spec::{DataField, DataType, Predicate, RowType}; -use paimon::table::{ArrowRecordBatchStream, DataSplit, IncrementalScanMode, RowRange, Table}; +use paimon::table::{ + ArrowRecordBatchStream, ChunkShuffle, DataSplit, IncrementalScanMode, RowRange, Table, +}; use paimon_datafusion::runtime::runtime; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; @@ -412,6 +414,7 @@ impl PyReadBuilder { incremental_range: None, row_position_slice: None, row_position_shard: None, + chunk_shuffle: None, } } @@ -449,6 +452,14 @@ pub struct PyTableScan { incremental_range: Option<(i64, i64)>, row_position_slice: Option<(u64, u64)>, row_position_shard: Option<(u64, u64)>, + chunk_shuffle: Option, +} + +#[derive(Clone)] +struct PyChunkShuffle { + seed: String, + chunk_size: u64, + shard: Option<(usize, usize)>, } impl PyTableScan { @@ -464,6 +475,15 @@ impl PyTableScan { .with_row_position_shard(index, count) .map_err(to_py_err)?; } + if let Some(chunk_shuffle) = &self.chunk_shuffle { + let mut config = + ChunkShuffle::from_decimal_seed(&chunk_shuffle.seed, chunk_shuffle.chunk_size) + .map_err(to_py_err)?; + if let Some((index, count)) = chunk_shuffle.shard { + config = config.with_shard(index, count).map_err(to_py_err)?; + } + scan = scan.with_chunk_shuffle(config).map_err(to_py_err)?; + } Ok(scan) } @@ -485,6 +505,15 @@ impl PyTableScan { .with_row_position_shard(index, count) .map_err(to_py_err)?; } + if let Some(chunk_shuffle) = &self.chunk_shuffle { + let mut config = + ChunkShuffle::from_decimal_seed(&chunk_shuffle.seed, chunk_shuffle.chunk_size) + .map_err(to_py_err)?; + if let Some((index, count)) = chunk_shuffle.shard { + config = config.with_shard(index, count).map_err(to_py_err)?; + } + scan = scan.with_chunk_shuffle(config).map_err(to_py_err)?; + } Ok(scan) } @@ -533,6 +562,42 @@ impl PyTableScan { Ok(slf) } + /// Deterministically shuffle fixed-live-row chunks, optionally selecting + /// one balanced worker shard. `seed` is a decimal Python integer string so + /// arbitrarily large seeds retain Python's `random.Random` semantics. + #[pyo3(signature = (seed, chunk_size, shard_index=None, shard_count=None))] + fn with_chunk_shuffle( + mut slf: PyRefMut<'_, Self>, + seed: String, + chunk_size: u64, + shard_index: Option, + shard_count: Option, + ) -> PyResult> { + let shard = match (shard_index, shard_count) { + (None, None) => None, + (Some(index), Some(count)) => Some((index, count)), + _ => { + return Err(PyValueError::new_err( + "chunk_shuffle shard_index and shard_count must be set together", + )); + } + }; + let mut config = ChunkShuffle::from_decimal_seed(&seed, chunk_size).map_err(to_py_err)?; + if let Some((index, count)) = shard { + config = config.with_shard(index, count).map_err(to_py_err)?; + } + // Validate every combination immediately, not only when plan() runs. + slf.core_scan()? + .with_chunk_shuffle(config) + .map_err(to_py_err)?; + slf.chunk_shuffle = Some(PyChunkShuffle { + seed, + chunk_size, + shard, + }); + Ok(slf) + } + fn plan(&self, py: Python<'_>) -> PyResult { py.detach(|| { runtime().block_on(async { @@ -773,6 +838,40 @@ impl PySplit { Ok(PyBytes::new(py, &bytes)) } + /// Serialize only the Java-compatible metadata view. Native-only file + /// ranges remain on this object and must be used for physical reading. + fn serialize_metadata<'py>(&self, py: Python<'py>) -> PyResult> { + let bytes = self + .inner + .serialize_split_v1_metadata_view() + .map_err(to_py_err)?; + Ok(PyBytes::new(py, &bytes)) + } + + /// Per-file local half-open ranges carried by native chunk planning. + fn file_row_ranges(&self) -> Option> { + let ranges = self.inner.file_row_ranges()?; + Some( + self.inner + .data_files() + .iter() + .zip(ranges) + .filter_map(|(file, range)| { + range.as_ref().map(|range| { + ( + file.file_name.clone(), + (range.from(), range.to().saturating_add(1)), + ) + }) + }) + .collect(), + ) + } + + fn exact_merged_row_count(&self) -> Option { + self.inner.exact_merged_row_count() + } + /// Reconstruct a native split from the stable, cross-language /// `SplitSerializer` v1 wire format. /// diff --git a/crates/paimon/src/table/chunk_shuffle.rs b/crates/paimon/src/table/chunk_shuffle.rs new file mode 100644 index 000000000..8e69d420d --- /dev/null +++ b/crates/paimon/src/table/chunk_shuffle.rs @@ -0,0 +1,959 @@ +// 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. + +//! Deterministic, fixed-row chunk planning for native Python reads. +//! +//! The shuffle intentionally matches `random.Random(seed).shuffle` in CPython. +//! PyPaimon exposed that ordering before native planning existed, so using a +//! different Rust RNG would silently assign different chunks to workers. + +use std::cmp::Ordering; +use std::collections::{BTreeMap, HashMap}; + +use crate::deletion_vector::{DeletionVector, DeletionVectorFactory}; +use crate::spec::{BinaryRow, DataField, DataFileMeta, Datum}; +use crate::table::source::{data_evolution_anchor_file, is_data_evolution_normal_file}; +use crate::table::stats_filter::group_by_overlapping_row_id; +use crate::table::{DataSplit, DataSplitBuilder, DeletionFile, RowRange, Table}; + +/// Native chunk-shuffle configuration. The seed is stored as the unsigned +/// little-endian 32-bit words consumed by CPython's MT19937 initializer. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ChunkShuffle { + seed_words: Vec, + chunk_size: i64, + shard: Option<(usize, usize)>, +} + +impl ChunkShuffle { + /// Build from a Python integer's decimal spelling. Negative integers use + /// their absolute value, matching `random.Random`. + pub fn from_decimal_seed(seed: &str, chunk_size: u64) -> crate::Result { + let chunk_size = i64::try_from(chunk_size).map_err(|_| crate::Error::DataInvalid { + message: format!("chunk_shuffle chunk_size {chunk_size} exceeds i64::MAX"), + source: None, + })?; + if chunk_size == 0 { + return Err(crate::Error::DataInvalid { + message: "chunk_shuffle chunk_size must be positive".to_string(), + source: None, + }); + } + Ok(Self { + seed_words: decimal_seed_words(seed)?, + chunk_size, + shard: None, + }) + } + + pub fn with_shard(mut self, index: usize, count: usize) -> crate::Result { + if count == 0 || index >= count { + return Err(crate::Error::DataInvalid { + message: "chunk_shuffle shard count must be positive and index less than count" + .to_string(), + source: None, + }); + } + self.shard = Some((index, count)); + Ok(self) + } +} + +#[derive(Debug, Clone)] +struct InputFile { + file: DataFileMeta, + deletion_file: Option, +} + +#[derive(Debug)] +struct InputGroup { + partition: BinaryRow, + bucket: i32, + bucket_path: String, + total_buckets: i32, + snapshot_id: i64, + is_streaming: bool, + files: Vec, +} + +#[derive(Debug)] +struct AppendSegment { + input: InputFile, + range: Option, + live_rows: i64, +} + +#[derive(Debug)] +struct EvolutionSegment { + files: Vec, + range: RowRange, + live_rows: i64, +} + +/// Repack planned files into shuffled, fixed-live-row chunks. Normal scan +/// planning runs first, so partition/stats/projection pruning and deletion-file +/// resolution remain centralized in `TableScan`. +pub(crate) async fn chunk_shuffle_splits( + table: &Table, + splits: Vec, + config: &ChunkShuffle, +) -> crate::Result> { + if !table.schema().primary_keys().is_empty() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle only supports append tables".to_string(), + }); + } + if splits + .iter() + .any(|split| split.row_ranges().is_some() || split.file_row_ranges().is_some()) + { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle cannot combine with row-range selection".to_string(), + }); + } + + let partition_fields = partition_fields(table)?; + let mut groups = flatten_groups(splits)?; + groups.sort_by(|left, right| { + compare_partitions(&left.partition, &right.partition, &partition_fields) + .unwrap_or_else(|_| { + left.partition + .to_serialized_bytes() + .cmp(&right.partition.to_serialized_bytes()) + }) + .then_with(|| left.bucket.cmp(&right.bucket)) + }); + + let data_evolution = table.schema().core_options().data_evolution_enabled(); + let mut chunks = Vec::new(); + for mut group in groups { + if data_evolution { + group.files.sort_by(|left, right| { + left.file + .first_row_id + .cmp(&right.file.first_row_id) + .then_with(|| { + is_data_evolution_normal_file(&right.file) + .cmp(&is_data_evolution_normal_file(&left.file)) + }) + .then_with(|| left.file.file_name.cmp(&right.file.file_name)) + }); + chunks.extend(evolution_chunks(table, group, config.chunk_size).await?); + } else { + group + .files + .sort_by(|left, right| left.file.file_name.cmp(&right.file.file_name)); + chunks.extend(append_chunks(table, group, config.chunk_size).await?); + } + } + + PythonRandom::new(&config.seed_words).shuffle(&mut chunks); + if let Some((index, count)) = config.shard { + let (start, end) = shard_range(chunks.len(), index, count); + chunks = chunks.drain(start..end).collect(); + } + Ok(chunks) +} + +fn partition_fields(table: &Table) -> crate::Result> { + let fields = table.schema().fields(); + table + .schema() + .partition_keys() + .iter() + .map(|name| { + fields + .iter() + .find(|field| field.name() == name) + .cloned() + .ok_or_else(|| crate::Error::DataInvalid { + message: format!("partition field '{name}' does not exist"), + source: None, + }) + }) + .collect() +} + +fn compare_partitions( + left: &BinaryRow, + right: &BinaryRow, + fields: &[DataField], +) -> crate::Result { + for (index, field) in fields.iter().enumerate() { + let left = left.get_datum(index, field.data_type())?; + let right = right.get_datum(index, field.data_type())?; + // Python's key is `(value is None, value)`: non-null sorts first. + let ordering = match (left, right) { + (None, None) => Ordering::Equal, + (None, Some(_)) => Ordering::Greater, + (Some(_), None) => Ordering::Less, + (Some(left), Some(right)) => compare_datums(&left, &right), + }; + if ordering != Ordering::Equal { + return Ok(ordering); + } + } + Ok(Ordering::Equal) +} + +fn compare_datums(left: &Datum, right: &Datum) -> Ordering { + left.partial_cmp(right).unwrap_or_else(|| { + // NaN has no total order. Python's stable sort leaves incomparable + // values in input order; the binary representation is a deterministic + // fallback when manifest concurrency changed that input order. + left.to_string().cmp(&right.to_string()) + }) +} + +fn flatten_groups(splits: Vec) -> crate::Result> { + let mut grouped: BTreeMap<(Vec, i32), InputGroup> = BTreeMap::new(); + for split in splits { + let deletion_files = split.data_deletion_files(); + let key = (split.partition().to_serialized_bytes(), split.bucket()); + let group = grouped.entry(key).or_insert_with(|| InputGroup { + partition: split.partition().clone(), + bucket: split.bucket(), + bucket_path: split.bucket_path().to_string(), + total_buckets: split.total_buckets(), + snapshot_id: split.snapshot_id(), + is_streaming: split.is_streaming(), + files: Vec::new(), + }); + if group.bucket_path != split.bucket_path() + || group.total_buckets != split.total_buckets() + || group.snapshot_id != split.snapshot_id() + || group.is_streaming != split.is_streaming() + { + return Err(crate::Error::DataInvalid { + message: "inconsistent split metadata within a partition bucket".to_string(), + source: None, + }); + } + for (index, file) in split.data_files().iter().cloned().enumerate() { + group.files.push(InputFile { + file, + deletion_file: deletion_files + .and_then(|files| files.get(index)) + .cloned() + .flatten(), + }); + } + } + Ok(grouped.into_values().collect()) +} + +async fn append_chunks( + table: &Table, + mut group: InputGroup, + chunk_size: i64, +) -> crate::Result> { + let mut chunks: Vec> = Vec::new(); + let mut current = Vec::new(); + let mut current_rows = 0; + + let inputs = std::mem::take(&mut group.files); + for input in inputs { + let mut slicer = match live_row_slicer(table, &input).await? { + Some(slicer) => slicer, + None => continue, + }; + loop { + if current_rows == chunk_size { + chunks.push(std::mem::take(&mut current)); + current_rows = 0; + } + let Some(slice) = slicer.take(chunk_size - current_rows)? else { + break; + }; + let range = if slice.start == 0 && slice.end == input.file.row_count { + None + } else { + Some(RowRange::new(slice.start, slice.end - 1)) + }; + current_rows += slice.live_rows; + current.push(AppendSegment { + input: input.clone(), + range, + live_rows: slice.live_rows, + }); + } + } + if !current.is_empty() { + chunks.push(current); + } + + chunks + .into_iter() + .map(|segments| build_append_split(&group, segments)) + .collect() +} + +fn build_append_split( + group: &InputGroup, + segments: Vec, +) -> crate::Result { + let exact_count = segments.iter().map(|segment| segment.live_rows).sum(); + let files = segments + .iter() + .map(|segment| segment.input.file.clone()) + .collect(); + let deletion_files: Vec<_> = segments + .iter() + .map(|segment| segment.input.deletion_file.clone()) + .collect(); + let file_ranges: Vec<_> = segments + .iter() + .map(|segment| segment.range.clone()) + .collect(); + let has_ranges = file_ranges.iter().any(Option::is_some); + + let mut builder = base_builder(group, files, true).with_exact_merged_row_count(exact_count); + if deletion_files.iter().any(Option::is_some) { + builder = builder.with_data_deletion_files(deletion_files); + } + if has_ranges { + builder = builder.with_file_row_ranges(file_ranges); + } + builder.build() +} + +async fn evolution_chunks( + table: &Table, + group: InputGroup, + chunk_size: i64, +) -> crate::Result> { + if group + .files + .iter() + .any(|input| input.file.first_row_id.is_none()) + { + return Err(crate::Error::DataInvalid { + message: "chunk_shuffle for data evolution requires first_row_id on every file" + .to_string(), + source: None, + }); + } + let deletion_by_name: HashMap<_, _> = group + .files + .iter() + .map(|input| (input.file.file_name.clone(), input.deletion_file.clone())) + .collect(); + let files = group.files.iter().map(|input| input.file.clone()).collect(); + let aligned_groups = group_by_overlapping_row_id(files); + let mut chunks: Vec> = Vec::new(); + let mut current = Vec::new(); + let mut current_rows = 0; + + for files in aligned_groups { + let from = files + .iter() + .filter_map(|file| file.first_row_id) + .min() + .ok_or_else(|| crate::Error::DataInvalid { + message: "data evolution chunk group has no first_row_id".to_string(), + source: None, + })?; + let to = files + .iter() + .filter_map(DataFileMeta::row_id_range) + .map(|(_, to)| to) + .max() + .ok_or_else(|| crate::Error::DataInvalid { + message: "data evolution chunk group has no row range".to_string(), + source: None, + })?; + let anchor = data_evolution_anchor_file(&files)?; + let anchor_deletion = deletion_by_name.get(&anchor.file_name).cloned().flatten(); + let (first_row_id, physical_count) = if anchor_deletion.is_some() { + let (anchor_from, anchor_to) = anchor.row_id_range().unwrap(); + if anchor_from > from || anchor_to < to { + return Err(crate::Error::DataInvalid { + message: format!( + "data evolution anchor range [{anchor_from}, {anchor_to}] does not contain group [{from}, {to}]" + ), + source: None, + }); + } + (anchor_from, anchor.row_count) + } else { + (from, to - from + 1) + }; + let anchor_input = InputFile { + file: anchor.clone(), + deletion_file: anchor_deletion.clone(), + }; + let slicer = if anchor_deletion.is_some() { + live_row_slicer(table, &anchor_input).await? + } else { + LiveRowSlicer::new(physical_count, Vec::new())? + }; + let mut slicer = match slicer { + Some(slicer) => slicer, + None => continue, + }; + if slicer.physical_count != physical_count { + return Err(crate::Error::DataInvalid { + message: format!( + "data evolution anchor row count {} does not match group row count {physical_count}", + slicer.physical_count + ), + source: None, + }); + } + let inputs: Vec<_> = files + .into_iter() + .map(|file| InputFile { + deletion_file: deletion_by_name.get(&file.file_name).cloned().flatten(), + file, + }) + .collect(); + loop { + if current_rows == chunk_size { + chunks.push(std::mem::take(&mut current)); + current_rows = 0; + } + let Some(slice) = slicer.take(chunk_size - current_rows)? else { + break; + }; + current_rows += slice.live_rows; + current.push(EvolutionSegment { + files: inputs.clone(), + range: RowRange::new(first_row_id + slice.start, first_row_id + slice.end - 1), + live_rows: slice.live_rows, + }); + } + } + if !current.is_empty() { + chunks.push(current); + } + + chunks + .into_iter() + .map(|segments| build_evolution_split(&group, segments)) + .collect() +} + +fn build_evolution_split( + group: &InputGroup, + segments: Vec, +) -> crate::Result { + let exact_count = segments.iter().map(|segment| segment.live_rows).sum(); + let mut files = Vec::new(); + let mut deletion_files = Vec::new(); + let mut ranges = Vec::new(); + for segment in segments { + ranges.push(segment.range); + for input in segment.files { + files.push(input.file); + deletion_files.push(input.deletion_file); + } + } + ranges.sort_by_key(RowRange::from); + let mut builder = base_builder(group, files, false) + .with_row_ranges(ranges) + .with_exact_merged_row_count(exact_count); + if deletion_files.iter().any(Option::is_some) { + builder = builder.with_data_deletion_files(deletion_files); + } + builder.build() +} + +fn base_builder(group: &InputGroup, files: Vec, raw: bool) -> DataSplitBuilder { + DataSplitBuilder::new() + .with_snapshot(group.snapshot_id) + .with_partition(group.partition.clone()) + .with_bucket(group.bucket) + .with_bucket_path(group.bucket_path.clone()) + .with_total_buckets(group.total_buckets) + .with_data_files(files) + .with_raw_convertible(raw) + .with_streaming(group.is_streaming) +} + +#[derive(Debug)] +struct PhysicalSlice { + start: i64, + end: i64, + live_rows: i64, +} + +#[derive(Debug)] +struct LiveRowSlicer { + physical_count: i64, + deleted: Vec, + deleted_index: usize, + position: i64, +} + +impl LiveRowSlicer { + fn new(physical_count: i64, deleted: Vec) -> crate::Result> { + if physical_count < 0 { + return Err(crate::Error::DataInvalid { + message: format!("negative physical row count {physical_count}"), + source: None, + }); + } + if deleted + .iter() + .any(|position| *position < 0 || *position >= physical_count) + { + return Err(crate::Error::DataInvalid { + message: "deletion vector position is outside the data file".to_string(), + source: None, + }); + } + if deleted.windows(2).any(|pair| pair[0] >= pair[1]) { + return Err(crate::Error::DataInvalid { + message: "deletion vector positions must be strictly increasing".to_string(), + source: None, + }); + } + if deleted.len() as i64 == physical_count { + return Ok(None); + } + Ok(Some(Self { + physical_count, + deleted, + deleted_index: 0, + position: 0, + })) + } + + fn take(&mut self, expected_live_rows: i64) -> crate::Result> { + if expected_live_rows <= 0 { + return Err(crate::Error::DataInvalid { + message: "chunk slice must request a positive number of live rows".to_string(), + source: None, + }); + } + if self.position >= self.physical_count { + return Ok(None); + } + let start = self.position; + let mut live_rows = 0; + while self.position < self.physical_count { + let next_deleted = self.deleted.get(self.deleted_index).copied(); + if let Some(deleted) = next_deleted { + let live_run = deleted - self.position; + let take = (expected_live_rows - live_rows).min(live_run); + self.position += take; + live_rows += take; + } else { + let take = + (expected_live_rows - live_rows).min(self.physical_count - self.position); + self.position += take; + live_rows += take; + } + if live_rows == expected_live_rows { + self.skip_deleted_at_cursor(); + return Ok(Some(PhysicalSlice { + start, + end: self.position, + live_rows, + })); + } + self.skip_deleted_at_cursor(); + } + if live_rows == 0 { + Ok(None) + } else { + Ok(Some(PhysicalSlice { + start, + end: self.position, + live_rows, + })) + } + } + + fn skip_deleted_at_cursor(&mut self) { + while self.deleted.get(self.deleted_index) == Some(&self.position) { + self.position += 1; + self.deleted_index += 1; + } + } +} + +async fn live_row_slicer(table: &Table, input: &InputFile) -> crate::Result> { + let Some(deletion_file) = &input.deletion_file else { + return LiveRowSlicer::new(input.file.row_count, Vec::new()); + }; + if deletion_file.cardinality() == Some(0) { + return LiveRowSlicer::new(input.file.row_count, Vec::new()); + } + if let Some(cardinality) = deletion_file.cardinality() { + if cardinality < 0 || cardinality > input.file.row_count { + return Err(crate::Error::DataInvalid { + message: format!( + "deletion vector cardinality {cardinality} is outside [0, {}]", + input.file.row_count + ), + source: None, + }); + } + } + let vector = DeletionVectorFactory::read(table.file_io(), deletion_file).await?; + validate_cardinality(deletion_file, &vector)?; + let deleted = vector.iter().map(|position| position as i64).collect(); + LiveRowSlicer::new(input.file.row_count, deleted) +} + +fn validate_cardinality(file: &DeletionFile, vector: &DeletionVector) -> crate::Result<()> { + if let Some(expected) = file.cardinality() { + if expected as u64 != vector.cardinality() { + return Err(crate::Error::DataInvalid { + message: format!( + "deletion vector cardinality mismatch: metadata {expected}, bitmap {}", + vector.cardinality() + ), + source: None, + }); + } + } + Ok(()) +} + +fn shard_range(total: usize, index: usize, count: usize) -> (usize, usize) { + let base = total / count; + let remainder = total % count; + let start = index * base + index.min(remainder); + (start, start + base + usize::from(index < remainder)) +} + +fn decimal_seed_words(seed: &str) -> crate::Result> { + let digits = seed + .strip_prefix('-') + .or_else(|| seed.strip_prefix('+')) + .unwrap_or(seed); + if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(crate::Error::DataInvalid { + message: format!("invalid integer chunk_shuffle seed '{seed}'"), + source: None, + }); + } + let mut words = vec![0u32]; + for digit in digits.bytes().map(|byte| u64::from(byte - b'0')) { + let mut carry = digit; + for word in &mut words { + let value = u64::from(*word) * 10 + carry; + *word = value as u32; + carry = value >> 32; + } + if carry != 0 { + words.push(carry as u32); + } + } + while words.len() > 1 && words.last() == Some(&0) { + words.pop(); + } + Ok(words) +} + +/// MT19937 plus CPython's integer-seed and `_randbelow_with_getrandbits` +/// conventions. Only up to 64 random bits are needed because slice lengths are +/// Rust `usize` values. +struct PythonRandom { + state: [u32; 624], + index: usize, +} + +impl PythonRandom { + fn new(seed_words: &[u32]) -> Self { + let mut random = Self { + state: [0; 624], + index: 624, + }; + random.init_genrand(19_650_218); + let mut i = 1usize; + let mut j = 0usize; + for _ in 0..624usize.max(seed_words.len()) { + random.state[i] = (random.state[i] + ^ (random.state[i - 1] ^ (random.state[i - 1] >> 30)).wrapping_mul(1_664_525)) + .wrapping_add(seed_words[j]) + .wrapping_add(j as u32); + i += 1; + j += 1; + if i >= 624 { + random.state[0] = random.state[623]; + i = 1; + } + if j >= seed_words.len() { + j = 0; + } + } + for _ in 0..623 { + random.state[i] = (random.state[i] + ^ (random.state[i - 1] ^ (random.state[i - 1] >> 30)).wrapping_mul(1_566_083_941)) + .wrapping_sub(i as u32); + i += 1; + if i >= 624 { + random.state[0] = random.state[623]; + i = 1; + } + } + random.state[0] = 0x8000_0000; + random + } + + fn init_genrand(&mut self, seed: u32) { + self.state[0] = seed; + for i in 1..624 { + self.state[i] = 1_812_433_253u32 + .wrapping_mul(self.state[i - 1] ^ (self.state[i - 1] >> 30)) + .wrapping_add(i as u32); + } + } + + fn gen_u32(&mut self) -> u32 { + if self.index >= 624 { + for i in 0..624 { + let y = (self.state[i] & 0x8000_0000) | (self.state[(i + 1) % 624] & 0x7fff_ffff); + self.state[i] = self.state[(i + 397) % 624] + ^ (y >> 1) + ^ if y & 1 == 0 { 0 } else { 0x9908_b0df }; + } + self.index = 0; + } + let mut y = self.state[self.index]; + self.index += 1; + y ^= y >> 11; + y ^= (y << 7) & 0x9d2c_5680; + y ^= (y << 15) & 0xefc6_0000; + y ^ (y >> 18) + } + + fn getrandbits(&mut self, bits: u32) -> u64 { + if bits <= 32 { + return u64::from(self.gen_u32() >> (32 - bits)); + } + let low = u64::from(self.gen_u32()); + let high_bits = bits - 32; + let high = u64::from(self.gen_u32() >> (32 - high_bits)); + low | (high << 32) + } + + fn randbelow(&mut self, n: usize) -> usize { + debug_assert!(n > 0); + let bits = usize::BITS - n.leading_zeros(); + loop { + let value = self.getrandbits(bits) as usize; + if value < n { + return value; + } + } + } + + fn shuffle(&mut self, values: &mut [T]) { + for index in (1..values.len()).rev() { + let other = self.randbelow(index + 1); + values.swap(index, other); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::catalog::Identifier; + use crate::io::FileIOBuilder; + use crate::spec::{stats::BinaryTableStats, DataType, IntType, Schema, TableSchema}; + + fn test_table(data_evolution: bool) -> Table { + let mut schema = Schema::builder().column("id", DataType::Int(IntType::new())); + if data_evolution { + schema = schema + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true"); + } + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "chunk_test"), + "memory:/chunk-test".to_string(), + TableSchema::new(0, &schema.build().unwrap()), + None, + ) + } + + fn file(name: &str, row_count: i64, first_row_id: Option) -> DataFileMeta { + DataFileMeta { + file_name: name.to_string(), + file_size: 100, + row_count, + min_key: Vec::new(), + max_key: Vec::new(), + key_stats: BinaryTableStats::new(Vec::new(), Vec::new(), Vec::new()), + value_stats: BinaryTableStats::new(Vec::new(), Vec::new(), Vec::new()), + min_sequence_number: 0, + max_sequence_number: 0, + schema_id: 0, + level: 0, + extra_files: Vec::new(), + creation_time: None, + delete_row_count: None, + embedded_index: None, + first_row_id, + write_cols: None, + external_path: None, + file_source: None, + value_stats_cols: None, + column_max_sequence_numbers: None, + } + } + + fn split(files: Vec, raw: bool) -> DataSplit { + DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path("memory:/chunk-test/bucket-0".to_string()) + .with_total_buckets(1) + .with_data_files(files) + .with_raw_convertible(raw) + .build() + .unwrap() + } + + #[test] + fn python_shuffle_matches_cpython_for_signed_and_large_integer_seeds() { + for (seed, expected) in [ + ("0", vec![7, 8, 1, 5, 3, 4, 2, 0, 9, 6]), + ("42", vec![7, 3, 2, 8, 5, 6, 9, 4, 0, 1]), + ("-11", vec![2, 6, 0, 1, 5, 4, 3, 9, 8, 7]), + ("1180591620717411303424", vec![4, 9, 5, 2, 1, 8, 0, 7, 6, 3]), + ] { + let words = decimal_seed_words(seed).unwrap(); + let mut actual: Vec<_> = (0..10).collect(); + PythonRandom::new(&words).shuffle(&mut actual); + assert_eq!(actual, expected, "seed={seed}"); + } + } + + #[test] + fn live_row_slicer_counts_visible_rows_and_attaches_boundary_deletes() { + let mut slicer = LiveRowSlicer::new(12, vec![0, 3, 4, 8, 11]) + .unwrap() + .unwrap(); + let first = slicer.take(3).unwrap().unwrap(); + assert_eq!((first.start, first.end, first.live_rows), (0, 6, 3)); + let second = slicer.take(3).unwrap().unwrap(); + assert_eq!((second.start, second.end, second.live_rows), (6, 10, 3)); + let last = slicer.take(3).unwrap().unwrap(); + assert_eq!((last.start, last.end, last.live_rows), (10, 12, 1)); + assert!(slicer.take(1).unwrap().is_none()); + } + + #[test] + fn shard_ranges_are_balanced_disjoint_and_cover_all_chunks() { + let ranges: Vec<_> = (0..5).map(|index| shard_range(12, index, 5)).collect(); + assert_eq!(ranges, vec![(0, 3), (3, 6), (6, 8), (8, 10), (10, 12)]); + } + + #[tokio::test] + async fn append_chunks_cover_every_file_position_once_and_shard_by_chunk() { + let table = test_table(false); + // File-name order, not the input split order, defines chunk positions. + let input = split( + vec![file("b.parquet", 4, None), file("a.parquet", 5, None)], + true, + ); + let config = ChunkShuffle::from_decimal_seed("42", 3).unwrap(); + let chunks = chunk_shuffle_splits(&table, vec![input.clone()], &config) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert!(chunks + .iter() + .all(|chunk| chunk.exact_merged_row_count() == Some(3))); + + let mut covered = Vec::new(); + for chunk in &chunks { + for (index, file) in chunk.data_files().iter().enumerate() { + let (from, to) = chunk + .file_row_range(index) + .map(|range| (range.from(), range.to())) + .unwrap_or((0, file.row_count - 1)); + covered.push((file.file_name.clone(), from, to)); + } + } + covered.sort(); + assert_eq!( + covered, + vec![ + ("a.parquet".to_string(), 0, 2), + ("a.parquet".to_string(), 3, 4), + ("b.parquet".to_string(), 0, 0), + ("b.parquet".to_string(), 1, 3), + ] + ); + + let left = chunk_shuffle_splits( + &table, + vec![input.clone()], + &config.clone().with_shard(0, 2).unwrap(), + ) + .await + .unwrap(); + let right = chunk_shuffle_splits( + &table, + vec![input], + &config.clone().with_shard(1, 2).unwrap(), + ) + .await + .unwrap(); + assert_eq!([left, right].concat(), chunks); + } + + #[tokio::test] + async fn data_evolution_chunks_keep_aligned_files_and_global_ranges() { + let table = test_table(true); + let input = split( + vec![ + file("base.parquet", 5, Some(10)), + file("payload.blob", 5, Some(10)), + file("next.parquet", 4, Some(20)), + ], + false, + ); + let chunks = chunk_shuffle_splits( + &table, + vec![input], + &ChunkShuffle::from_decimal_seed("0", 3).unwrap(), + ) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert!(chunks + .iter() + .all(|chunk| chunk.exact_merged_row_count() == Some(3))); + + let mut ranges: Vec<_> = chunks + .iter() + .flat_map(|chunk| chunk.row_ranges().unwrap()) + .map(|range| (range.from(), range.to())) + .collect(); + ranges.sort(); + assert_eq!(ranges, vec![(10, 12), (13, 14), (20, 20), (21, 23)]); + for chunk in chunks { + for range in chunk.row_ranges().unwrap() { + if range.from() < 20 { + assert!(chunk + .data_files() + .iter() + .any(|file| file.file_name == "payload.blob")); + } + } + } + } +} diff --git a/crates/paimon/src/table/data_evolution_reader.rs b/crates/paimon/src/table/data_evolution_reader.rs index ddd80975c..8bfe68f56 100644 --- a/crates/paimon/src/table/data_evolution_reader.rs +++ b/crates/paimon/src/table/data_evolution_reader.rs @@ -34,6 +34,7 @@ use crate::spec::{ use crate::table::dedicated_format_file_writer::is_blob_file_name; use crate::table::schema_manager::SchemaManager; use crate::table::source::any_range_overlaps_file; +use crate::table::stats_filter::group_by_overlapping_row_id; use crate::table::{ArrowRecordBatchStream, RESTEnv, RowRange}; use crate::{DataSplit, Error}; use arrow_array::{Array, Int64Array, LargeBinaryArray, RecordBatch}; @@ -85,6 +86,20 @@ fn is_raw_convertible(files: &[DataFileMeta]) -> bool { true } +/// Split a reader input into independently readable row-id groups. +/// +/// Native chunk shuffle can deliberately combine several disjoint groups in +/// one split. If any group needs column merging, checking only the whole split +/// is insufficient: all normal files would then be sent to one merge group, +/// which requires a single shared row-id range. +fn reader_file_groups(files: &[DataFileMeta]) -> Vec> { + if is_raw_convertible(files) { + vec![files.to_vec()] + } else { + group_by_overlapping_row_id(files.to_vec()) + } +} + /// Reads data files in data evolution mode, merging columns from files /// that share the same logical row range. pub(crate) struct DataEvolutionReader { @@ -291,61 +306,133 @@ impl DataEvolutionReader { for split in splits { let row_ranges = split.row_ranges().map(|r| r.to_vec()); + // A chunk may span several disjoint row-id groups while one + // group still needs column-wise merging. Process each aligned + // group independently; treating the whole chunk as one merge + // group would require unrelated normal files to share a range. + let file_groups = reader_file_groups(split.data_files()); + + for files in file_groups { + if is_raw_convertible(&files) { + for file_meta in files { + let deletion_vector = read_file_deletion_vector( + &self.file_io, + &split, + &file_meta, + ) + .await?; + let data_fields = raw_file_physical_fields( + &self.schema_manager, + self.table_schema_id, + &self.table_fields, + &file_meta, + ) + .await?; + + let has_row_id = file_meta.first_row_id.is_some(); + let effective_row_ranges = if has_row_id { row_ranges.clone() } else { None }; - if is_raw_convertible(split.data_files()) { - for file_meta in split.data_files().to_vec() { - let deletion_vector = read_file_deletion_vector( + let selected_row_ids = if self.row_id_index.is_some() && has_row_id { + selected_absolute_row_ranges_for_file( + file_meta.first_row_id.unwrap(), + file_meta.row_count, + effective_row_ranges.as_deref(), + deletion_vector.as_deref(), + )? + .map(|ranges| { + expand_selected_row_ids( + file_meta.first_row_id.unwrap(), + file_meta.row_count, + &ranges, + ) + }) + } else { + None + }; + let file_base_row_id = file_meta.first_row_id.unwrap_or(0); + let mut row_id_cursor = file_base_row_id; + let mut row_id_offset: usize = 0; + + let mut stream = raw_file_reader.read_single_file_stream( + &split, + file_meta, + data_fields, + deletion_vector, + effective_row_ranges, + )?; + while let Some(batch) = stream.next().await { + let batch = batch?; + let num_rows = batch.num_rows(); + let batch = if let Some(idx) = self.row_id_index { + if !has_row_id { + append_null_row_id_column(batch, idx, &self.wide_output_schema)? + } else if let Some(ref ids) = selected_row_ids { + attach_row_id(batch, idx, ids, &mut row_id_offset, &self.wide_output_schema)? + } else { + let row_ids: Vec = (row_id_cursor..row_id_cursor + num_rows as i64).collect(); + row_id_cursor += num_rows as i64; + let array: Arc = Arc::new(Int64Array::from(row_ids)); + insert_column_at(batch, array, idx, &self.wide_output_schema)? + } + } else { + batch + }; + yield self.finish_wide_batch( + batch, + blob_view_lookup.as_ref(), + &descriptor_fields, + filter_before_blob_resolution, + ).await?; + } + } + } else { + let prepared_group = PreparedMergeGroup::new(&files)?; + let anchor_deletion_vector = read_anchor_deletion_vector( &self.file_io, &split, - &file_meta, - ) - .await?; - let data_fields = raw_file_physical_fields( - &self.schema_manager, - self.table_schema_id, - &self.table_fields, - &file_meta, + &prepared_group.files, ) .await?; + let effective_row_ranges = row_ranges.clone(); + let selected_ranges = selected_absolute_row_ranges_for_file( + prepared_group.first_row_id, + prepared_group.logical_row_count, + effective_row_ranges.as_deref(), + anchor_deletion_vector + .as_ref() + .map(|ctx| ctx.deletion_vector.as_ref()), + )?; + let expected_output_rows = match selected_ranges.as_ref() { + Some(ranges) => ranges.iter().map(|r| r.count() as usize).sum(), + None => prepared_group.logical_row_count as usize, + }; - let has_row_id = file_meta.first_row_id.is_some(); - let effective_row_ranges = if has_row_id { row_ranges.clone() } else { None }; - - let selected_row_ids = if self.row_id_index.is_some() && has_row_id { - selected_absolute_row_ranges_for_file( - file_meta.first_row_id.unwrap(), - file_meta.row_count, - effective_row_ranges.as_deref(), - deletion_vector.as_deref(), - )? - .map(|ranges| { + let selected_row_ids = if self.row_id_index.is_some() { + selected_ranges.as_ref().map(|ranges| { expand_selected_row_ids( - file_meta.first_row_id.unwrap(), - file_meta.row_count, - &ranges, + prepared_group.first_row_id, + prepared_group.logical_row_count, + ranges, ) }) } else { None }; - let file_base_row_id = file_meta.first_row_id.unwrap_or(0); - let mut row_id_cursor = file_base_row_id; + let mut row_id_cursor = prepared_group.first_row_id; let mut row_id_offset: usize = 0; - let mut stream = raw_file_reader.read_single_file_stream( + let mut merge_stream = self.merge_files_by_columns( &split, - file_meta, - data_fields, - deletion_vector, + &prepared_group, effective_row_ranges, + expected_output_rows, + anchor_deletion_vector, )?; - while let Some(batch) = stream.next().await { + while let Some(batch) = merge_stream.next().await { let batch = batch?; let num_rows = batch.num_rows(); let batch = if let Some(idx) = self.row_id_index { - if !has_row_id { - append_null_row_id_column(batch, idx, &self.wide_output_schema)? - } else if let Some(ref ids) = selected_row_ids { + if let Some(ref ids) = selected_row_ids { attach_row_id(batch, idx, ids, &mut row_id_offset, &self.wide_output_schema)? } else { let row_ids: Vec = (row_id_cursor..row_id_cursor + num_rows as i64).collect(); @@ -364,71 +451,6 @@ impl DataEvolutionReader { ).await?; } } - } else { - let prepared_group = PreparedMergeGroup::new(split.data_files())?; - let anchor_deletion_vector = read_anchor_deletion_vector( - &self.file_io, - &split, - &prepared_group.files, - ) - .await?; - let effective_row_ranges = row_ranges.clone(); - let selected_ranges = selected_absolute_row_ranges_for_file( - prepared_group.first_row_id, - prepared_group.logical_row_count, - effective_row_ranges.as_deref(), - anchor_deletion_vector - .as_ref() - .map(|ctx| ctx.deletion_vector.as_ref()), - )?; - let expected_output_rows = match selected_ranges.as_ref() { - Some(ranges) => ranges.iter().map(|r| r.count() as usize).sum(), - None => prepared_group.logical_row_count as usize, - }; - - let selected_row_ids = if self.row_id_index.is_some() { - selected_ranges.as_ref().map(|ranges| { - expand_selected_row_ids( - prepared_group.first_row_id, - prepared_group.logical_row_count, - ranges, - ) - }) - } else { - None - }; - let mut row_id_cursor = prepared_group.first_row_id; - let mut row_id_offset: usize = 0; - - let mut merge_stream = self.merge_files_by_columns( - &split, - &prepared_group, - effective_row_ranges, - expected_output_rows, - anchor_deletion_vector, - )?; - while let Some(batch) = merge_stream.next().await { - let batch = batch?; - let num_rows = batch.num_rows(); - let batch = if let Some(idx) = self.row_id_index { - if let Some(ref ids) = selected_row_ids { - attach_row_id(batch, idx, ids, &mut row_id_offset, &self.wide_output_schema)? - } else { - let row_ids: Vec = (row_id_cursor..row_id_cursor + num_rows as i64).collect(); - row_id_cursor += num_rows as i64; - let array: Arc = Arc::new(Int64Array::from(row_ids)); - insert_column_at(batch, array, idx, &self.wide_output_schema)? - } - } else { - batch - }; - yield self.finish_wide_batch( - batch, - blob_view_lookup.as_ref(), - &descriptor_fields, - filter_before_blob_resolution, - ).await?; - } } } } @@ -3055,6 +3077,27 @@ mod tests { assert!(!is_raw_convertible(&files)); } + #[test] + fn test_reader_file_groups_separates_disjoint_merge_groups() { + let files = vec![ + data_file("base-0.parquet", 0, 10, 1, Some(vec!["id"])), + data_file("partial-0.parquet", 0, 10, 2, Some(vec!["value"])), + data_file("base-10.parquet", 10, 10, 1, Some(vec!["id"])), + data_file("partial-10.parquet", 10, 10, 2, Some(vec!["value"])), + ]; + + assert!(!is_raw_convertible(&files)); + let groups = reader_file_groups(&files); + assert_eq!(groups.len(), 2); + assert_eq!(groups[0].len(), 2); + assert_eq!(groups[1].len(), 2); + assert_eq!(PreparedMergeGroup::new(&groups[0]).unwrap().first_row_id, 0); + assert_eq!( + PreparedMergeGroup::new(&groups[1]).unwrap().first_row_id, + 10 + ); + } + #[test] fn test_prepared_merge_group_rejects_vector_only_split() { // No normal anchor file -> DataInvalid. diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 0b3cb03e5..305f51eb9 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -278,7 +278,7 @@ impl DataFileReader { // Create DV factory for this split only. let dv_factory = reader.build_split_dv_factory(&split).await?; - for file_meta in split.data_files().to_vec() { + for (file_index, file_meta) in split.data_files().to_vec().into_iter().enumerate() { let dv = DataFileReader::deletion_vector_for_file( dv_factory.as_ref(), &file_meta.file_name, @@ -310,13 +310,24 @@ impl DataFileReader { FileIndexResult::Remain }; - let split_ranges = split.row_ranges().map(|ranges| { + let global_ranges = split.row_ranges().map(|ranges| { to_local_row_ranges( ranges, file_meta.first_row_id.unwrap_or(0), file_meta.row_count, ) }); + let local_ranges = split + .file_row_range(file_index) + .map(|range| vec![range.clone()]); + let split_ranges = match (global_ranges, local_ranges) { + (Some(global), Some(local)) => { + Some(intersect_sorted_ranges(&global, &local)) + } + (Some(global), None) => Some(global), + (None, Some(local)) => Some(local), + (None, None) => None, + }; let selected_ranges = match file_index_result { FileIndexResult::Remain => split_ranges, FileIndexResult::Skip => Some(Vec::new()), diff --git a/crates/paimon/src/table/incremental_scan.rs b/crates/paimon/src/table/incremental_scan.rs index c4e640409..c71deb177 100644 --- a/crates/paimon/src/table/incremental_scan.rs +++ b/crates/paimon/src/table/incremental_scan.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{DataSplit, Plan, SnapshotManager, Table, TableScan}; +use super::{ChunkShuffle, DataSplit, Plan, SnapshotManager, Table, TableScan}; use crate::spec::{CommitKind, CoreOptions}; /// Batch incremental scan mode. @@ -256,11 +256,17 @@ impl<'a> IncrementalScan<'a> { Ok(self) } + /// Repack the combined APPEND-delta batch into deterministic chunks. + pub fn with_chunk_shuffle(mut self, config: ChunkShuffle) -> crate::Result { + self.scan = self.scan.with_chunk_shuffle(config)?; + Ok(self) + } + pub async fn plan(&self) -> crate::Result { crate::spec::CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; - if self.scan.has_row_position_selection() { + if self.scan.has_row_position_selection() || self.scan.has_chunk_shuffle() { return Err(crate::Error::Unsupported { - message: "Incremental row-position selection requires combined delta planning" + message: "Incremental row-position selection and chunk_shuffle require combined delta planning" .into(), }); } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 407f9e0b2..a4b51a7eb 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -33,6 +33,7 @@ mod bucket_assigner_dynamic; mod bucket_assigner_fixed; mod bucket_filter; mod bucket_function; +mod chunk_shuffle; mod commit_message; mod consumer_manager; pub(crate) mod cow_writer; @@ -130,6 +131,7 @@ pub use audit_log_table::AuditLogTable; pub use batch_vector_search_builder::BatchVectorSearchBuilder; pub use blob_resolver::{BlobReader, BlobStream}; pub use branch_manager::BranchManager; +pub use chunk_shuffle::ChunkShuffle; pub use commit_message::CommitMessage; pub use consumer_manager::ConsumerManager; pub use cow_writer::{CopyOnWriteMergeWriter, FileInfo}; diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 4499fd011..b77dc182e 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -496,6 +496,17 @@ pub struct DataSplit { /// `None` at index `i` means no deletion file for `data_files[i]` (matches Java getDeletionFiles() / List with null elements). data_deletion_files: Option]>>, row_ranges: Option>, + /// Optional file-local inclusive row range for each data file. `None` at + /// index `i` means the complete file. This is native-only planning + /// metadata used by chunk-shuffled append scans; Java's split wire format + /// cannot represent it. + #[serde(default)] + file_row_ranges: Option]>>, + /// Exact number of visible rows when planning had to inspect deletion + /// vectors. This is a planning hint, not part of the stable split wire + /// format. + #[serde(default)] + exact_merged_row_count: Option, /// Whether the split can be read raw, without the merge reader: its /// physical rows are exactly its logical rows (modulo deletion files). /// Mirrors Java `DataSplit#rawConvertible`. @@ -539,6 +550,23 @@ impl DataSplit { self.row_ranges.as_deref() } + /// File-local inclusive row ranges aligned with [`Self::data_files`]. + pub fn file_row_ranges(&self) -> Option<&[Option]> { + self.file_row_ranges.as_deref() + } + + /// File-local range for the data file at `index`; `None` means full file. + pub fn file_row_range(&self, index: usize) -> Option<&RowRange> { + self.file_row_ranges + .as_deref() + .and_then(|ranges| ranges.get(index)) + .and_then(Option::as_ref) + } + + pub fn exact_merged_row_count(&self) -> Option { + self.exact_merged_row_count + } + /// Whether this split can be read raw (no sort-merge needed); see the /// field doc. Mirrors Java `DataSplit#rawConvertible`. pub fn raw_convertible(&self) -> bool { @@ -619,6 +647,9 @@ impl DataSplit { /// /// Reference: [DataSplit.mergedRowCount()](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/source/DataSplit.java#L133) pub fn merged_row_count(&self) -> Option { + if let Some(count) = self.exact_merged_row_count { + return Some(count); + } if !self.row_counts_known() { return None; } @@ -694,6 +725,19 @@ impl DataSplit { /// Byte-compatible with `compatibility/datasplit-v9`. Row ranges are not part of the /// format; `serialize_split_v1` wraps a row-range split as an `IndexedSplit` instead. pub fn serialize(&self) -> crate::Result> { + if self.file_row_ranges.is_some() { + return Err(crate::Error::Unsupported { + message: "Java DataSplit serialization cannot represent file-local row ranges" + .to_string(), + }); + } + self.serialize_metadata_view() + } + + /// Serialize the Java-compatible base split while intentionally omitting + /// native-only planning metadata. Callers must retain the original native + /// split for physical reading. + pub fn serialize_metadata_view(&self) -> crate::Result> { let mut out = Vec::new(); out.extend_from_slice(&SPLIT_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_VERSION.to_be_bytes()); @@ -849,20 +893,31 @@ impl DataSplit { /// `IndexedSplit` (type 3) wrapping the DataSplit body plus the ranges. Byte-compatible with /// `compatibility/split-v1-data` / `split-v1-indexed`. pub fn serialize_split_v1(&self) -> crate::Result> { + if self.file_row_ranges.is_some() { + return Err(crate::Error::Unsupported { + message: "SplitSerializer v1 cannot represent file-local row ranges".to_string(), + }); + } + self.serialize_split_v1_metadata_view() + } + + /// Serialize a cross-language metadata view of this split. Native-only + /// file ranges and exact row counts are intentionally omitted. + pub fn serialize_split_v1_metadata_view(&self) -> crate::Result> { let mut out = Vec::new(); out.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); match &self.row_ranges { None => { out.extend_from_slice(&SPLIT_SER_TYPE_DATA_SPLIT.to_be_bytes()); - out.extend_from_slice(&self.serialize()?); + out.extend_from_slice(&self.serialize_metadata_view()?); } Some(ranges) => { out.extend_from_slice(&SPLIT_SER_TYPE_INDEXED_SPLIT.to_be_bytes()); // IndexedSplit#serialize: magic + version + DataSplit body + ranges + scores. out.extend_from_slice(&INDEXED_SPLIT_MAGIC.to_be_bytes()); out.extend_from_slice(&INDEXED_SPLIT_VERSION.to_be_bytes()); - out.extend_from_slice(&self.serialize()?); + out.extend_from_slice(&self.serialize_metadata_view()?); out.extend_from_slice(&(ranges.len() as i32).to_be_bytes()); for r in ranges.iter() { out.extend_from_slice(&r.from().to_be_bytes()); @@ -1165,6 +1220,8 @@ pub struct DataSplitBuilder { /// Same length as data_files; `None` at index i = no deletion file for data_files[i]. data_deletion_files: Option>>, row_ranges: Option>, + file_row_ranges: Option>>, + exact_merged_row_count: Option, raw_convertible: bool, is_streaming: bool, } @@ -1180,6 +1237,8 @@ impl DataSplitBuilder { data_files: None, data_deletion_files: None, row_ranges: None, + file_row_ranges: None, + exact_merged_row_count: None, // Splits with no merge semantics (append tables, single-file // utility splits) are raw by nature; the merge-tree and // data-evolution scan paths set this explicitly per split group. @@ -1227,6 +1286,18 @@ impl DataSplitBuilder { self } + /// Set file-local inclusive ranges aligned with `data_files`. `None` + /// selects the whole file at that position. + pub fn with_file_row_ranges(mut self, ranges: Vec>) -> Self { + self.file_row_ranges = Some(ranges); + self + } + + pub fn with_exact_merged_row_count(mut self, count: i64) -> Self { + self.exact_merged_row_count = Some(count); + self + } + /// Preserve the Java DataSplit event-reading contract. pub fn with_streaming(mut self, is_streaming: bool) -> Self { self.is_streaming = is_streaming; @@ -1282,6 +1353,40 @@ impl DataSplitBuilder { }); } } + if let Some(ranges) = &self.file_row_ranges { + if ranges.len() != data_files.len() { + return Err(crate::Error::DataInvalid { + message: format!( + "file_row_ranges length {} does not match data_files length {}", + ranges.len(), + data_files.len() + ), + source: None, + }); + } + for (file, range) in data_files.iter().zip(ranges) { + if let Some(range) = range { + if range.from() < 0 || range.to() >= file.row_count { + return Err(crate::Error::DataInvalid { + message: format!( + "file-local row range [{}, {}] is outside file '{}' row count {}", + range.from(), + range.to(), + file.file_name, + file.row_count + ), + source: None, + }); + } + } + } + } + if self.exact_merged_row_count.is_some_and(|count| count < 0) { + return Err(crate::Error::DataInvalid { + message: "exact_merged_row_count must be non-negative".to_string(), + source: None, + }); + } Ok(DataSplit { snapshot_id: self.snapshot_id, partition: Arc::new(partition), @@ -1291,6 +1396,8 @@ impl DataSplitBuilder { data_files: data_files.into(), data_deletion_files: self.data_deletion_files.map(Into::into), row_ranges: self.row_ranges.map(Into::into), + file_row_ranges: self.file_row_ranges.map(Into::into), + exact_merged_row_count: self.exact_merged_row_count, raw_convertible: self.raw_convertible, is_streaming: self.is_streaming, }) @@ -2230,6 +2337,45 @@ mod tests { ); } + #[test] + fn native_file_ranges_are_validated_preserved_and_never_silently_serialized() { + let split = v1_data_split_builder() + .with_file_row_ranges(vec![Some(RowRange::new(2, 5)), None]) + .with_exact_merged_row_count(15) + .build() + .unwrap(); + assert_eq!(split.file_row_range(0), Some(&RowRange::new(2, 5))); + assert_eq!(split.merged_row_count(), Some(15)); + assert!(matches!( + split.serialize_split_v1(), + Err(crate::Error::Unsupported { .. }) + )); + + let metadata = split.serialize_split_v1_metadata_view().unwrap(); + let decoded = DataSplit::deserialize_split_v1(&metadata).unwrap(); + assert!(decoded.file_row_ranges().is_none()); + assert_eq!(decoded.exact_merged_row_count(), None); + assert_eq!(decoded.data_files(), split.data_files()); + + let pickle = serde_json::to_vec(&split).unwrap(); + let restored: DataSplit = serde_json::from_slice(&pickle).unwrap(); + assert_eq!(restored, split); + + assert!(v1_data_split_builder() + .with_file_row_ranges(vec![Some(RowRange::new(0, 10))]) + .build() + .is_err()); + assert!(v1_data_split_builder() + .with_file_row_ranges(vec![Some(RowRange::new(0, 100)), None]) + .build() + .is_err()); + assert!(v1_data_split_builder() + .with_file_row_ranges(vec![Some(RowRange::new(0, 10)), None]) + .with_exact_merged_row_count(-1) + .build() + .is_err()); + } + #[test] fn serialize_indexed_split_v1_matches_golden() { // Row ranges -> IndexedSplit (type 3): same DataSplit as split-v1-data + ranges [1,4],[11,13]. diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index bfc9f2b1c..ddf5139ee 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -21,6 +21,7 @@ //! and [FullStartingScanner](https://github.com/apache/paimon/blob/release-1.3/paimon-python/pypaimon/read/scanner/full_starting_scanner.py). use super::bucket_filter::compute_target_buckets; +use super::chunk_shuffle::{chunk_shuffle_splits, ChunkShuffle}; use super::format_table_scan::FormatTableScan; use super::global_index_scanner::RowRangeIndex; use super::global_index_types::normalize_queryable_global_index_type; @@ -1136,20 +1137,67 @@ impl<'a> TableScan<'a> { self.with_row_position_selection(RowPositionSelection::shard(index, count)?) } + /// Repack an append scan into deterministic fixed-live-row chunks. + pub fn with_chunk_shuffle(self, config: ChunkShuffle) -> crate::Result { + match self.0 { + TableScanKind::Paimon(mut scan) => { + if !scan.table.schema().primary_keys().is_empty() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle only supports append tables".to_string(), + }); + } + if scan.limit.is_some() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle cannot combine with limit".to_string(), + }); + } + if scan.row_ranges.is_some() || scan.row_position_selection().is_some() { + return Err(crate::Error::Unsupported { + message: + "chunk_shuffle cannot combine with row ranges or positional selection" + .to_string(), + }); + } + if !scan.data_predicates.is_empty() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle only supports partition predicates".to_string(), + }); + } + scan.split_selection = Some(Box::new(ScanSplitSelection::ChunkShuffle(config))); + Ok(Self(TableScanKind::Paimon(scan))) + } + TableScanKind::Format(_) => Err(crate::Error::Unsupported { + message: "format tables do not support chunk_shuffle".to_string(), + }), + } + } + pub(crate) fn has_row_position_selection(&self) -> bool { match &self.0 { - TableScanKind::Paimon(scan) => scan.row_position_selection.is_some(), + TableScanKind::Paimon(scan) => scan.row_position_selection().is_some(), TableScanKind::Format(_) => false, } } + pub(crate) fn has_chunk_shuffle(&self) -> bool { + matches!(&self.0, TableScanKind::Paimon(scan) if scan.chunk_shuffle().is_some()) + } + fn with_row_position_selection(self, selection: RowPositionSelection) -> crate::Result { match self.0 { TableScanKind::Paimon(mut scan) if scan.table.schema().core_options().data_evolution_enabled() => { + if scan.chunk_shuffle().is_some() { + return Err(crate::Error::DataInvalid { + message: + "row-position selection and chunk_shuffle cannot be used simultaneously" + .into(), + source: None, + }); + } if scan - .row_position_selection + .row_position_selection() .is_some_and(|previous| previous.is_slice() != selection.is_slice()) { return Err(crate::Error::DataInvalid { @@ -1158,7 +1206,7 @@ impl<'a> TableScan<'a> { source: None, }); } - scan.row_position_selection = Some(selection); + scan.split_selection = Some(Box::new(ScanSplitSelection::RowPosition(selection))); Ok(Self(TableScanKind::Paimon(scan))) } _ => Err(crate::Error::Unsupported { @@ -1264,6 +1312,12 @@ impl<'a> TableScan<'a> { /// Paimon table scan: resolves snapshots, reads manifests, and builds data splits. /// /// Reference: [pypaimon.read.table_scan.TableScan](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/read/table_scan.py) +#[derive(Debug, Clone)] +enum ScanSplitSelection { + RowPosition(RowPositionSelection), + ChunkShuffle(ChunkShuffle), +} + #[derive(Debug, Clone)] struct PaimonTableScan<'a> { table: &'a Table, @@ -1274,7 +1328,10 @@ struct PaimonTableScan<'a> { /// When set, the scan will try to return only enough splits to satisfy the limit. limit: Option, row_ranges: Option>, - row_position_selection: Option, + /// Mutually exclusive row-position or fixed-row chunk transformation. + /// Boxed because scans normally use neither, and keeping the cold payload + /// out of the scan preserves the compact `TableScanKind` representation. + split_selection: Option>, /// Diff compares complete logical states, so it must not accept physical /// row-range pruning from an explicit range or a global-index lookup. row_range_optimization_disabled: bool, @@ -1287,6 +1344,20 @@ struct PaimonTableScan<'a> { } impl<'a> PaimonTableScan<'a> { + fn row_position_selection(&self) -> Option { + match self.split_selection.as_deref() { + Some(ScanSplitSelection::RowPosition(selection)) => Some(*selection), + _ => None, + } + } + + fn chunk_shuffle(&self) -> Option<&ChunkShuffle> { + match self.split_selection.as_deref() { + Some(ScanSplitSelection::ChunkShuffle(config)) => Some(config), + _ => None, + } + } + fn is_streaming(&self) -> bool { self.incremental_split_mode.is_some() } @@ -1306,7 +1377,7 @@ impl<'a> PaimonTableScan<'a> { bucket_predicate, limit, row_ranges, - row_position_selection: None, + split_selection: None, row_range_optimization_disabled: false, scan_all_files: false, incremental_split_mode: None, @@ -1335,7 +1406,7 @@ impl<'a> PaimonTableScan<'a> { fn without_row_range_optimization(mut self) -> Self { self.row_ranges = None; - self.row_position_selection = None; + self.split_selection = None; self.row_range_optimization_disabled = true; self } @@ -1826,7 +1897,7 @@ impl<'a> PaimonTableScan<'a> { // Positional scans count all candidate row IDs. Pruning a preceding // manifest or file here would renumber the surviving rows; intersect // explicit/global-index ranges after assigning positions instead. - let row_range_index = if data_evolution_enabled && self.row_position_selection.is_none() { + let row_range_index = if data_evolution_enabled && self.row_position_selection().is_none() { manifest_row_ranges.clone().map(RowRangeIndex::create) } else { None @@ -2121,7 +2192,7 @@ impl<'a> PaimonTableScan<'a> { // Positional scans count all candidate row IDs. Pruning a preceding // manifest or file here would renumber the surviving rows; intersect // explicit/global-index ranges after assigning positions instead. - let row_range_index = if data_evolution_enabled && self.row_position_selection.is_none() { + let row_range_index = if data_evolution_enabled && self.row_position_selection().is_none() { manifest_row_ranges.clone().map(RowRangeIndex::create) } else { None @@ -2177,7 +2248,7 @@ impl<'a> PaimonTableScan<'a> { // Assign row positions using the full candidate file ranges, before // group stats, projection or DVs change visible rows. Intersect explicit // and index-selected ranges only after assigning the positional range. - let effective_row_ranges = if let Some(selection) = self.row_position_selection { + let effective_row_ranges = if let Some(selection) = self.row_position_selection() { Some(selection.select(&entries, effective_row_ranges.as_deref())?) } else { effective_row_ranges @@ -2395,7 +2466,7 @@ impl<'a> PaimonTableScan<'a> { row_id_groups }; - if self.row_position_selection.is_some() { + if self.row_position_selection().is_some() { // Positional scans promise row-id order across groups. A // projected group can become a singleton, so moving all // multi-file groups first would change which rows a limit @@ -2557,6 +2628,11 @@ impl<'a> PaimonTableScan<'a> { let split_candidates_built = splits.len(); (splits, split_candidates_built, false) }; + let splits = if let Some(config) = self.chunk_shuffle() { + chunk_shuffle_splits(self.table, splits, config).await? + } else { + splits + }; let splits_before_limit = split_candidates_built; if let Some(trace) = trace { let final_files = splits.iter().map(|split| split.data_files().len()).sum(); @@ -2581,8 +2657,8 @@ mod tests { prune_data_evolution_group_by_read_fields, retain_index_manifest_entry, retain_index_manifest_entry_for_scan, retain_manifest_buckets, retain_manifest_entry_row_ranges, retain_manifest_row_ranges, scan_predicate_field_ids, - should_skip_level_zero_for_scan, split_row_ranges_for_files, LimitPushdownAccumulator, - PaimonTableScan, RowRangeIndex, TableScan, + should_skip_level_zero_for_scan, split_row_ranges_for_files, ChunkShuffle, + LimitPushdownAccumulator, PaimonTableScan, RowRangeIndex, TableScan, }; use crate::catalog::Identifier; use crate::io::FileIOBuilder; @@ -3680,6 +3756,18 @@ mod tests { .unwrap() .with_row_position_shard(1, 2) .is_ok()); + assert!(reader + .new_scan() + .with_row_position_shard(0, 1) + .unwrap() + .with_chunk_shuffle(ChunkShuffle::from_decimal_seed("0", 1).unwrap()) + .is_err()); + assert!(reader + .new_scan() + .with_chunk_shuffle(ChunkShuffle::from_decimal_seed("0", 1).unwrap()) + .unwrap() + .with_row_position_shard(0, 1) + .is_err()); } #[tokio::test] From 0ed36be1338008bd43b27d8311cf4604c83b8e0c Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Sun, 20 Sep 2026 22:48:45 +0800 Subject: [PATCH 2/5] refactor(python): simplify chunk shuffle scan API --- bindings/python/src/read.rs | 68 ++++++++++----------- bindings/python/tests/test_read.py | 35 +++++++++++ crates/paimon/src/table/chunk_shuffle.rs | 28 ++++----- crates/paimon/src/table/incremental_scan.rs | 16 ++++- crates/paimon/src/table/mod.rs | 1 - crates/paimon/src/table/table_scan.rs | 46 ++++++++++++-- 6 files changed, 137 insertions(+), 57 deletions(-) diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index dd71a1ad1..87dec1d95 100644 --- a/bindings/python/src/read.rs +++ b/bindings/python/src/read.rs @@ -23,9 +23,7 @@ use arrow::pyarrow::ToPyArrow; use arrow::record_batch::RecordBatch; use futures::TryStreamExt; use paimon::spec::{DataField, DataType, Predicate, RowType}; -use paimon::table::{ - ArrowRecordBatchStream, ChunkShuffle, DataSplit, IncrementalScanMode, RowRange, Table, -}; +use paimon::table::{ArrowRecordBatchStream, DataSplit, IncrementalScanMode, RowRange, Table}; use paimon_datafusion::runtime::runtime; use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; use pyo3::prelude::*; @@ -476,13 +474,14 @@ impl PyTableScan { .map_err(to_py_err)?; } if let Some(chunk_shuffle) = &self.chunk_shuffle { - let mut config = - ChunkShuffle::from_decimal_seed(&chunk_shuffle.seed, chunk_shuffle.chunk_size) - .map_err(to_py_err)?; + scan = scan + .with_chunk_shuffle(&chunk_shuffle.seed, chunk_shuffle.chunk_size) + .map_err(to_py_err)?; if let Some((index, count)) = chunk_shuffle.shard { - config = config.with_shard(index, count).map_err(to_py_err)?; + scan = scan + .with_chunk_shuffle_shard(index, count) + .map_err(to_py_err)?; } - scan = scan.with_chunk_shuffle(config).map_err(to_py_err)?; } Ok(scan) } @@ -506,13 +505,14 @@ impl PyTableScan { .map_err(to_py_err)?; } if let Some(chunk_shuffle) = &self.chunk_shuffle { - let mut config = - ChunkShuffle::from_decimal_seed(&chunk_shuffle.seed, chunk_shuffle.chunk_size) - .map_err(to_py_err)?; + scan = scan + .with_chunk_shuffle(&chunk_shuffle.seed, chunk_shuffle.chunk_size) + .map_err(to_py_err)?; if let Some((index, count)) = chunk_shuffle.shard { - config = config.with_shard(index, count).map_err(to_py_err)?; + scan = scan + .with_chunk_shuffle_shard(index, count) + .map_err(to_py_err)?; } - scan = scan.with_chunk_shuffle(config).map_err(to_py_err)?; } Ok(scan) } @@ -562,42 +562,42 @@ impl PyTableScan { Ok(slf) } - /// Deterministically shuffle fixed-live-row chunks, optionally selecting - /// one balanced worker shard. `seed` is a decimal Python integer string so - /// arbitrarily large seeds retain Python's `random.Random` semantics. - #[pyo3(signature = (seed, chunk_size, shard_index=None, shard_count=None))] + /// Deterministically shuffle fixed-live-row chunks. `seed` is a decimal + /// Python integer string so arbitrarily large seeds retain Python's + /// `random.Random` semantics. fn with_chunk_shuffle( mut slf: PyRefMut<'_, Self>, seed: String, chunk_size: u64, - shard_index: Option, - shard_count: Option, ) -> PyResult> { - let shard = match (shard_index, shard_count) { - (None, None) => None, - (Some(index), Some(count)) => Some((index, count)), - _ => { - return Err(PyValueError::new_err( - "chunk_shuffle shard_index and shard_count must be set together", - )); - } - }; - let mut config = ChunkShuffle::from_decimal_seed(&seed, chunk_size).map_err(to_py_err)?; - if let Some((index, count)) = shard { - config = config.with_shard(index, count).map_err(to_py_err)?; - } // Validate every combination immediately, not only when plan() runs. slf.core_scan()? - .with_chunk_shuffle(config) + .with_chunk_shuffle(&seed, chunk_size) .map_err(to_py_err)?; slf.chunk_shuffle = Some(PyChunkShuffle { seed, chunk_size, - shard, + shard: None, }); Ok(slf) } + /// Select one balanced worker shard after chunk shuffling. + fn with_chunk_shuffle_shard( + mut slf: PyRefMut<'_, Self>, + index: usize, + count: usize, + ) -> PyResult> { + slf.core_scan()? + .with_chunk_shuffle_shard(index, count) + .map_err(to_py_err)?; + let chunk_shuffle = slf.chunk_shuffle.as_mut().ok_or_else(|| { + PyValueError::new_err("with_chunk_shuffle_shard requires with_chunk_shuffle first") + })?; + chunk_shuffle.shard = Some((index, count)); + Ok(slf) + } + fn plan(&self, py: Python<'_>) -> PyResult { py.detach(|| { runtime().block_on(async { diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index c072f32ed..e9337fdbb 100644 --- a/bindings/python/tests/test_read.py +++ b/bindings/python/tests/test_read.py @@ -79,6 +79,41 @@ def test_with_blob_parallelism(): table.new_read_builder().with_blob_parallelism(0) +def test_chunk_shuffle_takes_seed_and_chunk_size_before_optional_shard(): + with tempfile.TemporaryDirectory() as warehouse: + table = _make_table_with_data(warehouse) + builder = table.new_read_builder().with_projection(["id"]) + + scan = builder.new_scan() + assert scan.with_chunk_shuffle(str(2 ** 70), 2) is scan + plan = scan.plan() + chunks = [ + pa.Table.from_batches(builder.new_read().read([split])) + .column("id").to_pylist() + for split in plan.splits() + ] + assert sorted(value for chunk in chunks for value in chunk) == [1, 2, 3] + assert all(0 < len(chunk) <= 2 for chunk in chunks) + + sharded = [] + for index in range(2): + shard = ( + builder.new_scan() + .with_chunk_shuffle(str(2 ** 70), 2) + .with_chunk_shuffle_shard(index, 2) + .plan() + ) + sharded.extend( + pa.Table.from_batches(builder.new_read().read([split])) + .column("id").to_pylist() + for split in shard.splits() + ) + assert sharded == chunks + + with pytest.raises(ValueError, match="requires with_chunk_shuffle first"): + builder.new_scan().with_chunk_shuffle_shard(0, 2) + + def test_with_row_ranges(): with tempfile.TemporaryDirectory() as warehouse: ctx = SQLContext() diff --git a/crates/paimon/src/table/chunk_shuffle.rs b/crates/paimon/src/table/chunk_shuffle.rs index 8e69d420d..7baa4b567 100644 --- a/crates/paimon/src/table/chunk_shuffle.rs +++ b/crates/paimon/src/table/chunk_shuffle.rs @@ -32,7 +32,7 @@ use crate::table::{DataSplit, DataSplitBuilder, DeletionFile, RowRange, Table}; /// Native chunk-shuffle configuration. The seed is stored as the unsigned /// little-endian 32-bit words consumed by CPython's MT19937 initializer. #[derive(Debug, Clone, PartialEq, Eq)] -pub struct ChunkShuffle { +pub(crate) struct ChunkShuffle { seed_words: Vec, chunk_size: i64, shard: Option<(usize, usize)>, @@ -41,7 +41,7 @@ pub struct ChunkShuffle { impl ChunkShuffle { /// Build from a Python integer's decimal spelling. Negative integers use /// their absolute value, matching `random.Random`. - pub fn from_decimal_seed(seed: &str, chunk_size: u64) -> crate::Result { + pub(crate) fn from_decimal_seed(seed: &str, chunk_size: u64) -> crate::Result { let chunk_size = i64::try_from(chunk_size).map_err(|_| crate::Error::DataInvalid { message: format!("chunk_shuffle chunk_size {chunk_size} exceeds i64::MAX"), source: None, @@ -59,7 +59,7 @@ impl ChunkShuffle { }) } - pub fn with_shard(mut self, index: usize, count: usize) -> crate::Result { + pub(crate) fn set_shard(&mut self, index: usize, count: usize) -> crate::Result<()> { if count == 0 || index >= count { return Err(crate::Error::DataInvalid { message: "chunk_shuffle shard count must be positive and index less than count" @@ -68,7 +68,7 @@ impl ChunkShuffle { }); } self.shard = Some((index, count)); - Ok(self) + Ok(()) } } @@ -898,18 +898,18 @@ mod tests { ] ); - let left = chunk_shuffle_splits( - &table, - vec![input.clone()], - &config.clone().with_shard(0, 2).unwrap(), - ) + let left = chunk_shuffle_splits(&table, vec![input.clone()], &{ + let mut sharded = config.clone(); + sharded.set_shard(0, 2).unwrap(); + sharded + }) .await .unwrap(); - let right = chunk_shuffle_splits( - &table, - vec![input], - &config.clone().with_shard(1, 2).unwrap(), - ) + let right = chunk_shuffle_splits(&table, vec![input], &{ + let mut sharded = config.clone(); + sharded.set_shard(1, 2).unwrap(); + sharded + }) .await .unwrap(); assert_eq!([left, right].concat(), chunks); diff --git a/crates/paimon/src/table/incremental_scan.rs b/crates/paimon/src/table/incremental_scan.rs index c71deb177..8c2225f66 100644 --- a/crates/paimon/src/table/incremental_scan.rs +++ b/crates/paimon/src/table/incremental_scan.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{ChunkShuffle, DataSplit, Plan, SnapshotManager, Table, TableScan}; +use super::{DataSplit, Plan, SnapshotManager, Table, TableScan}; use crate::spec::{CommitKind, CoreOptions}; /// Batch incremental scan mode. @@ -257,8 +257,18 @@ impl<'a> IncrementalScan<'a> { } /// Repack the combined APPEND-delta batch into deterministic chunks. - pub fn with_chunk_shuffle(mut self, config: ChunkShuffle) -> crate::Result { - self.scan = self.scan.with_chunk_shuffle(config)?; + pub fn with_chunk_shuffle( + mut self, + seed: impl ToString, + chunk_size: u64, + ) -> crate::Result { + self.scan = self.scan.with_chunk_shuffle(seed, chunk_size)?; + Ok(self) + } + + /// Select one balanced worker shard after chunk shuffling. + pub fn with_chunk_shuffle_shard(mut self, index: usize, count: usize) -> crate::Result { + self.scan = self.scan.with_chunk_shuffle_shard(index, count)?; Ok(self) } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index a4b51a7eb..74581c019 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -131,7 +131,6 @@ pub use audit_log_table::AuditLogTable; pub use batch_vector_search_builder::BatchVectorSearchBuilder; pub use blob_resolver::{BlobReader, BlobStream}; pub use branch_manager::BranchManager; -pub use chunk_shuffle::ChunkShuffle; pub use commit_message::CommitMessage; pub use consumer_manager::ConsumerManager; pub use cow_writer::{CopyOnWriteMergeWriter, FileInfo}; diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index ddf5139ee..0b4753bf5 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1138,7 +1138,36 @@ impl<'a> TableScan<'a> { } /// Repack an append scan into deterministic fixed-live-row chunks. - pub fn with_chunk_shuffle(self, config: ChunkShuffle) -> crate::Result { + /// + /// `seed` must render as a decimal integer. Accepting any `ToString` seed + /// keeps ordinary Rust integer calls ergonomic while allowing language + /// bindings to preserve arbitrary-precision integer seeds. + pub fn with_chunk_shuffle(self, seed: impl ToString, chunk_size: u64) -> crate::Result { + let config = ChunkShuffle::from_decimal_seed(&seed.to_string(), chunk_size)?; + self.with_chunk_shuffle_config(config) + } + + /// Select one balanced worker shard after chunk shuffling. + pub fn with_chunk_shuffle_shard(mut self, index: usize, count: usize) -> crate::Result { + match &mut self.0 { + TableScanKind::Paimon(scan) => match scan.split_selection.as_deref_mut() { + Some(ScanSplitSelection::ChunkShuffle(config)) => { + config.set_shard(index, count)?; + Ok(self) + } + _ => Err(crate::Error::DataInvalid { + message: "with_chunk_shuffle_shard requires with_chunk_shuffle first" + .to_string(), + source: None, + }), + }, + TableScanKind::Format(_) => Err(crate::Error::Unsupported { + message: "format tables do not support chunk_shuffle".to_string(), + }), + } + } + + fn with_chunk_shuffle_config(self, config: ChunkShuffle) -> crate::Result { match self.0 { TableScanKind::Paimon(mut scan) => { if !scan.table.schema().primary_keys().is_empty() { @@ -2657,8 +2686,8 @@ mod tests { prune_data_evolution_group_by_read_fields, retain_index_manifest_entry, retain_index_manifest_entry_for_scan, retain_manifest_buckets, retain_manifest_entry_row_ranges, retain_manifest_row_ranges, scan_predicate_field_ids, - should_skip_level_zero_for_scan, split_row_ranges_for_files, ChunkShuffle, - LimitPushdownAccumulator, PaimonTableScan, RowRangeIndex, TableScan, + should_skip_level_zero_for_scan, split_row_ranges_for_files, LimitPushdownAccumulator, + PaimonTableScan, RowRangeIndex, TableScan, }; use crate::catalog::Identifier; use crate::io::FileIOBuilder; @@ -3760,14 +3789,21 @@ mod tests { .new_scan() .with_row_position_shard(0, 1) .unwrap() - .with_chunk_shuffle(ChunkShuffle::from_decimal_seed("0", 1).unwrap()) + .with_chunk_shuffle(0, 1) .is_err()); assert!(reader .new_scan() - .with_chunk_shuffle(ChunkShuffle::from_decimal_seed("0", 1).unwrap()) + .with_chunk_shuffle(0, 1) .unwrap() .with_row_position_shard(0, 1) .is_err()); + assert!(reader.new_scan().with_chunk_shuffle_shard(0, 1).is_err()); + assert!(reader + .new_scan() + .with_chunk_shuffle(0, 1) + .unwrap() + .with_chunk_shuffle_shard(0, 1) + .is_ok()); } #[tokio::test] From 67dfb4e504d7203b6b092dc80bf8cc0de74e0bf4 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 21 Sep 2026 00:00:21 +0800 Subject: [PATCH 3/5] refactor(python): reuse indexed ranges for chunk shuffle --- bindings/python/src/read.rs | 73 ++----- bindings/python/tests/test_read.py | 67 ++++++- crates/paimon/src/table/chunk_shuffle.rs | 205 ++++++++++---------- crates/paimon/src/table/data_file_reader.rs | 43 ++-- crates/paimon/src/table/incremental_scan.rs | 13 +- crates/paimon/src/table/source.rs | 163 ++-------------- crates/paimon/src/table/table_scan.rs | 133 ++++++++++--- 7 files changed, 343 insertions(+), 354 deletions(-) diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index 87dec1d95..c6605a2b7 100644 --- a/bindings/python/src/read.rs +++ b/bindings/python/src/read.rs @@ -413,6 +413,7 @@ impl PyReadBuilder { row_position_slice: None, row_position_shard: None, chunk_shuffle: None, + shard: None, } } @@ -451,13 +452,13 @@ pub struct PyTableScan { row_position_slice: Option<(u64, u64)>, row_position_shard: Option<(u64, u64)>, chunk_shuffle: Option, + shard: Option<(usize, usize)>, } #[derive(Clone)] struct PyChunkShuffle { seed: String, chunk_size: u64, - shard: Option<(usize, usize)>, } impl PyTableScan { @@ -477,11 +478,9 @@ impl PyTableScan { scan = scan .with_chunk_shuffle(&chunk_shuffle.seed, chunk_shuffle.chunk_size) .map_err(to_py_err)?; - if let Some((index, count)) = chunk_shuffle.shard { - scan = scan - .with_chunk_shuffle_shard(index, count) - .map_err(to_py_err)?; - } + } + if let Some((index, count)) = self.shard { + scan = scan.with_shard(index, count).map_err(to_py_err)?; } Ok(scan) } @@ -508,11 +507,9 @@ impl PyTableScan { scan = scan .with_chunk_shuffle(&chunk_shuffle.seed, chunk_shuffle.chunk_size) .map_err(to_py_err)?; - if let Some((index, count)) = chunk_shuffle.shard { - scan = scan - .with_chunk_shuffle_shard(index, count) - .map_err(to_py_err)?; - } + } + if let Some((index, count)) = self.shard { + scan = scan.with_shard(index, count).map_err(to_py_err)?; } Ok(scan) } @@ -574,27 +571,20 @@ impl PyTableScan { slf.core_scan()? .with_chunk_shuffle(&seed, chunk_size) .map_err(to_py_err)?; - slf.chunk_shuffle = Some(PyChunkShuffle { - seed, - chunk_size, - shard: None, - }); + slf.chunk_shuffle = Some(PyChunkShuffle { seed, chunk_size }); Ok(slf) } - /// Select one balanced worker shard after chunk shuffling. - fn with_chunk_shuffle_shard( + /// Select one balanced worker shard for a distributed scan. + fn with_shard( mut slf: PyRefMut<'_, Self>, index: usize, count: usize, ) -> PyResult> { slf.core_scan()? - .with_chunk_shuffle_shard(index, count) + .with_shard(index, count) .map_err(to_py_err)?; - let chunk_shuffle = slf.chunk_shuffle.as_mut().ok_or_else(|| { - PyValueError::new_err("with_chunk_shuffle_shard requires with_chunk_shuffle first") - })?; - chunk_shuffle.shard = Some((index, count)); + slf.shard = Some((index, count)); Ok(slf) } @@ -821,7 +811,8 @@ impl PySplit { #[pymethods] impl PySplit { - /// Physical row count: sum of data-file row counts (not a logical result count). + /// Selected row count for IndexedSplit-compatible row ranges, otherwise + /// the sum of physical data-file row counts. fn row_count(&self) -> i64 { self.inner.row_count() } @@ -838,40 +829,6 @@ impl PySplit { Ok(PyBytes::new(py, &bytes)) } - /// Serialize only the Java-compatible metadata view. Native-only file - /// ranges remain on this object and must be used for physical reading. - fn serialize_metadata<'py>(&self, py: Python<'py>) -> PyResult> { - let bytes = self - .inner - .serialize_split_v1_metadata_view() - .map_err(to_py_err)?; - Ok(PyBytes::new(py, &bytes)) - } - - /// Per-file local half-open ranges carried by native chunk planning. - fn file_row_ranges(&self) -> Option> { - let ranges = self.inner.file_row_ranges()?; - Some( - self.inner - .data_files() - .iter() - .zip(ranges) - .filter_map(|(file, range)| { - range.as_ref().map(|range| { - ( - file.file_name.clone(), - (range.from(), range.to().saturating_add(1)), - ) - }) - }) - .collect(), - ) - } - - fn exact_merged_row_count(&self) -> Option { - self.inner.exact_merged_row_count() - } - /// Reconstruct a native split from the stable, cross-language /// `SplitSerializer` v1 wire format. /// diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index e9337fdbb..3f265ceaa 100644 --- a/bindings/python/tests/test_read.py +++ b/bindings/python/tests/test_read.py @@ -87,6 +87,9 @@ def test_chunk_shuffle_takes_seed_and_chunk_size_before_optional_shard(): scan = builder.new_scan() assert scan.with_chunk_shuffle(str(2 ** 70), 2) is scan plan = scan.plan() + for split in plan.splits(): + restored = Split.deserialize(split.serialize()) + assert restored.row_count() == split.row_count() chunks = [ pa.Table.from_batches(builder.new_read().read([split])) .column("id").to_pylist() @@ -100,7 +103,7 @@ def test_chunk_shuffle_takes_seed_and_chunk_size_before_optional_shard(): shard = ( builder.new_scan() .with_chunk_shuffle(str(2 ** 70), 2) - .with_chunk_shuffle_shard(index, 2) + .with_shard(index, 2) .plan() ) sharded.extend( @@ -110,8 +113,64 @@ def test_chunk_shuffle_takes_seed_and_chunk_size_before_optional_shard(): ) assert sharded == chunks - with pytest.raises(ValueError, match="requires with_chunk_shuffle first"): - builder.new_scan().with_chunk_shuffle_shard(0, 2) + # Shard is scan-level state, so it may be configured before shuffle. + before_shuffle = ( + builder.new_scan() + .with_shard(0, 2) + .with_chunk_shuffle(str(2 ** 70), 2) + .plan() + ) + before_rows = [ + pa.Table.from_batches(builder.new_read().read([split])) + .column("id").to_pylist() + for split in before_shuffle.splits() + ] + after_rows = [ + pa.Table.from_batches(builder.new_read().read([split])) + .column("id").to_pylist() + for split in ( + builder.new_scan() + .with_chunk_shuffle(str(2 ** 70), 2) + .with_shard(0, 2) + .plan() + .splits() + ) + ] + assert before_rows == after_rows + + with pytest.raises(ValueError, match="count must be positive"): + builder.new_scan().with_shard(0, 0) + with pytest.raises(RuntimeError, match="requires chunk_shuffle"): + builder.new_scan().with_shard(0, 2).plan() + + +def test_chunk_shuffle_reads_split_local_ranges_across_files(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + ctx.sql("CREATE SCHEMA paimon.rdb") + ctx.sql("CREATE TABLE paimon.rdb.t (id INT)") + ctx.sql("INSERT INTO paimon.rdb.t VALUES (1), (2)") + ctx.sql("INSERT INTO paimon.rdb.t VALUES (3), (4)") + table = PaimonCatalog({"warehouse": warehouse}).get_table("rdb.t") + builder = table.new_read_builder().with_projection(["id"]) + splits = ( + builder.new_scan() + .with_chunk_shuffle("7", 3) + .plan() + .splits() + ) + + chunks = [] + for split in splits: + restored = Split.deserialize(split.serialize()) + chunks.append( + pa.Table.from_batches(builder.new_read().read([restored])) + .column("id").to_pylist() + ) + + assert sorted(len(chunk) for chunk in chunks) == [1, 3] + assert sorted(value for chunk in chunks for value in chunk) == [1, 2, 3, 4] def test_with_row_ranges(): @@ -240,7 +299,7 @@ def test_indexed_split_wire_roundtrip_preserves_row_ranges(): restored = Split.deserialize(split.serialize()) rows = pa.Table.from_batches(builder.new_read().read([restored])) - assert restored.row_count() == 3 + assert restored.row_count() == 1 assert rows.column("id").to_pylist() == [2] diff --git a/crates/paimon/src/table/chunk_shuffle.rs b/crates/paimon/src/table/chunk_shuffle.rs index 7baa4b567..f6a096636 100644 --- a/crates/paimon/src/table/chunk_shuffle.rs +++ b/crates/paimon/src/table/chunk_shuffle.rs @@ -27,7 +27,7 @@ use crate::deletion_vector::{DeletionVector, DeletionVectorFactory}; use crate::spec::{BinaryRow, DataField, DataFileMeta, Datum}; use crate::table::source::{data_evolution_anchor_file, is_data_evolution_normal_file}; use crate::table::stats_filter::group_by_overlapping_row_id; -use crate::table::{DataSplit, DataSplitBuilder, DeletionFile, RowRange, Table}; +use crate::table::{merge_row_ranges, DataSplit, DataSplitBuilder, DeletionFile, RowRange, Table}; /// Native chunk-shuffle configuration. The seed is stored as the unsigned /// little-endian 32-bit words consumed by CPython's MT19937 initializer. @@ -35,7 +35,6 @@ use crate::table::{DataSplit, DataSplitBuilder, DeletionFile, RowRange, Table}; pub(crate) struct ChunkShuffle { seed_words: Vec, chunk_size: i64, - shard: Option<(usize, usize)>, } impl ChunkShuffle { @@ -55,21 +54,8 @@ impl ChunkShuffle { Ok(Self { seed_words: decimal_seed_words(seed)?, chunk_size, - shard: None, }) } - - pub(crate) fn set_shard(&mut self, index: usize, count: usize) -> crate::Result<()> { - if count == 0 || index >= count { - return Err(crate::Error::DataInvalid { - message: "chunk_shuffle shard count must be positive and index less than count" - .to_string(), - source: None, - }); - } - self.shard = Some((index, count)); - Ok(()) - } } #[derive(Debug, Clone)] @@ -92,15 +78,13 @@ struct InputGroup { #[derive(Debug)] struct AppendSegment { input: InputFile, - range: Option, - live_rows: i64, + ranges: Vec, } #[derive(Debug)] struct EvolutionSegment { files: Vec, - range: RowRange, - live_rows: i64, + ranges: Vec, } /// Repack planned files into shuffled, fixed-live-row chunks. Normal scan @@ -110,16 +94,14 @@ pub(crate) async fn chunk_shuffle_splits( table: &Table, splits: Vec, config: &ChunkShuffle, + shard: Option<(usize, usize)>, ) -> crate::Result> { if !table.schema().primary_keys().is_empty() { return Err(crate::Error::Unsupported { message: "chunk_shuffle only supports append tables".to_string(), }); } - if splits - .iter() - .any(|split| split.row_ranges().is_some() || split.file_row_ranges().is_some()) - { + if splits.iter().any(|split| split.row_ranges().is_some()) { return Err(crate::Error::Unsupported { message: "chunk_shuffle cannot combine with row-range selection".to_string(), }); @@ -161,7 +143,7 @@ pub(crate) async fn chunk_shuffle_splits( } PythonRandom::new(&config.seed_words).shuffle(&mut chunks); - if let Some((index, count)) = config.shard { + if let Some((index, count)) = shard { let (start, end) = shard_range(chunks.len(), index, count); chunks = chunks.drain(start..end).collect(); } @@ -278,16 +260,10 @@ async fn append_chunks( let Some(slice) = slicer.take(chunk_size - current_rows)? else { break; }; - let range = if slice.start == 0 && slice.end == input.file.row_count { - None - } else { - Some(RowRange::new(slice.start, slice.end - 1)) - }; current_rows += slice.live_rows; current.push(AppendSegment { input: input.clone(), - range, - live_rows: slice.live_rows, + ranges: slice.ranges, }); } } @@ -305,28 +281,26 @@ fn build_append_split( group: &InputGroup, segments: Vec, ) -> crate::Result { - let exact_count = segments.iter().map(|segment| segment.live_rows).sum(); - let files = segments - .iter() - .map(|segment| segment.input.file.clone()) - .collect(); - let deletion_files: Vec<_> = segments - .iter() - .map(|segment| segment.input.deletion_file.clone()) - .collect(); - let file_ranges: Vec<_> = segments - .iter() - .map(|segment| segment.range.clone()) - .collect(); - let has_ranges = file_ranges.iter().any(Option::is_some); + let mut files = Vec::with_capacity(segments.len()); + let mut deletion_files = Vec::with_capacity(segments.len()); + let mut ranges = Vec::new(); + let mut split_offset = 0; + for segment in segments { + ranges.extend( + segment + .ranges + .into_iter() + .map(|range| RowRange::new(split_offset + range.from(), split_offset + range.to())), + ); + split_offset += segment.input.file.row_count; + files.push(segment.input.file); + deletion_files.push(segment.input.deletion_file); + } - let mut builder = base_builder(group, files, true).with_exact_merged_row_count(exact_count); + let mut builder = base_builder(group, files, true).with_row_ranges(merge_row_ranges(ranges)); if deletion_files.iter().any(Option::is_some) { builder = builder.with_data_deletion_files(deletion_files); } - if has_ranges { - builder = builder.with_file_row_ranges(file_ranges); - } builder.build() } @@ -431,8 +405,13 @@ async fn evolution_chunks( current_rows += slice.live_rows; current.push(EvolutionSegment { files: inputs.clone(), - range: RowRange::new(first_row_id + slice.start, first_row_id + slice.end - 1), - live_rows: slice.live_rows, + ranges: slice + .ranges + .into_iter() + .map(|range| { + RowRange::new(first_row_id + range.from(), first_row_id + range.to()) + }) + .collect(), }); } } @@ -450,21 +429,17 @@ fn build_evolution_split( group: &InputGroup, segments: Vec, ) -> crate::Result { - let exact_count = segments.iter().map(|segment| segment.live_rows).sum(); let mut files = Vec::new(); let mut deletion_files = Vec::new(); let mut ranges = Vec::new(); for segment in segments { - ranges.push(segment.range); + ranges.extend(segment.ranges); for input in segment.files { files.push(input.file); deletion_files.push(input.deletion_file); } } - ranges.sort_by_key(RowRange::from); - let mut builder = base_builder(group, files, false) - .with_row_ranges(ranges) - .with_exact_merged_row_count(exact_count); + let mut builder = base_builder(group, files, false).with_row_ranges(merge_row_ranges(ranges)); if deletion_files.iter().any(Option::is_some) { builder = builder.with_data_deletion_files(deletion_files); } @@ -485,8 +460,7 @@ fn base_builder(group: &InputGroup, files: Vec, raw: bool) -> Data #[derive(Debug)] struct PhysicalSlice { - start: i64, - end: i64, + ranges: Vec, live_rows: i64, } @@ -542,39 +516,37 @@ impl LiveRowSlicer { if self.position >= self.physical_count { return Ok(None); } - let start = self.position; let mut live_rows = 0; + let mut ranges = Vec::new(); while self.position < self.physical_count { let next_deleted = self.deleted.get(self.deleted_index).copied(); if let Some(deleted) = next_deleted { let live_run = deleted - self.position; let take = (expected_live_rows - live_rows).min(live_run); - self.position += take; - live_rows += take; + if take > 0 { + ranges.push(RowRange::new(self.position, self.position + take - 1)); + self.position += take; + live_rows += take; + } } else { let take = (expected_live_rows - live_rows).min(self.physical_count - self.position); - self.position += take; - live_rows += take; + if take > 0 { + ranges.push(RowRange::new(self.position, self.position + take - 1)); + self.position += take; + live_rows += take; + } } if live_rows == expected_live_rows { self.skip_deleted_at_cursor(); - return Ok(Some(PhysicalSlice { - start, - end: self.position, - live_rows, - })); + return Ok(Some(PhysicalSlice { ranges, live_rows })); } self.skip_deleted_at_cursor(); } if live_rows == 0 { Ok(None) } else { - Ok(Some(PhysicalSlice { - start, - end: self.position, - live_rows, - })) + Ok(Some(PhysicalSlice { ranges, live_rows })) } } @@ -841,17 +813,40 @@ mod tests { } #[test] - fn live_row_slicer_counts_visible_rows_and_attaches_boundary_deletes() { + fn live_row_slicer_returns_only_visible_physical_ranges() { let mut slicer = LiveRowSlicer::new(12, vec![0, 3, 4, 8, 11]) .unwrap() .unwrap(); let first = slicer.take(3).unwrap().unwrap(); - assert_eq!((first.start, first.end, first.live_rows), (0, 6, 3)); + assert_eq!( + (first.ranges, first.live_rows), + (vec![RowRange::new(1, 2), RowRange::new(5, 5)], 3) + ); let second = slicer.take(3).unwrap().unwrap(); - assert_eq!((second.start, second.end, second.live_rows), (6, 10, 3)); + assert_eq!( + (second.ranges, second.live_rows), + (vec![RowRange::new(6, 7), RowRange::new(9, 9)], 3) + ); let last = slicer.take(3).unwrap().unwrap(); - assert_eq!((last.start, last.end, last.live_rows), (10, 12, 1)); + assert_eq!( + (last.ranges, last.live_rows), + (vec![RowRange::new(10, 10)], 1) + ); assert!(slicer.take(1).unwrap().is_none()); + + let mut alternating = LiveRowSlicer::new(8, vec![1, 3, 5, 7]).unwrap().unwrap(); + assert_eq!( + alternating.take(3).unwrap().unwrap().ranges, + vec![ + RowRange::new(0, 0), + RowRange::new(2, 2), + RowRange::new(4, 4) + ] + ); + assert_eq!( + alternating.take(3).unwrap().unwrap().ranges, + vec![RowRange::new(6, 6)] + ); } #[test] @@ -869,23 +864,38 @@ mod tests { true, ); let config = ChunkShuffle::from_decimal_seed("42", 3).unwrap(); - let chunks = chunk_shuffle_splits(&table, vec![input.clone()], &config) + let chunks = chunk_shuffle_splits(&table, vec![input.clone()], &config, None) .await .unwrap(); assert_eq!(chunks.len(), 3); assert!(chunks .iter() - .all(|chunk| chunk.exact_merged_row_count() == Some(3))); + .all(|chunk| chunk.row_count() == 3 && chunk.merged_row_count() == Some(3))); let mut covered = Vec::new(); for chunk in &chunks { - for (index, file) in chunk.data_files().iter().enumerate() { - let (from, to) = chunk - .file_row_range(index) - .map(|range| (range.from(), range.to())) - .unwrap_or((0, file.row_count - 1)); - covered.push((file.file_name.clone(), from, to)); + let mut split_offset = 0; + for file in chunk.data_files() { + for range in chunk.row_ranges().unwrap() { + let from = range.from().max(split_offset); + let to = range.to().min(split_offset + file.row_count - 1); + if from <= to { + covered.push(( + file.file_name.clone(), + from - split_offset, + to - split_offset, + )); + } + } + split_offset += file.row_count; } + let serialized = chunk.serialize_split_v1().unwrap(); + assert_eq!( + DataSplit::deserialize_split_v1(&serialized) + .unwrap() + .row_ranges(), + chunk.row_ranges() + ); } covered.sort(); assert_eq!( @@ -898,20 +908,12 @@ mod tests { ] ); - let left = chunk_shuffle_splits(&table, vec![input.clone()], &{ - let mut sharded = config.clone(); - sharded.set_shard(0, 2).unwrap(); - sharded - }) - .await - .unwrap(); - let right = chunk_shuffle_splits(&table, vec![input], &{ - let mut sharded = config.clone(); - sharded.set_shard(1, 2).unwrap(); - sharded - }) - .await - .unwrap(); + let left = chunk_shuffle_splits(&table, vec![input.clone()], &config, Some((0, 2))) + .await + .unwrap(); + let right = chunk_shuffle_splits(&table, vec![input], &config, Some((1, 2))) + .await + .unwrap(); assert_eq!([left, right].concat(), chunks); } @@ -930,13 +932,14 @@ mod tests { &table, vec![input], &ChunkShuffle::from_decimal_seed("0", 3).unwrap(), + None, ) .await .unwrap(); assert_eq!(chunks.len(), 3); assert!(chunks .iter() - .all(|chunk| chunk.exact_merged_row_count() == Some(3))); + .all(|chunk| chunk.row_count() == 3 && chunk.merged_row_count() == Some(3))); let mut ranges: Vec<_> = chunks .iter() diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 305f51eb9..fee3b4a68 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -277,8 +277,11 @@ impl DataFileReader { for split in splits { // Create DV factory for this split only. let dv_factory = reader.build_split_dv_factory(&split).await?; + let data_evolution = + crate::spec::CoreOptions::new(&reader.table_options).data_evolution_enabled(); + let mut split_file_offset = 0; - for (file_index, file_meta) in split.data_files().to_vec().into_iter().enumerate() { + for file_meta in split.data_files().to_vec() { let dv = DataFileReader::deletion_vector_for_file( dv_factory.as_ref(), &file_meta.file_name, @@ -310,24 +313,19 @@ impl DataFileReader { FileIndexResult::Remain }; - let global_ranges = split.row_ranges().map(|ranges| { + let range_base = if data_evolution { + file_meta.first_row_id.unwrap_or(0) + } else { + split_file_offset + }; + let split_ranges = split.row_ranges().map(|ranges| { to_local_row_ranges( ranges, - file_meta.first_row_id.unwrap_or(0), + range_base, file_meta.row_count, ) }); - let local_ranges = split - .file_row_range(file_index) - .map(|range| vec![range.clone()]); - let split_ranges = match (global_ranges, local_ranges) { - (Some(global), Some(local)) => { - Some(intersect_sorted_ranges(&global, &local)) - } - (Some(global), None) => Some(global), - (None, Some(local)) => Some(local), - (None, None) => None, - }; + split_file_offset += file_meta.row_count; let selected_ranges = match file_index_result { FileIndexResult::Remain => split_ranges, FileIndexResult::Skip => Some(Vec::new()), @@ -1035,7 +1033,9 @@ fn is_row_file(file_meta: &DataFileMeta) -> bool { .is_some_and(|path| path.to_ascii_lowercase().ends_with(".row")) } -/// Convert absolute RowRanges to normalized file-local 0-based ranges. +/// Convert ranges from their read-path coordinate system to file-local ranges. +/// `first_row_id` is a stable row ID for data evolution, or the file's +/// cumulative split-local physical offset for raw append reads. fn to_local_row_ranges( row_ranges: &[RowRange], first_row_id: i64, @@ -1702,6 +1702,19 @@ mod tests { use roaring::RoaringBitmap; use std::io; + #[test] + fn split_local_ranges_map_across_file_boundaries() { + let ranges = [RowRange::new(3, 4), RowRange::new(5, 7)]; + assert_eq!( + to_local_row_ranges(&ranges, 0, 5), + vec![RowRange::new(3, 4)] + ); + assert_eq!( + to_local_row_ranges(&ranges, 5, 4), + vec![RowRange::new(0, 2)] + ); + } + #[test] fn test_data_file_read_timing_aggregates_file_waits() { let timing = DataFileReadTiming::default(); diff --git a/crates/paimon/src/table/incremental_scan.rs b/crates/paimon/src/table/incremental_scan.rs index 8c2225f66..196964624 100644 --- a/crates/paimon/src/table/incremental_scan.rs +++ b/crates/paimon/src/table/incremental_scan.rs @@ -266,17 +266,20 @@ impl<'a> IncrementalScan<'a> { Ok(self) } - /// Select one balanced worker shard after chunk shuffling. - pub fn with_chunk_shuffle_shard(mut self, index: usize, count: usize) -> crate::Result { - self.scan = self.scan.with_chunk_shuffle_shard(index, count)?; + /// Select one balanced worker shard for a distributed scan. + pub fn with_shard(mut self, index: usize, count: usize) -> crate::Result { + self.scan = self.scan.with_shard(index, count)?; Ok(self) } pub async fn plan(&self) -> crate::Result { crate::spec::CoreOptions::new(self.table.schema().options()).ensure_read_authorized()?; - if self.scan.has_row_position_selection() || self.scan.has_chunk_shuffle() { + if self.scan.has_row_position_selection() + || self.scan.has_chunk_shuffle() + || self.scan.has_shard() + { return Err(crate::Error::Unsupported { - message: "Incremental row-position selection and chunk_shuffle require combined delta planning" + message: "Incremental row-position selection, chunk_shuffle and sharding require combined delta planning" .into(), }); } diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index b77dc182e..9f51cd0df 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -45,7 +45,9 @@ pub(crate) fn data_evolution_anchor_file(files: &[DataFileMeta]) -> crate::Resul } // ======================= RowRange =============================== -/// An inclusive row ID range `[from, to]` for filtering reads in data evolution mode. +/// An inclusive row range `[from, to]` in the coordinate system of the read +/// path: stable row IDs for data evolution, or physical positions for raw +/// primary-key and append reads. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct RowRange { from: i64, @@ -495,18 +497,10 @@ pub struct DataSplit { /// Deletion file for each data file, same order as `data_files`. /// `None` at index `i` means no deletion file for `data_files[i]` (matches Java getDeletionFiles() / List with null elements). data_deletion_files: Option]>>, + /// IndexedSplit-compatible ranges. Data-evolution reads interpret these as + /// stable row IDs; raw reads interpret them as split-local physical + /// positions over `data_files` in list order. row_ranges: Option>, - /// Optional file-local inclusive row range for each data file. `None` at - /// index `i` means the complete file. This is native-only planning - /// metadata used by chunk-shuffled append scans; Java's split wire format - /// cannot represent it. - #[serde(default)] - file_row_ranges: Option]>>, - /// Exact number of visible rows when planning had to inspect deletion - /// vectors. This is a planning hint, not part of the stable split wire - /// format. - #[serde(default)] - exact_merged_row_count: Option, /// Whether the split can be read raw, without the merge reader: its /// physical rows are exactly its logical rows (modulo deletion files). /// Mirrors Java `DataSplit#rawConvertible`. @@ -550,23 +544,6 @@ impl DataSplit { self.row_ranges.as_deref() } - /// File-local inclusive row ranges aligned with [`Self::data_files`]. - pub fn file_row_ranges(&self) -> Option<&[Option]> { - self.file_row_ranges.as_deref() - } - - /// File-local range for the data file at `index`; `None` means full file. - pub fn file_row_range(&self, index: usize) -> Option<&RowRange> { - self.file_row_ranges - .as_deref() - .and_then(|ranges| ranges.get(index)) - .and_then(Option::as_ref) - } - - pub fn exact_merged_row_count(&self) -> Option { - self.exact_merged_row_count - } - /// Whether this split can be read raw (no sort-merge needed); see the /// field doc. Mirrors Java `DataSplit#rawConvertible`. pub fn raw_convertible(&self) -> bool { @@ -615,6 +592,9 @@ impl DataSplit { /// nothing, so the result is a lower bound, not a total. Ask /// [`Self::row_counts_known`] before presenting it as one. pub fn row_count(&self) -> i64 { + if let Some(ranges) = &self.row_ranges { + return ranges.iter().map(RowRange::count).sum(); + } self.data_files .iter() .filter(|f| f.row_count_known()) @@ -647,8 +627,8 @@ impl DataSplit { /// /// Reference: [DataSplit.mergedRowCount()](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/table/source/DataSplit.java#L133) pub fn merged_row_count(&self) -> Option { - if let Some(count) = self.exact_merged_row_count { - return Some(count); + if let Some(ranges) = &self.row_ranges { + return Some(ranges.iter().map(RowRange::count).sum()); } if !self.row_counts_known() { return None; @@ -725,19 +705,6 @@ impl DataSplit { /// Byte-compatible with `compatibility/datasplit-v9`. Row ranges are not part of the /// format; `serialize_split_v1` wraps a row-range split as an `IndexedSplit` instead. pub fn serialize(&self) -> crate::Result> { - if self.file_row_ranges.is_some() { - return Err(crate::Error::Unsupported { - message: "Java DataSplit serialization cannot represent file-local row ranges" - .to_string(), - }); - } - self.serialize_metadata_view() - } - - /// Serialize the Java-compatible base split while intentionally omitting - /// native-only planning metadata. Callers must retain the original native - /// split for physical reading. - pub fn serialize_metadata_view(&self) -> crate::Result> { let mut out = Vec::new(); out.extend_from_slice(&SPLIT_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_VERSION.to_be_bytes()); @@ -893,31 +860,20 @@ impl DataSplit { /// `IndexedSplit` (type 3) wrapping the DataSplit body plus the ranges. Byte-compatible with /// `compatibility/split-v1-data` / `split-v1-indexed`. pub fn serialize_split_v1(&self) -> crate::Result> { - if self.file_row_ranges.is_some() { - return Err(crate::Error::Unsupported { - message: "SplitSerializer v1 cannot represent file-local row ranges".to_string(), - }); - } - self.serialize_split_v1_metadata_view() - } - - /// Serialize a cross-language metadata view of this split. Native-only - /// file ranges and exact row counts are intentionally omitted. - pub fn serialize_split_v1_metadata_view(&self) -> crate::Result> { let mut out = Vec::new(); out.extend_from_slice(&SPLIT_SER_MAGIC.to_be_bytes()); out.extend_from_slice(&SPLIT_SER_VERSION.to_be_bytes()); match &self.row_ranges { None => { out.extend_from_slice(&SPLIT_SER_TYPE_DATA_SPLIT.to_be_bytes()); - out.extend_from_slice(&self.serialize_metadata_view()?); + out.extend_from_slice(&self.serialize()?); } Some(ranges) => { out.extend_from_slice(&SPLIT_SER_TYPE_INDEXED_SPLIT.to_be_bytes()); // IndexedSplit#serialize: magic + version + DataSplit body + ranges + scores. out.extend_from_slice(&INDEXED_SPLIT_MAGIC.to_be_bytes()); out.extend_from_slice(&INDEXED_SPLIT_VERSION.to_be_bytes()); - out.extend_from_slice(&self.serialize_metadata_view()?); + out.extend_from_slice(&self.serialize()?); out.extend_from_slice(&(ranges.len() as i32).to_be_bytes()); for r in ranges.iter() { out.extend_from_slice(&r.from().to_be_bytes()); @@ -1220,8 +1176,6 @@ pub struct DataSplitBuilder { /// Same length as data_files; `None` at index i = no deletion file for data_files[i]. data_deletion_files: Option>>, row_ranges: Option>, - file_row_ranges: Option>>, - exact_merged_row_count: Option, raw_convertible: bool, is_streaming: bool, } @@ -1237,8 +1191,6 @@ impl DataSplitBuilder { data_files: None, data_deletion_files: None, row_ranges: None, - file_row_ranges: None, - exact_merged_row_count: None, // Splits with no merge semantics (append tables, single-file // utility splits) are raw by nature; the merge-tree and // data-evolution scan paths set this explicitly per split group. @@ -1286,18 +1238,6 @@ impl DataSplitBuilder { self } - /// Set file-local inclusive ranges aligned with `data_files`. `None` - /// selects the whole file at that position. - pub fn with_file_row_ranges(mut self, ranges: Vec>) -> Self { - self.file_row_ranges = Some(ranges); - self - } - - pub fn with_exact_merged_row_count(mut self, count: i64) -> Self { - self.exact_merged_row_count = Some(count); - self - } - /// Preserve the Java DataSplit event-reading contract. pub fn with_streaming(mut self, is_streaming: bool) -> Self { self.is_streaming = is_streaming; @@ -1353,40 +1293,6 @@ impl DataSplitBuilder { }); } } - if let Some(ranges) = &self.file_row_ranges { - if ranges.len() != data_files.len() { - return Err(crate::Error::DataInvalid { - message: format!( - "file_row_ranges length {} does not match data_files length {}", - ranges.len(), - data_files.len() - ), - source: None, - }); - } - for (file, range) in data_files.iter().zip(ranges) { - if let Some(range) = range { - if range.from() < 0 || range.to() >= file.row_count { - return Err(crate::Error::DataInvalid { - message: format!( - "file-local row range [{}, {}] is outside file '{}' row count {}", - range.from(), - range.to(), - file.file_name, - file.row_count - ), - source: None, - }); - } - } - } - } - if self.exact_merged_row_count.is_some_and(|count| count < 0) { - return Err(crate::Error::DataInvalid { - message: "exact_merged_row_count must be non-negative".to_string(), - source: None, - }); - } Ok(DataSplit { snapshot_id: self.snapshot_id, partition: Arc::new(partition), @@ -1396,8 +1302,6 @@ impl DataSplitBuilder { data_files: data_files.into(), data_deletion_files: self.data_deletion_files.map(Into::into), row_ranges: self.row_ranges.map(Into::into), - file_row_ranges: self.file_row_ranges.map(Into::into), - exact_merged_row_count: self.exact_merged_row_count, raw_convertible: self.raw_convertible, is_streaming: self.is_streaming, }) @@ -2337,45 +2241,6 @@ mod tests { ); } - #[test] - fn native_file_ranges_are_validated_preserved_and_never_silently_serialized() { - let split = v1_data_split_builder() - .with_file_row_ranges(vec![Some(RowRange::new(2, 5)), None]) - .with_exact_merged_row_count(15) - .build() - .unwrap(); - assert_eq!(split.file_row_range(0), Some(&RowRange::new(2, 5))); - assert_eq!(split.merged_row_count(), Some(15)); - assert!(matches!( - split.serialize_split_v1(), - Err(crate::Error::Unsupported { .. }) - )); - - let metadata = split.serialize_split_v1_metadata_view().unwrap(); - let decoded = DataSplit::deserialize_split_v1(&metadata).unwrap(); - assert!(decoded.file_row_ranges().is_none()); - assert_eq!(decoded.exact_merged_row_count(), None); - assert_eq!(decoded.data_files(), split.data_files()); - - let pickle = serde_json::to_vec(&split).unwrap(); - let restored: DataSplit = serde_json::from_slice(&pickle).unwrap(); - assert_eq!(restored, split); - - assert!(v1_data_split_builder() - .with_file_row_ranges(vec![Some(RowRange::new(0, 10))]) - .build() - .is_err()); - assert!(v1_data_split_builder() - .with_file_row_ranges(vec![Some(RowRange::new(0, 100)), None]) - .build() - .is_err()); - assert!(v1_data_split_builder() - .with_file_row_ranges(vec![Some(RowRange::new(0, 10)), None]) - .with_exact_merged_row_count(-1) - .build() - .is_err()); - } - #[test] fn serialize_indexed_split_v1_matches_golden() { // Row ranges -> IndexedSplit (type 3): same DataSplit as split-v1-data + ranges [1,4],[11,13]. @@ -2389,6 +2254,8 @@ mod tests { .with_row_ranges(vec![RowRange::new(1, 4), RowRange::new(11, 13)]) .build() .unwrap(); + assert_eq!(split.row_count(), 7); + assert_eq!(split.merged_row_count(), Some(7)); assert_eq!(split.serialize_split_v1().unwrap(), expected); } diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index 0b4753bf5..ed8ef6225 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1147,22 +1147,40 @@ impl<'a> TableScan<'a> { self.with_chunk_shuffle_config(config) } - /// Select one balanced worker shard after chunk shuffling. - pub fn with_chunk_shuffle_shard(mut self, index: usize, count: usize) -> crate::Result { + /// Select one balanced worker shard for a distributed scan. + /// + /// Sharding is scan-level state, independent of the selected planning + /// strategy, so callers may configure it before or after chunk shuffling. + pub fn with_shard(mut self, index: usize, count: usize) -> crate::Result { + if count == 0 || index >= count { + return Err(crate::Error::DataInvalid { + message: "shard count must be positive and index less than count".to_string(), + source: None, + }); + } match &mut self.0 { - TableScanKind::Paimon(scan) => match scan.split_selection.as_deref_mut() { - Some(ScanSplitSelection::ChunkShuffle(config)) => { - config.set_shard(index, count)?; - Ok(self) + TableScanKind::Paimon(scan) => { + if scan.row_position_selection().is_some() { + return Err(crate::Error::DataInvalid { + message: + "with_shard and row-position selection cannot be used simultaneously" + .to_string(), + source: None, + }); } - _ => Err(crate::Error::DataInvalid { - message: "with_chunk_shuffle_shard requires with_chunk_shuffle first" - .to_string(), - source: None, - }), - }, + match scan.split_selection.as_deref_mut() { + Some(selection) => selection.shard = Some((index, count)), + None => { + scan.split_selection = Some(Box::new(ScanSplitSelection { + mode: None, + shard: Some((index, count)), + })); + } + } + Ok(self) + } TableScanKind::Format(_) => Err(crate::Error::Unsupported { - message: "format tables do not support chunk_shuffle".to_string(), + message: "format tables do not support sharding".to_string(), }), } } @@ -1192,7 +1210,10 @@ impl<'a> TableScan<'a> { message: "chunk_shuffle only supports partition predicates".to_string(), }); } - scan.split_selection = Some(Box::new(ScanSplitSelection::ChunkShuffle(config))); + scan.split_selection = Some(Box::new(ScanSplitSelection { + mode: Some(ScanSplitMode::ChunkShuffle(config)), + shard: scan.shard(), + })); Ok(Self(TableScanKind::Paimon(scan))) } TableScanKind::Format(_) => Err(crate::Error::Unsupported { @@ -1212,6 +1233,10 @@ impl<'a> TableScan<'a> { matches!(&self.0, TableScanKind::Paimon(scan) if scan.chunk_shuffle().is_some()) } + pub(crate) fn has_shard(&self) -> bool { + matches!(&self.0, TableScanKind::Paimon(scan) if scan.shard().is_some()) + } + fn with_row_position_selection(self, selection: RowPositionSelection) -> crate::Result { match self.0 { TableScanKind::Paimon(mut scan) @@ -1225,6 +1250,14 @@ impl<'a> TableScan<'a> { source: None, }); } + if scan.shard().is_some() { + return Err(crate::Error::DataInvalid { + message: + "row-position selection and with_shard cannot be used simultaneously" + .into(), + source: None, + }); + } if scan .row_position_selection() .is_some_and(|previous| previous.is_slice() != selection.is_slice()) @@ -1235,7 +1268,10 @@ impl<'a> TableScan<'a> { source: None, }); } - scan.split_selection = Some(Box::new(ScanSplitSelection::RowPosition(selection))); + scan.split_selection = Some(Box::new(ScanSplitSelection { + mode: Some(ScanSplitMode::RowPosition(selection)), + shard: None, + })); Ok(Self(TableScanKind::Paimon(scan))) } _ => Err(crate::Error::Unsupported { @@ -1342,11 +1378,17 @@ impl<'a> TableScan<'a> { /// /// Reference: [pypaimon.read.table_scan.TableScan](https://github.com/apache/paimon/blob/master/paimon-python/pypaimon/read/table_scan.py) #[derive(Debug, Clone)] -enum ScanSplitSelection { +enum ScanSplitMode { RowPosition(RowPositionSelection), ChunkShuffle(ChunkShuffle), } +#[derive(Debug, Clone)] +struct ScanSplitSelection { + mode: Option, + shard: Option<(usize, usize)>, +} + #[derive(Debug, Clone)] struct PaimonTableScan<'a> { table: &'a Table, @@ -1374,19 +1416,33 @@ struct PaimonTableScan<'a> { impl<'a> PaimonTableScan<'a> { fn row_position_selection(&self) -> Option { - match self.split_selection.as_deref() { - Some(ScanSplitSelection::RowPosition(selection)) => Some(*selection), + match self + .split_selection + .as_deref() + .and_then(|selection| selection.mode.as_ref()) + { + Some(ScanSplitMode::RowPosition(selection)) => Some(*selection), _ => None, } } fn chunk_shuffle(&self) -> Option<&ChunkShuffle> { - match self.split_selection.as_deref() { - Some(ScanSplitSelection::ChunkShuffle(config)) => Some(config), + match self + .split_selection + .as_deref() + .and_then(|selection| selection.mode.as_ref()) + { + Some(ScanSplitMode::ChunkShuffle(config)) => Some(config), _ => None, } } + fn shard(&self) -> Option<(usize, usize)> { + self.split_selection + .as_deref() + .and_then(|selection| selection.shard) + } + fn is_streaming(&self) -> bool { self.incremental_split_mode.is_some() } @@ -1468,6 +1524,7 @@ impl<'a> PaimonTableScan<'a> { /// `scan.snapshot-id` / `scan.tag-name` handling. pub async fn plan(&self) -> crate::Result { self.ensure_query_auth_allowed()?; + self.validate_shard_strategy()?; let data_evolution_read_field_ids = self.projected_read_field_ids()?; let snapshot = match super::time_travel::resolve_snapshot(self.table).await? { Some(snapshot) => snapshot, @@ -1480,6 +1537,7 @@ impl<'a> PaimonTableScan<'a> { /// Plan the full scan and return metadata-pruning trace counters. pub async fn plan_with_trace(&self) -> crate::Result<(Plan, ScanTrace)> { self.ensure_query_auth_allowed()?; + self.validate_shard_strategy()?; let mut trace = ScanTrace { limit: self.limit, ..Default::default() @@ -1508,6 +1566,15 @@ impl<'a> PaimonTableScan<'a> { CoreOptions::new(self.table.schema().options()).ensure_read_authorized() } + fn validate_shard_strategy(&self) -> crate::Result<()> { + if self.shard().is_some() && self.chunk_shuffle().is_none() { + return Err(crate::Error::Unsupported { + message: "with_shard currently requires chunk_shuffle".to_string(), + }); + } + Ok(()) + } + fn projected_read_field_ids(&self) -> crate::Result>> { Ok(self.projected_read_field_ids.clone()) } @@ -1864,6 +1931,7 @@ impl<'a> PaimonTableScan<'a> { end_snapshot: &Snapshot, ) -> crate::Result { self.ensure_query_auth_allowed()?; + self.validate_shard_strategy()?; let data_evolution_read_field_ids = self.projected_read_field_ids()?; let mut scan = self.clone(); scan.incremental_split_mode = Some(IncrementalSplitMode::Batch); @@ -2658,7 +2726,7 @@ impl<'a> PaimonTableScan<'a> { (splits, split_candidates_built, false) }; let splits = if let Some(config) = self.chunk_shuffle() { - chunk_shuffle_splits(self.table, splits, config).await? + chunk_shuffle_splits(self.table, splits, config, self.shard()).await? } else { splits }; @@ -3791,19 +3859,38 @@ mod tests { .unwrap() .with_chunk_shuffle(0, 1) .is_err()); + assert!(reader + .new_scan() + .with_row_position_shard(0, 1) + .unwrap() + .with_shard(0, 1) + .is_err()); + assert!(reader + .new_scan() + .with_shard(0, 1) + .unwrap() + .with_row_position_shard(0, 1) + .is_err()); assert!(reader .new_scan() .with_chunk_shuffle(0, 1) .unwrap() .with_row_position_shard(0, 1) .is_err()); - assert!(reader.new_scan().with_chunk_shuffle_shard(0, 1).is_err()); + assert!(reader + .new_scan() + .with_shard(0, 1) + .unwrap() + .with_chunk_shuffle(0, 1) + .is_ok()); assert!(reader .new_scan() .with_chunk_shuffle(0, 1) .unwrap() - .with_chunk_shuffle_shard(0, 1) + .with_shard(0, 1) .is_ok()); + assert!(reader.new_scan().with_shard(0, 0).is_err()); + assert!(reader.new_scan().with_shard(1, 1).is_err()); } #[tokio::test] From 724d2d685b852519926e4fff181526c1efe7a421 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 21 Sep 2026 08:13:52 +0800 Subject: [PATCH 4/5] refactor: inline chunk shuffle scan configuration --- crates/paimon/src/table/table_scan.rs | 72 +++++++++++++-------------- 1 file changed, 34 insertions(+), 38 deletions(-) diff --git a/crates/paimon/src/table/table_scan.rs b/crates/paimon/src/table/table_scan.rs index ed8ef6225..46864b8e2 100644 --- a/crates/paimon/src/table/table_scan.rs +++ b/crates/paimon/src/table/table_scan.rs @@ -1144,7 +1144,40 @@ impl<'a> TableScan<'a> { /// bindings to preserve arbitrary-precision integer seeds. pub fn with_chunk_shuffle(self, seed: impl ToString, chunk_size: u64) -> crate::Result { let config = ChunkShuffle::from_decimal_seed(&seed.to_string(), chunk_size)?; - self.with_chunk_shuffle_config(config) + match self.0 { + TableScanKind::Paimon(mut scan) => { + if !scan.table.schema().primary_keys().is_empty() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle only supports append tables".to_string(), + }); + } + if scan.limit.is_some() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle cannot combine with limit".to_string(), + }); + } + if scan.row_ranges.is_some() || scan.row_position_selection().is_some() { + return Err(crate::Error::Unsupported { + message: + "chunk_shuffle cannot combine with row ranges or positional selection" + .to_string(), + }); + } + if !scan.data_predicates.is_empty() { + return Err(crate::Error::Unsupported { + message: "chunk_shuffle only supports partition predicates".to_string(), + }); + } + scan.split_selection = Some(Box::new(ScanSplitSelection { + mode: Some(ScanSplitMode::ChunkShuffle(config)), + shard: scan.shard(), + })); + Ok(Self(TableScanKind::Paimon(scan))) + } + TableScanKind::Format(_) => Err(crate::Error::Unsupported { + message: "format tables do not support chunk_shuffle".to_string(), + }), + } } /// Select one balanced worker shard for a distributed scan. @@ -1185,43 +1218,6 @@ impl<'a> TableScan<'a> { } } - fn with_chunk_shuffle_config(self, config: ChunkShuffle) -> crate::Result { - match self.0 { - TableScanKind::Paimon(mut scan) => { - if !scan.table.schema().primary_keys().is_empty() { - return Err(crate::Error::Unsupported { - message: "chunk_shuffle only supports append tables".to_string(), - }); - } - if scan.limit.is_some() { - return Err(crate::Error::Unsupported { - message: "chunk_shuffle cannot combine with limit".to_string(), - }); - } - if scan.row_ranges.is_some() || scan.row_position_selection().is_some() { - return Err(crate::Error::Unsupported { - message: - "chunk_shuffle cannot combine with row ranges or positional selection" - .to_string(), - }); - } - if !scan.data_predicates.is_empty() { - return Err(crate::Error::Unsupported { - message: "chunk_shuffle only supports partition predicates".to_string(), - }); - } - scan.split_selection = Some(Box::new(ScanSplitSelection { - mode: Some(ScanSplitMode::ChunkShuffle(config)), - shard: scan.shard(), - })); - Ok(Self(TableScanKind::Paimon(scan))) - } - TableScanKind::Format(_) => Err(crate::Error::Unsupported { - message: "format tables do not support chunk_shuffle".to_string(), - }), - } - } - pub(crate) fn has_row_position_selection(&self) -> bool { match &self.0 { TableScanKind::Paimon(scan) => scan.row_position_selection().is_some(), From 7596b12d00b531983ab932a7a2412f8ceca65e46 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Mon, 21 Sep 2026 09:01:46 +0800 Subject: [PATCH 5/5] fix: preserve global ranges for row-tracked tables --- bindings/python/tests/test_read.py | 37 ++++++++++ crates/paimon/src/table/chunk_shuffle.rs | 77 ++++++++++++++++++++- crates/paimon/src/table/data_file_reader.rs | 25 +++++-- crates/paimon/src/table/source.rs | 10 +-- 4 files changed, 136 insertions(+), 13 deletions(-) diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index 3f265ceaa..559b4f826 100644 --- a/bindings/python/tests/test_read.py +++ b/bindings/python/tests/test_read.py @@ -203,6 +203,43 @@ def test_with_row_ranges(): table.new_read_builder().with_row_ranges([(2, 1)]) +def test_row_tracking_append_row_ranges_keep_global_row_ids(): + with tempfile.TemporaryDirectory() as warehouse: + ctx = SQLContext() + ctx.register_catalog("paimon", {"warehouse": warehouse}) + ctx.sql("CREATE SCHEMA paimon.rdb") + ctx.sql("""CREATE TABLE paimon.rdb.tracked (id INT, pt STRING) + PARTITIONED BY (pt) WITH ('row-tracking.enabled' = 'true')""") + ctx.sql("""INSERT INTO paimon.rdb.tracked VALUES + (1, 'a'), (2, 'a'), (3, 'a')""") + ctx.sql("""INSERT INTO paimon.rdb.tracked VALUES + (4, 'b'), (5, 'b'), (6, 'b')""") + table = PaimonCatalog({"warehouse": warehouse}).get_table("rdb.tracked") + builder = table.new_read_builder().with_row_ranges([(3, 4)]) + + plan = builder.new_scan().plan() + rows = pa.Table.from_batches(builder.new_read().read(plan.splits())) + + assert rows.column("id").to_pylist() == [4, 5] + + chunk_builder = table.new_read_builder().with_projection(["id"]) + chunks = ( + chunk_builder.new_scan() + .with_chunk_shuffle("7", 2) + .plan() + .splits() + ) + chunk_rows = [ + pa.Table.from_batches(chunk_builder.new_read().read([split])) + .column("id").to_pylist() + for split in chunks + ] + assert all(0 < len(values) <= 2 for values in chunk_rows) + assert sorted(value for values in chunk_rows for value in values) == [ + 1, 2, 3, 4, 5, 6, + ] + + def test_format_table_rejects_row_ranges(): with tempfile.TemporaryDirectory() as warehouse: ctx = SQLContext() diff --git a/crates/paimon/src/table/chunk_shuffle.rs b/crates/paimon/src/table/chunk_shuffle.rs index f6a096636..d426f649d 100644 --- a/crates/paimon/src/table/chunk_shuffle.rs +++ b/crates/paimon/src/table/chunk_shuffle.rs @@ -242,6 +242,7 @@ async fn append_chunks( mut group: InputGroup, chunk_size: i64, ) -> crate::Result> { + let row_tracking = table.schema().core_options().row_tracking_enabled(); let mut chunks: Vec> = Vec::new(); let mut current = Vec::new(); let mut current_rows = 0; @@ -273,24 +274,40 @@ async fn append_chunks( chunks .into_iter() - .map(|segments| build_append_split(&group, segments)) + .map(|segments| build_append_split(&group, segments, row_tracking)) .collect() } fn build_append_split( group: &InputGroup, segments: Vec, + row_tracking: bool, ) -> crate::Result { let mut files = Vec::with_capacity(segments.len()); let mut deletion_files = Vec::with_capacity(segments.len()); let mut ranges = Vec::new(); let mut split_offset = 0; for segment in segments { + let range_base = if row_tracking { + segment + .input + .file + .first_row_id + .ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "Row-tracked file '{}' is missing first_row_id", + segment.input.file.file_name + ), + source: None, + })? + } else { + split_offset + }; ranges.extend( segment .ranges .into_iter() - .map(|range| RowRange::new(split_offset + range.from(), split_offset + range.to())), + .map(|range| RowRange::new(range_base + range.from(), range_base + range.to())), ); split_offset += segment.input.file.row_count; files.push(segment.input.file); @@ -758,6 +775,21 @@ mod tests { ) } + fn row_tracking_table() -> Table { + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .option("row-tracking.enabled", "true") + .build() + .unwrap(); + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "chunk_test"), + "memory:/chunk-test".to_string(), + TableSchema::new(0, &schema), + None, + ) + } + fn file(name: &str, row_count: i64, first_row_id: Option) -> DataFileMeta { DataFileMeta { file_name: name.to_string(), @@ -917,6 +949,47 @@ mod tests { assert_eq!([left, right].concat(), chunks); } + #[tokio::test] + async fn row_tracking_append_chunks_keep_global_row_ids() { + let table = row_tracking_table(); + let input = split( + vec![ + file("a.parquet", 2, Some(100)), + file("b.parquet", 2, Some(200)), + ], + true, + ); + let chunks = chunk_shuffle_splits( + &table, + vec![input], + &ChunkShuffle::from_decimal_seed("0", 3).unwrap(), + None, + ) + .await + .unwrap(); + + assert_eq!(chunks.len(), 2); + assert_eq!(chunks.iter().map(DataSplit::row_count).sum::(), 4); + let mut ranges: Vec<_> = chunks + .iter() + .flat_map(|chunk| chunk.row_ranges().unwrap()) + .map(|range| (range.from(), range.to())) + .collect(); + ranges.sort(); + assert_eq!(ranges, vec![(100, 101), (200, 200), (201, 201)]); + + let missing_row_id = split(vec![file("legacy.parquet", 2, None)], true); + let error = chunk_shuffle_splits( + &table, + vec![missing_row_id], + &ChunkShuffle::from_decimal_seed("0", 3).unwrap(), + None, + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("missing first_row_id")); + } + #[tokio::test] async fn data_evolution_chunks_keep_aligned_files_and_global_ranges() { let table = test_table(true); diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index fee3b4a68..cf429efc5 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -277,8 +277,9 @@ impl DataFileReader { for split in splits { // Create DV factory for this split only. let dv_factory = reader.build_split_dv_factory(&split).await?; - let data_evolution = - crate::spec::CoreOptions::new(&reader.table_options).data_evolution_enabled(); + let core_options = crate::spec::CoreOptions::new(&reader.table_options); + let ranges_use_row_ids = core_options.row_tracking_enabled() + || core_options.data_evolution_enabled(); let mut split_file_offset = 0; for file_meta in split.data_files().to_vec() { @@ -313,8 +314,14 @@ impl DataFileReader { FileIndexResult::Remain }; - let range_base = if data_evolution { - file_meta.first_row_id.unwrap_or(0) + let range_base = if ranges_use_row_ids { + file_meta.first_row_id.ok_or_else(|| crate::Error::DataInvalid { + message: format!( + "Row-tracked file '{}' is missing first_row_id", + file_meta.file_name + ), + source: None, + })? } else { split_file_offset }; @@ -1034,8 +1041,9 @@ fn is_row_file(file_meta: &DataFileMeta) -> bool { } /// Convert ranges from their read-path coordinate system to file-local ranges. -/// `first_row_id` is a stable row ID for data evolution, or the file's -/// cumulative split-local physical offset for raw append reads. +/// `first_row_id` is the coordinate base selected by the table's read path: +/// stable row ID for row-tracked tables, or cumulative split-local physical +/// offset for raw tables without row tracking. fn to_local_row_ranges( row_ranges: &[RowRange], first_row_id: i64, @@ -1713,6 +1721,11 @@ mod tests { to_local_row_ranges(&ranges, 5, 4), vec![RowRange::new(0, 2)] ); + + assert_eq!( + to_local_row_ranges(&[RowRange::new(103, 104)], 100, 6), + vec![RowRange::new(3, 4)] + ); } #[test] diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 9f51cd0df..e4c308665 100644 --- a/crates/paimon/src/table/source.rs +++ b/crates/paimon/src/table/source.rs @@ -46,8 +46,8 @@ pub(crate) fn data_evolution_anchor_file(files: &[DataFileMeta]) -> crate::Resul // ======================= RowRange =============================== /// An inclusive row range `[from, to]` in the coordinate system of the read -/// path: stable row IDs for data evolution, or physical positions for raw -/// primary-key and append reads. +/// path: stable row IDs for row-tracked tables, or physical positions for raw +/// tables without row tracking. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct RowRange { from: i64, @@ -497,9 +497,9 @@ pub struct DataSplit { /// Deletion file for each data file, same order as `data_files`. /// `None` at index `i` means no deletion file for `data_files[i]` (matches Java getDeletionFiles() / List with null elements). data_deletion_files: Option]>>, - /// IndexedSplit-compatible ranges. Data-evolution reads interpret these as - /// stable row IDs; raw reads interpret them as split-local physical - /// positions over `data_files` in list order. + /// IndexedSplit-compatible ranges. Row-tracked tables interpret these as + /// stable row IDs; raw tables without row tracking interpret them as + /// split-local physical positions over `data_files` in list order. row_ranges: Option>, /// Whether the split can be read raw, without the merge reader: its /// physical rows are exactly its logical rows (modulo deletion files).