diff --git a/bindings/python/src/read.rs b/bindings/python/src/read.rs index 22af9ee61..c6605a2b7 100644 --- a/bindings/python/src/read.rs +++ b/bindings/python/src/read.rs @@ -412,6 +412,8 @@ impl PyReadBuilder { incremental_range: None, row_position_slice: None, row_position_shard: None, + chunk_shuffle: None, + shard: None, } } @@ -449,6 +451,14 @@ pub struct PyTableScan { incremental_range: Option<(i64, i64)>, 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, } impl PyTableScan { @@ -464,6 +474,14 @@ impl PyTableScan { .with_row_position_shard(index, count) .map_err(to_py_err)?; } + if let Some(chunk_shuffle) = &self.chunk_shuffle { + scan = scan + .with_chunk_shuffle(&chunk_shuffle.seed, chunk_shuffle.chunk_size) + .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) } @@ -485,6 +503,14 @@ impl PyTableScan { .with_row_position_shard(index, count) .map_err(to_py_err)?; } + if let Some(chunk_shuffle) = &self.chunk_shuffle { + scan = scan + .with_chunk_shuffle(&chunk_shuffle.seed, chunk_shuffle.chunk_size) + .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) } @@ -533,6 +559,35 @@ impl PyTableScan { Ok(slf) } + /// 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, + ) -> PyResult> { + // Validate every combination immediately, not only when plan() runs. + slf.core_scan()? + .with_chunk_shuffle(&seed, chunk_size) + .map_err(to_py_err)?; + slf.chunk_shuffle = Some(PyChunkShuffle { seed, chunk_size }); + Ok(slf) + } + + /// 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_shard(index, count) + .map_err(to_py_err)?; + slf.shard = Some((index, count)); + Ok(slf) + } + fn plan(&self, py: Python<'_>) -> PyResult { py.detach(|| { runtime().block_on(async { @@ -756,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() } diff --git a/bindings/python/tests/test_read.py b/bindings/python/tests/test_read.py index c072f32ed..559b4f826 100644 --- a/bindings/python/tests/test_read.py +++ b/bindings/python/tests/test_read.py @@ -79,6 +79,100 @@ 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() + 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() + 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_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 + + # 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(): with tempfile.TemporaryDirectory() as warehouse: ctx = SQLContext() @@ -109,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() @@ -205,7 +336,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 new file mode 100644 index 000000000..d426f649d --- /dev/null +++ b/crates/paimon/src/table/chunk_shuffle.rs @@ -0,0 +1,1035 @@ +// 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::{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. +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct ChunkShuffle { + seed_words: Vec, + chunk_size: i64, +} + +impl ChunkShuffle { + /// Build from a Python integer's decimal spelling. Negative integers use + /// their absolute value, matching `random.Random`. + 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, + })?; + 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, + }) + } +} + +#[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, + ranges: Vec, +} + +#[derive(Debug)] +struct EvolutionSegment { + files: Vec, + ranges: Vec, +} + +/// 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, + 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()) { + 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)) = 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 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; + + 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; + }; + current_rows += slice.live_rows; + current.push(AppendSegment { + input: input.clone(), + ranges: slice.ranges, + }); + } + } + if !current.is_empty() { + chunks.push(current); + } + + chunks + .into_iter() + .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(range_base + range.from(), range_base + 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_row_ranges(merge_row_ranges(ranges)); + if deletion_files.iter().any(Option::is_some) { + builder = builder.with_data_deletion_files(deletion_files); + } + 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(), + ranges: slice + .ranges + .into_iter() + .map(|range| { + RowRange::new(first_row_id + range.from(), first_row_id + range.to()) + }) + .collect(), + }); + } + } + 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 mut files = Vec::new(); + let mut deletion_files = Vec::new(); + let mut ranges = Vec::new(); + for segment in segments { + ranges.extend(segment.ranges); + for input in segment.files { + files.push(input.file); + deletion_files.push(input.deletion_file); + } + } + 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); + } + 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 { + ranges: Vec, + 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 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); + 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); + 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 { ranges, live_rows })); + } + self.skip_deleted_at_cursor(); + } + if live_rows == 0 { + Ok(None) + } else { + Ok(Some(PhysicalSlice { ranges, 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 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(), + 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_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.ranges, first.live_rows), + (vec![RowRange::new(1, 2), RowRange::new(5, 5)], 3) + ); + let second = slicer.take(3).unwrap().unwrap(); + 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.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] + 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, None) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert!(chunks + .iter() + .all(|chunk| chunk.row_count() == 3 && chunk.merged_row_count() == Some(3))); + + let mut covered = Vec::new(); + for chunk in &chunks { + 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!( + 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, 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); + } + + #[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); + 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(), + None, + ) + .await + .unwrap(); + assert_eq!(chunks.len(), 3); + assert!(chunks + .iter() + .all(|chunk| chunk.row_count() == 3 && chunk.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..cf429efc5 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -277,6 +277,10 @@ impl DataFileReader { for split in splits { // Create DV factory for this split only. let dv_factory = reader.build_split_dv_factory(&split).await?; + 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() { let dv = DataFileReader::deletion_vector_for_file( @@ -310,13 +314,25 @@ impl DataFileReader { FileIndexResult::Remain }; + 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 + }; 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, ) }); + split_file_offset += file_meta.row_count; let selected_ranges = match file_index_result { FileIndexResult::Remain => split_ranges, FileIndexResult::Skip => Some(Vec::new()), @@ -1024,7 +1040,10 @@ 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 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, @@ -1691,6 +1710,24 @@ 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)] + ); + + assert_eq!( + to_local_row_ranges(&[RowRange::new(103, 104)], 100, 6), + vec![RowRange::new(3, 4)] + ); + } + #[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 c4e640409..196964624 100644 --- a/crates/paimon/src/table/incremental_scan.rs +++ b/crates/paimon/src/table/incremental_scan.rs @@ -256,11 +256,30 @@ impl<'a> IncrementalScan<'a> { Ok(self) } + /// Repack the combined APPEND-delta batch into deterministic chunks. + 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 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() { + 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 requires combined delta planning" + message: "Incremental row-position selection, chunk_shuffle and sharding require combined delta planning" .into(), }); } diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 407f9e0b2..74581c019 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; diff --git a/crates/paimon/src/table/source.rs b/crates/paimon/src/table/source.rs index 4499fd011..e4c308665 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 row-tracked tables, or physical positions for raw +/// tables without row tracking. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct RowRange { from: i64, @@ -495,6 +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. 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). @@ -587,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()) @@ -619,6 +627,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(ranges) = &self.row_ranges { + return Some(ranges.iter().map(RowRange::count).sum()); + } if !self.row_counts_known() { return None; } @@ -2243,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 bfc9f2b1c..46864b8e2 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,125 @@ impl<'a> TableScan<'a> { self.with_row_position_selection(RowPositionSelection::shard(index, count)?) } + /// Repack an append scan into deterministic fixed-live-row chunks. + /// + /// `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)?; + 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. + /// + /// 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) => { + 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, + }); + } + 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 sharding".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()) + } + + 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) 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.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 + .row_position_selection() .is_some_and(|previous| previous.is_slice() != selection.is_slice()) { return Err(crate::Error::DataInvalid { @@ -1158,7 +1264,10 @@ impl<'a> TableScan<'a> { source: None, }); } - scan.row_position_selection = Some(selection); + scan.split_selection = Some(Box::new(ScanSplitSelection { + mode: Some(ScanSplitMode::RowPosition(selection)), + shard: None, + })); Ok(Self(TableScanKind::Paimon(scan))) } _ => Err(crate::Error::Unsupported { @@ -1264,6 +1373,18 @@ 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 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, @@ -1274,7 +1395,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 +1411,34 @@ struct PaimonTableScan<'a> { } impl<'a> PaimonTableScan<'a> { + fn row_position_selection(&self) -> Option { + 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() + .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() } @@ -1306,7 +1458,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 +1487,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 } @@ -1368,6 +1520,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, @@ -1380,6 +1533,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() @@ -1408,6 +1562,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()) } @@ -1764,6 +1927,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); @@ -1826,7 +1990,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 +2285,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 +2341,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 +2559,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 +2721,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, self.shard()).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(); @@ -3680,6 +3849,44 @@ 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(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_shard(0, 1) + .unwrap() + .with_chunk_shuffle(0, 1) + .is_ok()); + assert!(reader + .new_scan() + .with_chunk_shuffle(0, 1) + .unwrap() + .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]