This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 8aaed4e2 feat: plan source-backed primary-key sorted indexes (#1086)
8aaed4e2 is described below
commit 8aaed4e20823d1b9434a8ca3f0f715466f52f035
Author: Jingsong Lee <[email protected]>
AuthorDate: Fri Oct 9 19:54:59 2026 +0800
feat: plan source-backed primary-key sorted indexes (#1086)
---
crates/paimon/src/table/data_file_reader.rs | 26 +-
crates/paimon/src/table/global_index_scanner.rs | 15 +-
.../src/table/global_index_scanner/row_ranges.rs | 2 +-
crates/paimon/src/table/kv_file_reader.rs | 6 +-
crates/paimon/src/table/mod.rs | 1 +
crates/paimon/src/table/pk_sorted_index_scan.rs | 511 ++++++++++++++
.../paimon/src/table/pk_sorted_index_scan/tests.rs | 743 +++++++++++++++++++++
crates/paimon/src/table/source.rs | 27 +-
crates/paimon/src/table/table_scan.rs | 101 ++-
9 files changed, 1406 insertions(+), 26 deletions(-)
diff --git a/crates/paimon/src/table/data_file_reader.rs
b/crates/paimon/src/table/data_file_reader.rs
index 4c32b659..1b4ac987 100644
--- a/crates/paimon/src/table/data_file_reader.rs
+++ b/crates/paimon/src/table/data_file_reader.rs
@@ -486,34 +486,30 @@ impl DataFileReader {
)
}
- /// Read one file with the complete physical data schema supplied by a
- /// caller such as the KV reader. The format chooses whether it needs the
- /// full schema or only the projected fields.
- pub(super) fn read_single_file_stream_with_schema(
+ /// Read one file with the complete physical data schema and optional
+ /// file-local row positions, independent of stable row IDs. The format
+ /// chooses whether it needs the full schema or only the projected fields.
+ pub(super) fn read_single_file_stream_with_schema_and_local_ranges(
&self,
split: &DataSplit,
file_meta: DataFileMeta,
data_fields: Option<Vec<DataField>>,
data_schema_fields: Vec<DataField>,
dv: Option<Arc<DeletionVector>>,
- row_ranges: Option<Vec<RowRange>>,
+ local_row_ranges: Option<Vec<RowRange>>,
) -> crate::Result<ArrowRecordBatchStream> {
- let local_ranges = row_ranges.as_ref().map(|ranges| {
- to_local_row_ranges(
- ranges,
- file_meta.first_row_id.unwrap_or(0),
- file_meta.row_count,
- )
- });
- let row_selection =
- merge_row_selection(file_meta.row_count, dv.as_deref(),
local_ranges.as_deref());
+ let row_selection = merge_row_selection(
+ file_meta.row_count,
+ dv.as_deref(),
+ local_row_ranges.as_deref(),
+ );
self.read_single_file_stream_with_selection(
split,
file_meta,
data_fields,
Some(data_schema_fields),
row_selection,
- local_ranges,
+ local_row_ranges,
)
}
diff --git a/crates/paimon/src/table/global_index_scanner.rs
b/crates/paimon/src/table/global_index_scanner.rs
index 008b323a..b9f7f058 100644
--- a/crates/paimon/src/table/global_index_scanner.rs
+++ b/crates/paimon/src/table/global_index_scanner.rs
@@ -32,6 +32,7 @@ mod row_ranges;
pub(crate) use deletion_vectors::deleted_row_ranges_for_data_evolution_dvs;
use entry::{validate_fm_file_sets, GlobalIndexEntry, GlobalIndexEntryMeta,
GlobalIndexFileKind};
+pub(super) use row_ranges::intersect_sorted_ranges;
use row_ranges::unindexed_ranges_for_indexed_coverage;
pub(crate) use row_ranges::{
search_limit_with_deleted_rows, unindexed_ranges_for_global_index_entries,
RowRangeIndex,
@@ -168,7 +169,7 @@ impl GlobalIndexScanner {
}
#[allow(clippy::too_many_arguments)]
- fn create_with_fm_options(
+ pub(super) fn create_with_fm_options(
file_io: &FileIO,
table_path: &str,
global_index_thread_num: usize,
@@ -364,6 +365,18 @@ impl GlobalIndexScanner {
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))),
)
}
+
+ /// Matching positions without data-evolution's unindexed row-id extension.
+ /// Source-backed PK indexes cover a validated immutable source group.
+ pub(super) async fn matching_ranges(
+ &self,
+ predicate: &Predicate,
+ ) -> Result<Option<Vec<RowRange>>> {
+ Ok(self
+ .evaluate(predicate)
+ .await?
+ .map(|result| result.row_ranges))
+ }
}
/// Create a GlobalIndexScanner and evaluate predicates, returning row ranges.
diff --git a/crates/paimon/src/table/global_index_scanner/row_ranges.rs
b/crates/paimon/src/table/global_index_scanner/row_ranges.rs
index a306d88a..3c4c7f12 100644
--- a/crates/paimon/src/table/global_index_scanner/row_ranges.rs
+++ b/crates/paimon/src/table/global_index_scanner/row_ranges.rs
@@ -46,7 +46,7 @@ pub(super) fn bitmap_to_ranges(bitmap: &RoaringTreemap) ->
Vec<RowRange> {
}
/// Intersect two sorted range lists using RowRangeIndex for efficient binary
search.
-pub(super) fn intersect_sorted_ranges(a: &[RowRange], b: &[RowRange]) ->
Vec<RowRange> {
+pub(crate) fn intersect_sorted_ranges(a: &[RowRange], b: &[RowRange]) ->
Vec<RowRange> {
let idx = RowRangeIndex::from_sorted_ranges(a.to_vec());
let mut result = Vec::new();
for r in b {
diff --git a/crates/paimon/src/table/kv_file_reader.rs
b/crates/paimon/src/table/kv_file_reader.rs
index cbe1aa6c..a9d9cdf6 100644
--- a/crates/paimon/src/table/kv_file_reader.rs
+++ b/crates/paimon/src/table/kv_file_reader.rs
@@ -717,7 +717,7 @@ impl KeyValueFileReader {
// Only the PK-only predicate projection may
run before
// sort-merge. A value index can match an old
version of
// a key, so it must not prune merge inputs
here.
- let row_ranges = if file_index_read_enabled {
+ let local_row_ranges = if
file_index_read_enabled {
match evaluate_file_index(
&run_file_io,
split.bucket_path(),
@@ -750,13 +750,13 @@ impl KeyValueFileReader {
)),
None => None,
};
- let mut file_stream =
reader.read_single_file_stream_with_schema(
+ let mut file_stream =
reader.read_single_file_stream_with_schema_and_local_ranges(
split.as_ref(),
file_meta,
data_fields,
data_schema_fields,
deletion_vector,
- row_ranges,
+ local_row_ranges,
)?;
while let Some(batch) =
file_stream.next().await {
yield batch?;
diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs
index 443f4848..5c577881 100644
--- a/crates/paimon/src/table/mod.rs
+++ b/crates/paimon/src/table/mod.rs
@@ -108,6 +108,7 @@ mod pk_full_text_read;
mod pk_full_text_scan;
mod pk_search_position;
mod pk_search_ranker;
+mod pk_sorted_index_scan;
mod pk_vector_bucket_split;
mod pk_vector_data_file_reader;
mod pk_vector_indexed_split_read;
diff --git a/crates/paimon/src/table/pk_sorted_index_scan.rs
b/crates/paimon/src/table/pk_sorted_index_scan.rs
new file mode 100644
index 00000000..75a17ab2
--- /dev/null
+++ b/crates/paimon/src/table/pk_sorted_index_scan.rs
@@ -0,0 +1,511 @@
+// 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.
+
+//! Source-backed BTREE/BITMAP planning for batch primary-key reads.
+//! Mirrors Java PrimaryKeySortedIndexScan and PkSortedBucketIndexState:
+//! immutable compacted files may use physical positions; merge inputs stay
whole.
+
+use super::global_index_scanner::{intersect_sorted_ranges, GlobalIndexScanner};
+use super::index_file_path::IndexFileLocation;
+use super::{merge_row_ranges, DataSplit, RowRange, Table};
+use crate::spec::{
+ should_read_pk_index_source, CoreOptions, DataField, DataFileMeta,
FileKind,
+ IndexManifestEntry, Predicate, PrimaryKeyIndexSourceMeta,
+};
+use crate::{Error, Result};
+use std::collections::{BTreeMap, HashMap, HashSet};
+use std::future::Future;
+use std::pin::Pin;
+
+const MAX_POSITION_RANGES: usize = 4096;
+
+type Options = HashMap<String, String>;
+type SourceKey = (Vec<u8>, i32, String, i32);
+type SourceGroups = HashMap<SourceKey, (usize, RowRange)>;
+type EvaluateFuture<'a> = Pin<Box<dyn Future<Output =
Result<Option<Vec<RowRange>>>> + Send + 'a>>;
+
+pub(super) struct Definition {
+ field_id: i32,
+ index_type: &'static str,
+ options: Options,
+}
+
+pub(super) fn definitions(fields: &[DataField], options: &Options) ->
Result<Vec<Definition>> {
+ let mut owners = HashMap::new();
+ for family in ["btree", "bitmap", "vector", "full-text", "multivalue",
"fm"] {
+ let key = format!("pk-{family}.index.columns");
+ if let Some(columns) = options.get(&key) {
+ let mut unique = HashSet::new();
+ for column in columns.split(',').map(str::trim) {
+ if !unique.insert(column) {
+ return Err(config_invalid(format!(
+ "{key} contains duplicate column '{column}'"
+ )));
+ }
+ if owners.insert(column, family).is_some() {
+ return Err(config_invalid(format!(
+ "Column '{column}' can own at most one primary-key
index"
+ )));
+ }
+ }
+ }
+ }
+ let mut result = Vec::new();
+ for field in fields {
+ let Some(&family @ ("btree" | "bitmap")) = owners.get(field.name())
else {
+ continue;
+ };
+ result.push(Definition {
+ field_id: field.id(),
+ index_type: family,
+ options: definition_options(options, field.name(), family)?,
+ });
+ }
+ Ok(result)
+}
+
+fn config_invalid(message: String) -> Error {
+ Error::ConfigInvalid { message }
+}
+
+fn definition_options(options: &Options, column: &str, family: &str) ->
Result<Options> {
+ let mut resolved = options.clone();
+ resolved.remove("sorted-index.records-per-file");
+ resolved.remove("sorted-index.records-per-range");
+ let option_key = format!("fields.{column}.pk-{family}.index.options");
+ let Some(raw) = options
+ .get(&option_key)
+ .filter(|raw| !raw.trim().is_empty())
+ else {
+ return Ok(resolved);
+ };
+ let parsed: serde_json::Map<String, serde_json::Value> =
+ serde_json::from_str(raw).map_err(|_| {
+ config_invalid(format!(
+ "{option_key} must be a JSON object of option key-value pairs"
+ ))
+ })?;
+ let prefix = format!("{family}-index.");
+ for (key, value) in parsed {
+ if key.trim().is_empty() {
+ return Err(config_invalid(format!(
+ "{option_key} contains an empty option key"
+ )));
+ }
+ let value = match value {
+ serde_json::Value::String(value) => value,
+ serde_json::Value::Number(_) | serde_json::Value::Bool(_) =>
value.to_string(),
+ _ => {
+ return Err(config_invalid(format!(
+ "{option_key} contains an invalid option value for {key}"
+ )))
+ }
+ };
+ let qualified = if key.starts_with(&prefix) ||
key.starts_with("fields.") {
+ key
+ } else {
+ format!("{prefix}{key}")
+ };
+ if resolved
+ .get(&qualified)
+ .is_some_and(|previous| previous != &value)
+ {
+ return Err(config_invalid(format!(
+ "{option_key} defines conflicting values for {qualified}"
+ )));
+ }
+ resolved.insert(qualified, value);
+ }
+ Ok(resolved)
+}
+
+struct Group {
+ entry: IndexManifestEntry,
+ definition: usize,
+ row_count: i64,
+ scanner: Option<GlobalIndexScanner>,
+ opened: bool,
+ // Shared across all active source files: each query reads the payload
once.
+ queries: Vec<(Predicate, Option<Vec<RowRange>>)>,
+}
+
+struct Candidate {
+ entry: IndexManifestEntry,
+ row_count: i64,
+ active_sources: Vec<(String, RowRange)>,
+}
+
+fn candidate(entry: &IndexManifestEntry, active: &[&DataFileMeta]) ->
Option<(i32, Candidate)> {
+ let meta = entry.index_file.global_index_meta.as_ref()?;
+ if meta
+ .extra_field_ids
+ .as_ref()
+ .is_some_and(|ids| !ids.is_empty())
+ {
+ return None;
+ }
+ let source = PrimaryKeyIndexSourceMeta::from_global_index_meta(meta).ok()?;
+ let active: HashMap<_, _> = active
+ .iter()
+ .filter(|file| file.level == source.data_level() &&
should_read_pk_index_source(file))
+ .map(|file| (file.file_name.as_str(), file.row_count))
+ .collect();
+ let mut previous: Option<&str> = None;
+ let mut offset = 0i64;
+ let mut active_sources = Vec::new();
+ for file in source.source_files() {
+ if previous.is_some_and(|name| name >= file.file_name()) {
+ return None;
+ }
+ previous = Some(file.file_name());
+ let end = offset.checked_add(file.row_count())?;
+ if let Some(&rows) = active.get(file.file_name()) {
+ if rows != file.row_count() {
+ return None;
+ }
+ if end > offset {
+ active_sources.push((file.file_name().to_string(),
RowRange::new(offset, end - 1)));
+ }
+ }
+ offset = end;
+ }
+ if active_sources.is_empty()
+ || offset <= 0
+ || meta.row_range_start != 0
+ || meta.row_range_end != offset - 1
+ || entry.index_file.row_count != offset
+ {
+ return None;
+ }
+ Some((
+ source.data_level(),
+ Candidate {
+ entry: entry.clone(),
+ row_count: offset,
+ active_sources,
+ },
+ ))
+}
+
+fn plan_groups(
+ table: &Table,
+ splits: &[DataSplit],
+ entries: &[IndexManifestEntry],
+ definitions: &[Definition],
+) -> (Vec<Group>, SourceGroups) {
+ let mut buckets: BTreeMap<_, (Vec<&DataFileMeta>, &str)> = BTreeMap::new();
+ for split in splits {
+ let (files, _) = buckets
+ .entry((split.partition().to_serialized_bytes(), split.bucket()))
+ .or_insert_with(|| (Vec::new(), split.bucket_path()));
+ files.extend(split.data_files());
+ }
+ let mut entries_by_bucket: HashMap<_, Vec<_>> = HashMap::new();
+ for entry in entries.iter().filter(|entry| entry.kind == FileKind::Add) {
+ entries_by_bucket
+ .entry((entry.partition.as_slice(), entry.bucket))
+ .or_default()
+ .push(entry);
+ }
+ let mut groups = Vec::new();
+ let mut sources = HashMap::new();
+ for ((partition, bucket), (files, bucket_path)) in buckets {
+ let location = IndexFileLocation::BucketLocal {
+ table_path: table.location().trim_end_matches('/'),
+ bucket_path,
+ index_file_in_data_file_dir: table
+ .schema()
+ .core_options()
+ .index_file_in_data_file_dir(),
+ };
+ for (definition_index, definition) in definitions.iter().enumerate() {
+ let mut levels: BTreeMap<i32, Vec<Candidate>> = BTreeMap::new();
+ for entry in entries_by_bucket
+ .get(&(partition.as_slice(), bucket))
+ .into_iter()
+ .flatten()
+ {
+ if entry.index_file.index_type != definition.index_type
+ || entry
+ .index_file
+ .global_index_meta
+ .as_ref()
+ .is_none_or(|meta| meta.index_field_id !=
definition.field_id)
+ {
+ continue;
+ }
+ if let Some((level, candidate)) = candidate(entry, &files) {
+ levels.entry(level).or_default().push(candidate);
+ }
+ }
+ for candidates in levels.into_values() {
+ // Java rejects ambiguous payload coverage at the same level.
+ if candidates.len() != 1 {
+ continue;
+ }
+ let Candidate {
+ mut entry,
+ row_count,
+ active_sources,
+ } = candidates.into_iter().next().unwrap();
+ entry.index_file.external_path = Some(location.resolve(
+ &entry.index_file.file_name,
+ entry.index_file.external_path.as_deref(),
+ ));
+ let group_id = groups.len();
+ for (file, range) in active_sources {
+ sources.insert(
+ (partition.clone(), bucket, file, definition.field_id),
+ (group_id, range),
+ );
+ }
+ groups.push(Group {
+ entry,
+ definition: definition_index,
+ row_count,
+ scanner: None,
+ opened: false,
+ queries: Vec::new(),
+ });
+ }
+ }
+ }
+ (groups, sources)
+}
+
+impl Group {
+ async fn query(
+ &mut self,
+ table: &Table,
+ definition: &Definition,
+ predicate: &Predicate,
+ ) -> Result<Option<Vec<RowRange>>> {
+ if let Some((_, result)) = self.queries.iter().find(|(query, _)| query
== predicate) {
+ return Ok(result.clone());
+ }
+ let result = self.evaluate(table, definition, predicate).await;
+ let ranges = match result {
+ Ok(Some(ranges))
+ if ranges
+ .iter()
+ .all(|range| range.from() >= 0 && range.to() <
self.row_count) =>
+ {
+ Some(ranges)
+ }
+ Ok(_) => None,
+ Err(error) => {
+ log::warn!(
+ "Ignoring source-backed PK index '{}': {error}",
+ self.entry.index_file.file_name
+ );
+ None
+ }
+ };
+ self.queries.push((predicate.clone(), ranges.clone()));
+ Ok(ranges)
+ }
+
+ async fn evaluate(
+ &mut self,
+ table: &Table,
+ definition: &Definition,
+ predicate: &Predicate,
+ ) -> Result<Option<Vec<RowRange>>> {
+ if !self.opened {
+ self.opened = true;
+ let options = CoreOptions::new(&definition.options);
+ self.scanner = GlobalIndexScanner::create_with_fm_options(
+ table.file_io(),
+ table.location(),
+ options.global_index_thread_num()?,
+ options.btree_index_fallback_scan_max_size()?,
+ options.btree_index_data_block_cache_size()?,
+ options.bitmap_index_fallback_scan_max_size()?,
+ std::slice::from_ref(&self.entry),
+ table.schema().fields(),
+ crate::fm_index::FMReadOptions::default(),
+ )?;
+ }
+ match &self.scanner {
+ Some(scanner) => scanner.matching_ranges(predicate).await,
+ None => Ok(None),
+ }
+ }
+}
+
+struct Evaluation<'a> {
+ table: &'a Table,
+ definitions: &'a [Definition],
+ groups: Vec<Group>,
+ sources: SourceGroups,
+}
+
+impl Evaluation<'_> {
+ fn evaluate<'a>(
+ &'a mut self,
+ split: &'a DataSplit,
+ file: &'a DataFileMeta,
+ predicate: &'a Predicate,
+ ) -> EvaluateFuture<'a> {
+ Box::pin(async move {
+ match predicate {
+ Predicate::Leaf { column, .. } => {
+ self.evaluate_field(split, file, column, predicate).await
+ }
+ Predicate::And(_) => {
+ let mut result: Option<Vec<RowRange>> = None;
+ let mut fields: BTreeMap<String, Vec<Predicate>> =
BTreeMap::new();
+ let mut other = Vec::new();
+ for child in predicate.clone().split_and() {
+ if let Predicate::Leaf { column, .. } = &child {
+
fields.entry(column.clone()).or_default().push(child);
+ } else {
+ other.push(child);
+ }
+ }
+ // Same-field bounds must reach the scalar scanner
together:
+ // it combines BETWEEN and reuses a single opened reader.
+ for (column, predicates) in fields {
+ let combined = Predicate::and(predicates);
+ if let Some(ranges) =
+ self.evaluate_field(split, file, &column,
&combined).await?
+ {
+ result = Some(match result {
+ None => ranges,
+ Some(previous) =>
intersect_sorted_ranges(&previous, &ranges),
+ });
+ }
+ }
+ for child in &other {
+ if let Some(ranges) = self.evaluate(split, file,
child).await? {
+ result = Some(match result {
+ None => ranges,
+ Some(previous) =>
intersect_sorted_ranges(&previous, &ranges),
+ });
+ }
+ }
+ Ok(result)
+ }
+ Predicate::Or(children) if !children.is_empty() => {
+ let mut result = Vec::new();
+ for child in children {
+ let Some(ranges) = self.evaluate(split, file,
child).await? else {
+ return Ok(None);
+ };
+ result.extend(ranges);
+ }
+ Ok(Some(merge_row_ranges(result)))
+ }
+ _ => Ok(None),
+ }
+ })
+ }
+ async fn evaluate_field(
+ &mut self,
+ split: &DataSplit,
+ file: &DataFileMeta,
+ column: &str,
+ predicate: &Predicate,
+ ) -> Result<Option<Vec<RowRange>>> {
+ let Some(field) = self
+ .table
+ .schema()
+ .fields()
+ .iter()
+ .find(|field| field.name() == column)
+ else {
+ return Ok(None);
+ };
+ let key = (
+ split.partition().to_serialized_bytes(),
+ split.bucket(),
+ file.file_name.clone(),
+ field.id(),
+ );
+ let Some((group_id, range)) = self.sources.get(&key).cloned() else {
+ return Ok(None);
+ };
+ let group = &mut self.groups[group_id];
+ let definition = &self.definitions[group.definition];
+ Ok(group
+ .query(self.table, definition, predicate)
+ .await?
+ .map(|ranges| localize(&ranges, range)))
+ }
+}
+
+fn localize(ranges: &[RowRange], source: RowRange) -> Vec<RowRange> {
+ let start = ranges.partition_point(|range| range.to() < source.from());
+ ranges[start..]
+ .iter()
+ .take_while(|range| range.from() <= source.to())
+ .filter_map(|range| {
+ let from = range.from().max(source.from());
+ let to = range.to().min(source.to());
+ (from <= to).then(|| RowRange::new(from - source.from(), to -
source.from()))
+ })
+ .collect()
+}
+
+pub(super) async fn refine(
+ table: &Table,
+ splits: Vec<DataSplit>,
+ entries: &[IndexManifestEntry],
+ definitions: &[Definition],
+ predicates: &[Predicate],
+) -> Result<Vec<DataSplit>> {
+ let (groups, sources) = plan_groups(table, &splits, entries, definitions);
+ if groups.is_empty() {
+ return Ok(splits);
+ }
+ let mut evaluation = Evaluation {
+ table,
+ definitions,
+ groups,
+ sources,
+ };
+ let predicate = Predicate::and(predicates.to_vec());
+ let mut result = Vec::new();
+ for split in splits {
+ if !split.raw_convertible() || split.is_streaming() ||
split.row_ranges().is_some() {
+ result.push(split);
+ continue;
+ }
+ for (file_index, file) in split.data_files().iter().enumerate() {
+ let ranges = evaluation.evaluate(&split, file, &predicate).await?;
+ if ranges.as_ref().is_some_and(Vec::is_empty) {
+ continue;
+ }
+ let mut selected = split.for_pk_index_file(file_index);
+ if let Some(ranges) = ranges.filter(|ranges| {
+ ranges.len() <= MAX_POSITION_RANGES
+ && ranges.iter().all(|range| {
+ range.from() >= 0
+ && range.to() < file.row_count
+ && range.to() <= i64::from(i32::MAX)
+ })
+ }) {
+ selected = selected.with_selected_row_ranges(ranges);
+ }
+ result.push(selected);
+ }
+ }
+ Ok(result)
+}
+
+#[cfg(test)]
+mod tests;
diff --git a/crates/paimon/src/table/pk_sorted_index_scan/tests.rs
b/crates/paimon/src/table/pk_sorted_index_scan/tests.rs
new file mode 100644
index 00000000..be7ffb65
--- /dev/null
+++ b/crates/paimon/src/table/pk_sorted_index_scan/tests.rs
@@ -0,0 +1,743 @@
+// 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.
+
+use super::*;
+use crate::btree::{
+ make_key_comparator, serialize_datum, test_util::VecFileWrite,
BTreeIndexWriter,
+ BlockCompressionType,
+};
+use crate::catalog::Identifier;
+use crate::io::FileIOBuilder;
+use crate::spec::{
+ stats::BinaryTableStats, BinaryRow, DataType, Datum, GlobalIndexMeta,
IndexFileMeta, IntType,
+ PredicateBuilder, Schema, TableSchema,
+};
+use crate::table::bitmap_global_index_format::{
+ make_bitmap_key_comparator, serialize_bitmap_datum,
+};
+use crate::table::bitmap_global_index_writer::BitmapGlobalIndexWriter;
+use crate::table::{DataSplitBuilder, DeletionFile};
+
+fn table(extra: &[(&str, &str)]) -> Table {
+ let schema = Schema::builder()
+ .column("id", DataType::Int(IntType::new()))
+ .column("a", DataType::Int(IntType::new()))
+ .column("b", DataType::Int(IntType::new()))
+ .primary_key(["id"])
+ .option("bucket", "1")
+ .option("pk-btree.index.columns", "a")
+ .option("pk-bitmap.index.columns", "b")
+ .options(extra.iter().map(|(k, v)| (k.to_string(), v.to_string())))
+ .build()
+ .unwrap();
+ Table::new(
+ FileIOBuilder::new("memory").build().unwrap(),
+ Identifier::new("db", "t"),
+ "memory:/table".into(),
+ TableSchema::new(0, &schema),
+ None,
+ )
+}
+
+fn file(name: &str, rows: i64) -> DataFileMeta {
+ DataFileMeta {
+ file_name: name.into(),
+ file_size: 100,
+ row_count: rows,
+ min_key: vec![],
+ max_key: vec![],
+ key_stats: BinaryTableStats::new(vec![], vec![], vec![]),
+ value_stats: BinaryTableStats::new(vec![], vec![], vec![]),
+ min_sequence_number: 0,
+ max_sequence_number: 0,
+ schema_id: 0,
+ level: 1,
+ extra_files: vec![],
+ creation_time: None,
+ delete_row_count: Some(0),
+ embedded_index: None,
+ file_source: Some(1),
+ value_stats_cols: None,
+ external_path: None,
+ first_row_id: None,
+ write_cols: None,
+ column_max_sequence_numbers: None,
+ }
+}
+
+fn split(files: Vec<DataFileMeta>, raw: bool) -> DataSplit {
+ DataSplitBuilder::new()
+ .with_snapshot(7)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/table/bucket-0".into())
+ .with_total_buckets(1)
+ .with_data_files(files)
+ .with_raw_convertible(raw)
+ .build()
+ .unwrap()
+}
+
+fn source_meta(files: &[(&str, i64)], level: i32) -> Vec<u8> {
+ let mut bytes = vec![];
+ bytes.extend_from_slice(&1i32.to_be_bytes());
+ bytes.extend_from_slice(&level.to_be_bytes());
+ bytes.extend_from_slice(&(files.len() as i32).to_be_bytes());
+ for (name, count) in files {
+ assert!(name.is_ascii());
+ bytes.extend_from_slice(&(name.len() as u16).to_be_bytes());
+ bytes.extend_from_slice(name.as_bytes());
+ bytes.extend_from_slice(&count.to_be_bytes());
+ }
+ bytes
+}
+
+async fn payload(
+ table: &Table,
+ field: &str,
+ sources: &[(&str, i64)],
+ values: &[Option<i32>],
+ name: &str,
+ path: &str,
+) -> IndexManifestEntry {
+ let field = table
+ .schema()
+ .fields()
+ .iter()
+ .find(|f| f.name() == field)
+ .unwrap();
+ let ty = field.data_type();
+ let output = VecFileWrite::new();
+ let (meta, row_count) = if field.name() == "a" {
+ let mut writer = BTreeIndexWriter::with_comparator(
+ Box::new(output.clone()),
+ 256,
+ BlockCompressionType::None,
+ make_key_comparator(ty),
+ );
+ let mut sorted_values: Vec<_> = values.iter().enumerate().collect();
+ sorted_values.sort_by_key(|(_, value)| **value);
+ for (position, value) in sorted_values {
+ let key = value.map(|v| serialize_datum(&Datum::Int(v), ty));
+ writer.write(key.as_deref(), position as i64).await.unwrap();
+ }
+ let result = writer.finish().await.unwrap();
+ (result.meta, result.row_count as i64)
+ } else {
+ let mut writer = BitmapGlobalIndexWriter::new(
+ Box::new(output.clone()),
+ 256,
+ BlockCompressionType::None,
+ make_bitmap_key_comparator(ty),
+ );
+ for (position, value) in values.iter().enumerate() {
+ let key = value.map(|v| serialize_bitmap_datum(&Datum::Int(v),
ty));
+ writer.write(key.as_deref(), position as i64).unwrap();
+ }
+ let result = writer.finish().await.unwrap();
+ (result.meta, result.row_count as i64)
+ };
+ let bytes = output.to_vec();
+ table
+ .file_io()
+ .new_output(path)
+ .unwrap()
+ .write(bytes::Bytes::from(bytes.clone()))
+ .await
+ .unwrap();
+ IndexManifestEntry {
+ version: 1,
+ kind: FileKind::Add,
+ partition: BinaryRow::new(0).to_serialized_bytes(),
+ bucket: 0,
+ index_file: IndexFileMeta {
+ index_type: if field.name() == "a" {
+ "btree"
+ } else {
+ "bitmap"
+ }
+ .into(),
+ file_name: name.into(),
+ file_size: bytes.len() as i64,
+ row_count,
+ deletion_vectors_ranges: None,
+ external_path: None,
+ global_index_meta: Some(GlobalIndexMeta {
+ row_range_start: 0,
+ row_range_end: row_count - 1,
+ index_field_id: field.id(),
+ extra_field_ids: None,
+ index_meta: Some(meta.serialize()),
+ source_meta: Some(source_meta(sources, 1)),
+ }),
+ },
+ }
+}
+
+async fn run(
+ table: &Table,
+ splits: Vec<DataSplit>,
+ entries: &[IndexManifestEntry],
+ predicate: Predicate,
+) -> Vec<DataSplit> {
+ let defs = definitions(table.schema().fields(),
table.schema().options()).unwrap();
+ refine(table, splits, entries, &defs, &[predicate])
+ .await
+ .unwrap()
+}
+
+fn equal(table: &Table, column: &str, value: i32) -> Predicate {
+ PredicateBuilder::new(table.schema().fields())
+ .equal(column, Datum::Int(value))
+ .unwrap()
+}
+
+#[tokio::test]
+async fn retired_sources_keep_offsets_and_uncovered_sources_scan() {
+ let table = table(&[]);
+ let entry = payload(
+ &table,
+ "a",
+ &[("a-retired", 2), ("b-active", 3)],
+ &[Some(9), Some(9), Some(1), Some(9), Some(1)],
+ "a.idx",
+ "memory:/table/index/a.idx",
+ )
+ .await;
+ let result = run(
+ &table,
+ vec![split(vec![file("b-active", 3), file("c-new", 2)], true)],
+ &[entry],
+ equal(&table, "a", 9),
+ )
+ .await;
+ assert_eq!(result.len(), 2);
+ assert_eq!(
+ result[0].row_ranges(),
+ Some([RowRange::new(1, 1)].as_slice())
+ );
+ assert_eq!(result[1].row_ranges(), None);
+ assert!(result.iter().all(|s| !s.raw_convertible()));
+}
+
+#[tokio::test]
+async fn predicates_combine_after_each_fields_source_offset_is_removed() {
+ let table = table(&[]);
+ let a = payload(
+ &table,
+ "a",
+ &[("a-retired", 1), ("c-active", 3)],
+ &[Some(5), Some(5), Some(5), Some(0)],
+ "a.idx",
+ "memory:/table/index/a.idx",
+ )
+ .await;
+ let b = payload(
+ &table,
+ "b",
+ &[("b-retired", 2), ("c-active", 3)],
+ &[Some(5), Some(5), Some(0), Some(5), Some(5)],
+ "b.idx",
+ "memory:/table/index/b.idx",
+ )
+ .await;
+ let input = vec![split(vec![file("c-active", 3)], true)];
+ let entries = [a, b];
+ let and = run(
+ &table,
+ input.clone(),
+ &entries,
+ Predicate::and(vec![equal(&table, "a", 5), equal(&table, "b", 5)]),
+ )
+ .await;
+ assert_eq!(and[0].row_ranges(), Some([RowRange::new(1, 1)].as_slice()));
+ let or = run(
+ &table,
+ input.clone(),
+ &entries,
+ Predicate::or(vec![equal(&table, "a", 5), equal(&table, "b", 5)]),
+ )
+ .await;
+ assert_eq!(or[0].row_ranges(), Some([RowRange::new(0, 2)].as_slice()));
+ let unsupported_or = run(
+ &table,
+ input,
+ &entries,
+ Predicate::or(vec![equal(&table, "a", 5), equal(&table, "id", 8)]),
+ )
+ .await;
+ assert_eq!(unsupported_or[0].row_ranges(), None);
+}
+
+#[tokio::test]
+async fn non_raw_merge_inputs_stay_whole_even_when_index_returns_no_matches() {
+ let table = table(&[]);
+ let entry = payload(
+ &table,
+ "a",
+ &[("old", 2)],
+ &[Some(0), Some(0)],
+ "a.idx",
+ "memory:/table/index/a.idx",
+ )
+ .await;
+ let mut newer = file("new", 1);
+ newer.level = 0;
+ newer.file_source = Some(0);
+ let original = split(vec![file("old", 2), newer], false);
+ let result = run(
+ &table,
+ vec![original.clone()],
+ &[entry],
+ equal(&table, "a", 7),
+ )
+ .await;
+ assert_eq!(result, vec![original]);
+}
+
+#[tokio::test]
+async fn index_paths_and_aligned_deletion_files_are_preserved() {
+ for (bucket_local, external) in [(false, false), (true, false), (true,
true)] {
+ let table = table(&[(
+ "index-file-in-data-file-dir",
+ if bucket_local { "true" } else { "false" },
+ )]);
+ let path = if external {
+ "memory:/outside/index"
+ } else if bucket_local {
+ "memory:/table/bucket-0/a.idx"
+ } else {
+ "memory:/table/index/a.idx"
+ };
+ let mut entry = payload(
+ &table,
+ "a",
+ &[("a", 2), ("b", 2)],
+ &[Some(0), Some(1), Some(1), Some(0)],
+ "a.idx",
+ path,
+ )
+ .await;
+ if external {
+ entry.index_file.external_path = Some(path.into());
+ }
+ let deletion = DeletionFile::new("memory:/dv".into(), 3, 9, Some(1));
+ let input = DataSplitBuilder::new()
+ .with_snapshot(7)
+ .with_partition(BinaryRow::new(0))
+ .with_bucket(0)
+ .with_bucket_path("memory:/table/bucket-0".into())
+ .with_data_files(vec![file("a", 2), file("b", 2)])
+ .with_data_deletion_files(vec![None, Some(deletion.clone())])
+ .build()
+ .unwrap();
+ let result = run(&table, vec![input], &[entry], equal(&table, "a",
1)).await;
+ assert_eq!(
+ result[0].row_ranges(),
+ Some([RowRange::new(1, 1)].as_slice())
+ );
+ assert_eq!(
+ result[1].row_ranges(),
+ Some([RowRange::new(0, 0)].as_slice())
+ );
+ assert_eq!(result[0].data_deletion_files(), Some([None].as_slice()));
+ assert_eq!(
+ result[1].data_deletion_files(),
+ Some([Some(deletion)].as_slice())
+ );
+ }
+}
+
+#[tokio::test]
+async fn malformed_or_ambiguous_coverage_never_drops_data() {
+ let table = table(&[]);
+ let entry = payload(
+ &table,
+ "a",
+ &[("a", 2)],
+ &[Some(1), Some(2)],
+ "a.idx",
+ "memory:/table/index/a.idx",
+ )
+ .await;
+ let input = vec![split(vec![file("a", 2)], true)];
+ for case in 0..12 {
+ let mut broken = entry.clone();
+ let mut active = input.clone();
+ let meta = broken.index_file.global_index_meta.as_mut().unwrap();
+ match case {
+ 0 => meta.source_meta = Some(source_meta(&[("a", 1)], 1)),
+ 1 => meta.source_meta = Some(source_meta(&[("retired", 2)], 1)),
+ 2 => meta.source_meta = Some(source_meta(&[("a", 1), ("a", 1)],
1)),
+ 3 => meta.source_meta = Some(source_meta(&[("z", 1), ("a", 1)],
1)),
+ 4 => meta.row_range_start = 1,
+ 5 => meta.row_range_end = 99,
+ 6 => broken.index_file.row_count = 1,
+ 7 => {
+ active = vec![split(
+ vec![{
+ let mut f = file("a", 2);
+ f.file_source = Some(0);
+ f
+ }],
+ true,
+ )]
+ }
+ 8 => meta.source_meta = Some(source_meta(&[("a", 2)], 2)),
+ 9 => broken.bucket = 1,
+ 10 => broken.partition = BinaryRow::new(1).to_serialized_bytes(),
+ 11 => meta.index_field_id = 0,
+ _ => unreachable!(),
+ }
+ assert_eq!(
+ run(&table, active.clone(), &[broken], equal(&table, "a",
99)).await,
+ active,
+ "case {case}"
+ );
+ }
+ let mut duplicate = entry.clone();
+ duplicate.index_file.file_name = "other.idx".into();
+ assert_eq!(
+ run(
+ &table,
+ input.clone(),
+ &[entry, duplicate],
+ equal(&table, "a", 99)
+ )
+ .await,
+ input
+ );
+}
+
+#[tokio::test]
+async fn corrupt_payload_and_unsupported_and_conjunct_fall_back_safely() {
+ let table = table(&[]);
+ let entry = payload(
+ &table,
+ "a",
+ &[("a", 2)],
+ &[Some(1), Some(2)],
+ "a.idx",
+ "memory:/table/index/a.idx",
+ )
+ .await;
+ let input = vec![split(vec![file("a", 2)], true)];
+ let and = run(
+ &table,
+ input.clone(),
+ std::slice::from_ref(&entry),
+ Predicate::and(vec![equal(&table, "a", 2), equal(&table, "id", 1)]),
+ )
+ .await;
+ assert_eq!(and[0].row_ranges(), Some([RowRange::new(1, 1)].as_slice()));
+ table
+ .file_io()
+ .new_output("memory:/table/index/a.idx")
+ .unwrap()
+ .write(bytes::Bytes::from_static(b"corrupt"))
+ .await
+ .unwrap();
+ let result = run(&table, input, &[entry], equal(&table, "a", 2)).await;
+ assert_eq!(result.len(), 1);
+ assert_eq!(result[0].row_ranges(), None);
+}
+
+#[test]
+fn per_column_read_options_follow_java_resolution_and_conflicts() {
+ let table = table(&[(
+ "fields.a.pk-btree.index.options",
+ r#"{"fallback-scan-max-size":"17 kb","cache-size":"0"}"#,
+ )]);
+ let defs = definitions(table.schema().fields(),
table.schema().options()).unwrap();
+ let opts = CoreOptions::new(&defs[0].options);
+ assert_eq!(
+ opts.btree_index_fallback_scan_max_size().unwrap(),
+ 17 * 1024
+ );
+ assert_eq!(opts.btree_index_data_block_cache_size().unwrap(), 0);
+ for json in [
+ "[]",
+ "null",
+ r#"{"a":null}"#,
+ r#"{" ":"v"}"#,
+ r#"{"fallback-scan-max-size":"10 kb"}"#,
+ ] {
+ let options = Options::from([
+ ("pk-btree.index.columns".into(), "a".into()),
+ ("btree-index.fallback-scan-max-size".into(), "1 kb".into()),
+ ("fields.a.pk-btree.index.options".into(), json.into()),
+ ]);
+ assert!(
+ definitions(table.schema().fields(), &options).is_err(),
+ "{json}"
+ );
+ }
+ for options in [
+ Options::from([("pk-btree.index.columns".into(), "a, a".into())]),
+ Options::from([
+ ("pk-btree.index.columns".into(), "a".into()),
+ ("pk-bitmap.index.columns".into(), "a".into()),
+ ]),
+ ] {
+ assert!(definitions(table.schema().fields(), &options).is_err());
+ }
+}
+
+#[tokio::test]
+async fn combined_bounds_respect_per_column_fallback_scan_budget() {
+ for budget in ["0", "17 kb"] {
+ let json = format!(r#"{{"fallback-scan-max-size":"{budget}"}}"#);
+ let table = table(&[("fields.a.pk-btree.index.options", &json)]);
+ let entry = payload(
+ &table,
+ "a",
+ &[("a", 4)],
+ &[Some(1), Some(2), Some(3), Some(4)],
+ "a.idx",
+ "memory:/table/index/a.idx",
+ )
+ .await;
+ let pb = PredicateBuilder::new(table.schema().fields());
+ let predicate = Predicate::and(vec![
+ pb.greater_or_equal("a", Datum::Int(2)).unwrap(),
+ pb.less_or_equal("a", Datum::Int(3)).unwrap(),
+ ]);
+ let result = run(
+ &table,
+ vec![split(vec![file("a", 4)], true)],
+ &[entry],
+ predicate,
+ )
+ .await;
+ if budget == "0" {
+ assert_eq!(result[0].row_ranges(), None);
+ } else {
+ assert_eq!(
+ result[0].row_ranges(),
+ Some([RowRange::new(1, 2)].as_slice())
+ );
+ }
+ }
+}
+
+#[tokio::test]
+async fn shared_query_is_cached_before_localizing_each_source() {
+ let table = table(&[]);
+ let entry = payload(
+ &table,
+ "b",
+ &[("a", 2), ("b", 2)],
+ &[Some(1), Some(0), Some(0), Some(1)],
+ "b.idx",
+ "memory:/table/index/b.idx",
+ )
+ .await;
+ let defs = definitions(table.schema().fields(),
table.schema().options()).unwrap();
+ let input = split(vec![file("a", 2), file("b", 2)], true);
+ let (groups, sources) = plan_groups(&table, std::slice::from_ref(&input),
&[entry], &defs);
+ let mut evaluation = Evaluation {
+ table: &table,
+ definitions: &defs,
+ groups,
+ sources,
+ };
+ let predicate = equal(&table, "b", 1);
+ assert_eq!(
+ evaluation
+ .evaluate_field(&input, &input.data_files()[0], "b", &predicate)
+ .await
+ .unwrap(),
+ Some(vec![RowRange::new(0, 0)])
+ );
+ table
+ .file_io()
+ .delete_file("memory:/table/index/b.idx")
+ .await
+ .unwrap();
+ assert_eq!(
+ evaluation
+ .evaluate_field(&input, &input.data_files()[1], "b", &predicate)
+ .await
+ .unwrap(),
+ Some(vec![RowRange::new(1, 1)])
+ );
+}
+
+#[tokio::test]
+async fn invalid_payload_positions_and_fragmentation_keep_full_file() {
+ let table = table(&[]);
+ let mut invalid = payload(
+ &table,
+ "b",
+ &[("a", 2)],
+ &[Some(0), Some(0), Some(1)],
+ "bad.idx",
+ "memory:/table/index/bad.idx",
+ )
+ .await;
+ invalid.index_file.row_count = 2;
+ invalid
+ .index_file
+ .global_index_meta
+ .as_mut()
+ .unwrap()
+ .row_range_end = 1;
+ let result = run(
+ &table,
+ vec![split(vec![file("a", 2)], true)],
+ &[invalid],
+ equal(&table, "b", 1),
+ )
+ .await;
+ assert_eq!(result.len(), 1);
+ assert_eq!(result[0].row_ranges(), None);
+ for matching_rows in [MAX_POSITION_RANGES, MAX_POSITION_RANGES + 1] {
+ let row_count = matching_rows * 2;
+ let values: Vec<_> = (0..row_count).map(|i| Some((i % 2) as
i32)).collect();
+ let entry = payload(
+ &table,
+ "b",
+ &[("a", row_count as i64)],
+ &values,
+ "fragmented.idx",
+ "memory:/table/index/fragmented.idx",
+ )
+ .await;
+ let result = run(
+ &table,
+ vec![split(vec![file("a", row_count as i64)], true)],
+ &[entry],
+ equal(&table, "b", 1),
+ )
+ .await;
+ if matching_rows == MAX_POSITION_RANGES {
+ assert_eq!(result[0].row_ranges().unwrap().len(),
MAX_POSITION_RANGES);
+ } else {
+ assert_eq!(result[0].row_ranges(), None);
+ }
+ }
+}
+
+#[tokio::test]
+async fn nulls_and_empty_matches_preserve_index_semantics() {
+ let table = table(&[]);
+ for column in ["a", "b"] {
+ let name = format!("{column}.idx");
+ let entry = payload(
+ &table,
+ column,
+ &[("a", 4)],
+ &[None, Some(1), None, Some(2)],
+ &name,
+ &format!("memory:/table/index/{name}"),
+ )
+ .await;
+ let pb = PredicateBuilder::new(table.schema().fields());
+ let result = run(
+ &table,
+ vec![split(vec![file("a", 4)], true)],
+ std::slice::from_ref(&entry),
+ pb.is_null(column).unwrap(),
+ )
+ .await;
+ assert_eq!(
+ result[0].row_ranges(),
+ Some([RowRange::new(0, 0), RowRange::new(2, 2)].as_slice())
+ );
+ assert!(run(
+ &table,
+ vec![split(vec![file("a", 4)], true)],
+ &[entry],
+ equal(&table, column, 99)
+ )
+ .await
+ .is_empty());
+ }
+}
+
+#[tokio::test]
+async fn
positions_at_java_integer_max_are_allowed_but_larger_positions_fall_back() {
+ let table = table(&[]);
+ let ty = table.schema().fields()[2].data_type();
+ for position in [i64::from(i32::MAX), i64::from(i32::MAX) + 1] {
+ let row_count = position + 1;
+ let output = VecFileWrite::new();
+ let mut writer = BitmapGlobalIndexWriter::new(
+ Box::new(output.clone()),
+ 256,
+ BlockCompressionType::None,
+ make_bitmap_key_comparator(ty),
+ );
+ // Sparse postings exercise the position boundary with a few bytes,
+ // without allocating a multi-billion-row data file.
+ writer
+ .write(Some(&serialize_bitmap_datum(&Datum::Int(0), ty)), 0)
+ .unwrap();
+ writer
+ .write(Some(&serialize_bitmap_datum(&Datum::Int(1), ty)), position)
+ .unwrap();
+ let result = writer
+ .finish_with_source_row_count(row_count as u64)
+ .await
+ .unwrap();
+ let bytes = output.to_vec();
+ table
+ .file_io()
+ .new_output("memory:/table/index/large.idx")
+ .unwrap()
+ .write(bytes::Bytes::from(bytes.clone()))
+ .await
+ .unwrap();
+ let entry = IndexManifestEntry {
+ version: 1,
+ kind: FileKind::Add,
+ partition: BinaryRow::new(0).to_serialized_bytes(),
+ bucket: 0,
+ index_file: IndexFileMeta {
+ index_type: "bitmap".into(),
+ file_name: "large.idx".into(),
+ file_size: bytes.len() as i64,
+ row_count,
+ deletion_vectors_ranges: None,
+ external_path: None,
+ global_index_meta: Some(GlobalIndexMeta {
+ row_range_start: 0,
+ row_range_end: row_count - 1,
+ index_field_id: 2,
+ extra_field_ids: None,
+ index_meta: Some(result.meta.serialize()),
+ source_meta: Some(source_meta(&[("a", row_count)], 1)),
+ }),
+ },
+ };
+ let selected = run(
+ &table,
+ vec![split(vec![file("a", row_count)], true)],
+ &[entry],
+ equal(&table, "b", 1),
+ )
+ .await;
+ if position == i64::from(i32::MAX) {
+ assert_eq!(
+ selected[0].row_ranges(),
+ Some([RowRange::new(position, position)].as_slice())
+ );
+ } else {
+ assert_eq!(selected[0].row_ranges(), None);
+ }
+ }
+}
diff --git a/crates/paimon/src/table/source.rs
b/crates/paimon/src/table/source.rs
index 7c106fc1..d50980ab 100644
--- a/crates/paimon/src/table/source.rs
+++ b/crates/paimon/src/table/source.rs
@@ -47,8 +47,9 @@ 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 row-tracked tables, or physical positions for raw
-/// tables without row tracking.
+/// path: stable row IDs for row-tracked append/Data Evolution tables,
+/// split-local positions for ordinary append tables, or file-local physical
+/// positions for primary-key IndexedSplits.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RowRange {
from: i64,
@@ -498,9 +499,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<DeletionFile> with null elements).
data_deletion_files: Option<Arc<[Option<DeletionFile>]>>,
- /// 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.
+ /// IndexedSplit-compatible ranges. Append/Data Evolution reads use stable
+ /// row IDs when tracked, otherwise split-local positions over
`data_files`.
+ /// Primary-key IndexedSplits use file-local physical positions,
independent
+ /// of row tracking.
row_ranges: Option<Arc<[RowRange]>>,
/// Whether the split can be read raw, without the merge reader: its
/// physical rows are exactly its logical rows (modulo deletion files).
@@ -587,6 +589,21 @@ impl DataSplit {
self
}
+ /// Java PrimaryKeySortedIndexResult preserves each non-raw split whole,
+ /// and turns raw splits into single-file merge-reader inputs. Its ranges
+ /// address physical positions, including when the file tracks row IDs.
+ pub(crate) fn for_pk_index_file(&self, file_index: usize) -> Self {
+ let mut split = self.clone();
+ split.data_files = vec![self.data_files[file_index].clone()].into();
+ split.data_deletion_files = self
+ .data_deletion_files
+ .as_ref()
+ .map(|files| vec![files[file_index].clone()].into());
+ split.row_ranges = None;
+ split.raw_convertible = false;
+ split
+ }
+
/// 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 {
diff --git a/crates/paimon/src/table/table_scan.rs
b/crates/paimon/src/table/table_scan.rs
index 367e0de1..b56bf4a4 100644
--- a/crates/paimon/src/table/table_scan.rs
+++ b/crates/paimon/src/table/table_scan.rs
@@ -2178,6 +2178,25 @@ impl<'a> PaimonTableScan<'a> {
}
}
+ fn pk_sorted_index_definitions(
+ &self,
+ ) -> crate::Result<Vec<super::pk_sorted_index_scan::Definition>> {
+ let schema = self.table.schema();
+ let options = schema.core_options();
+ if schema.primary_keys().is_empty()
+ || self.is_streaming()
+ || self.scan_all_files
+ || self.row_range_optimization_disabled
+ || self.snapshot_mode().is_some()
+ || options.query_auth_enabled()
+ || !options.global_index_enabled()
+ || self.data_predicates.is_empty()
+ {
+ return Ok(Vec::new());
+ }
+ super::pk_sorted_index_scan::definitions(schema.fields(),
schema.options())
+ }
+
async fn read_index_manifest_entries(
&self,
snapshot: &Snapshot,
@@ -2487,6 +2506,7 @@ impl<'a> PaimonTableScan<'a> {
None,
before_index_entries,
None,
+ &[],
None,
)
.await?;
@@ -2497,6 +2517,7 @@ impl<'a> PaimonTableScan<'a> {
None,
after_index_entries,
None,
+ &[],
None,
)
.await?;
@@ -2558,10 +2579,11 @@ impl<'a> PaimonTableScan<'a> {
let data_evolution_enabled = core_options.data_evolution_enabled();
let global_index_settings =
self.global_index_scan_settings(&core_options,
data_evolution_enabled)?;
+ let pk_sorted_index_definitions = self.pk_sorted_index_definitions()?;
let index_entries = self
.read_index_manifest_entries(
snapshot,
- global_index_settings.is_some(),
+ global_index_settings.is_some() ||
!pk_sorted_index_definitions.is_empty(),
core_options.deletion_vectors_enabled() &&
!self.is_streaming(),
)
.await?;
@@ -2604,11 +2626,13 @@ impl<'a> PaimonTableScan<'a> {
data_evolution_read_field_ids,
index_entries,
effective_row_ranges,
+ &pk_sorted_index_definitions,
trace,
)
.await
}
+ #[allow(clippy::too_many_arguments)]
async fn plan_snapshot_from_entries(
&self,
snapshot: Snapshot,
@@ -2616,6 +2640,7 @@ impl<'a> PaimonTableScan<'a> {
data_evolution_read_field_ids: Option<&HashSet<i32>>,
index_entries: Option<Vec<IndexManifestEntry>>,
effective_row_ranges: Option<Vec<RowRange>>,
+ pk_sorted_index_definitions:
&[super::pk_sorted_index_scan::Definition],
mut trace: Option<&mut ScanTrace>,
) -> crate::Result<Plan> {
let table_path = self.table.location();
@@ -3034,6 +3059,18 @@ impl<'a> PaimonTableScan<'a> {
} else {
splits
};
+ let splits = if !pk_sorted_index_definitions.is_empty() {
+ super::pk_sorted_index_scan::refine(
+ &self.table,
+ splits,
+ index_entries.as_deref().unwrap_or_default(),
+ pk_sorted_index_definitions,
+ &self.data_predicates,
+ )
+ .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();
@@ -3506,6 +3543,7 @@ mod tests {
None,
None,
None,
+ &[],
None,
)
.await
@@ -4974,6 +5012,7 @@ mod tests {
None,
None,
None,
+ &[],
None,
)
.await
@@ -5051,6 +5090,7 @@ mod tests {
None,
None,
None,
+ &[],
None,
)
.await
@@ -5147,6 +5187,7 @@ mod tests {
None,
None,
None,
+ &[],
None,
)
.await
@@ -6596,6 +6637,64 @@ mod tests {
assert!(entries.is_none());
}
+ #[tokio::test]
+ async fn
test_pk_index_scan_uses_resolved_definitions_instead_of_option_presence() {
+ let table = pk_stats_gate_table("memory:/pk_index_definition_gate");
+ setup_scan_trace_dirs(&table).await;
+ TableCommit::new(table.clone(), "definition-gate".into())
+ .commit(vec![CommitMessage::new(
+ BinaryRow::new(0).to_serialized_bytes(),
+ 0,
+ vec![pk_stats_file("data.parquet", (1, 2), (10, 20))],
+ )])
+ .await
+ .unwrap();
+ let snapshot = table
+ .snapshot_manager()
+ .get_latest_snapshot()
+ .await
+ .unwrap()
+ .unwrap();
+ let mut json = serde_json::to_value(&snapshot).unwrap();
+ json["indexManifest"] = serde_json::json!("missing-index-manifest");
+ let snapshot: Snapshot = serde_json::from_value(json).unwrap();
+ let predicate = PredicateBuilder::new(table.schema().fields())
+ .equal("value", Datum::Int(10))
+ .unwrap();
+ for (key, columns, enabled, uses_index) in [
+ ("pk-btree.index.columns", "", true, false),
+ ("pk-btree.index.columns", "unknown", true, false),
+ ("pk-bitmap.index.columns", "unknown", true, false),
+ ("pk-btree.index.columns", "value", false, false),
+ ("pk-btree.index.columns", "value", true, true),
+ ("pk-bitmap.index.columns", "value", true, true),
+ ] {
+ let configured = table.copy_with_options(HashMap::from([
+ (key.into(), columns.into()),
+ ("global-index.enabled".into(), enabled.to_string()),
+ ]));
+ let scan =
+ PaimonTableScan::new(&configured, None,
vec![predicate.clone()], None, None, None);
+ let plan = scan.plan_snapshot(snapshot.clone(), None, None).await;
+ if uses_index {
+ assert!(plan
+ .unwrap_err()
+ .to_string()
+ .contains("missing-index-manifest"));
+ } else {
+ let plan = plan.unwrap();
+ assert_eq!(plan.snapshot_id(), Some(snapshot.id()));
+ let files: Vec<_> = plan
+ .splits()
+ .iter()
+ .flat_map(|split| split.data_files())
+ .map(|file| file.file_name.as_str())
+ .collect();
+ assert_eq!(files, vec!["data.parquet"]);
+ }
+ }
+ }
+
// ======================== Bucket predicate filtering
========================
fn bucket_key_fields() -> Vec<DataField> {
