From 35157d8360f2897ccf4e3a3b37de6ae82aa44f38 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Wed, 2 Sep 2026 16:59:58 +0800 Subject: [PATCH] feat(file_index): integrate pruning into raw data-file reads Use embedded and sidecar indexes to prune raw scans while preserving residual predicates and schema-evolution safety. Signed-off-by: QuakeWang --- crates/paimon/src/file_index/evaluator.rs | 640 +++++++++++++++++ .../src/file_index/file_index_format.rs | 175 ++++- .../src/file_index/file_index_reader.rs | 2 +- .../src/file_index/file_indexer_factory.rs | 4 + crates/paimon/src/file_index/mod.rs | 3 +- crates/paimon/src/spec/core_options.rs | 21 + crates/paimon/src/table/data_file_reader.rs | 665 +++++++++++++++++- crates/paimon/src/table/table_read.rs | 111 ++- 8 files changed, 1551 insertions(+), 70 deletions(-) create mode 100644 crates/paimon/src/file_index/evaluator.rs diff --git a/crates/paimon/src/file_index/evaluator.rs b/crates/paimon/src/file_index/evaluator.rs new file mode 100644 index 000000000..ce84aa756 --- /dev/null +++ b/crates/paimon/src/file_index/evaluator.rs @@ -0,0 +1,640 @@ +// 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 std::collections::HashSet; + +use bytes::Bytes; + +use crate::file_index::file_index_predicate::FileIndexPredicate; +use crate::file_index::file_index_result::FileIndexResult; +use crate::file_index::FileIndexFormatReader; +use crate::io::FileIO; +use crate::spec::{DataField, DataFileMeta, DataType, Datum, Predicate, DATA_FILE_INDEX_SUFFIX}; +use crate::Error; + +/// Evaluate the usable data predicates against one data file's FileIndex. +pub(crate) async fn evaluate_file_index( + file_io: &FileIO, + bucket_path: &str, + file: &DataFileMeta, + table_fields: &[DataField], + data_fields: &[DataField], + predicates: &[Predicate], +) -> crate::Result { + if predicates.is_empty() || !(0..=i64::from(i32::MAX)).contains(&file.row_count) { + return Ok(FileIndexResult::Remain); + } + + let Some(predicate) = remap_predicates(table_fields, data_fields, predicates) else { + return Ok(FileIndexResult::Remain); + }; + let mut required_columns = HashSet::new(); + collect_required_columns(&predicate, &mut required_columns); + if required_columns.is_empty() { + return Ok(FileIndexResult::Remain); + } + + let file_index = if let Some(embedded) = &file.embedded_index { + FileIndexFormatReader::get_file_index_from_bytes(Bytes::copy_from_slice(embedded)).await? + } else { + let sidecars = file + .extra_files + .iter() + .filter(|name| name.ends_with(DATA_FILE_INDEX_SUFFIX)) + .collect::>(); + match sidecars.as_slice() { + [] => return Ok(FileIndexResult::Remain), + [sidecar] => { + let path = file.aligned_file_path(bucket_path, sidecar); + FileIndexFormatReader::get_file_index(file_io.new_input(&path)?).await? + } + _ => { + return Err(Error::DataInvalid { + message: format!( + "Found more than one index file for data file '{}': {}", + file.file_name, + sidecars + .iter() + .map(|name| name.as_str()) + .collect::>() + .join(" and ") + ), + source: None, + }) + } + } + }; + + let readers = file_index + .create_index_readers(data_fields, &required_columns) + .await?; + if readers.values().all(Vec::is_empty) { + return Ok(FileIndexResult::Remain); + } + Ok(FileIndexPredicate::new(readers).evaluate(&predicate)) +} + +fn remap_predicates( + table_fields: &[DataField], + data_fields: &[DataField], + predicates: &[Predicate], +) -> Option { + let remapped = predicates + .iter() + .filter_map(|predicate| remap_predicate(table_fields, data_fields, predicate)) + .collect::>(); + (!remapped.is_empty()).then(|| Predicate::and(remapped)) +} + +fn remap_predicate( + table_fields: &[DataField], + data_fields: &[DataField], + predicate: &Predicate, +) -> Option { + match predicate { + Predicate::Leaf { + column, + index, + data_type, + op, + literals, + } => { + let table_field = table_fields.get(*index)?; + if table_field.name() != column { + return None; + } + let (data_index, data_field) = data_fields + .iter() + .enumerate() + .find(|(_, field)| field.id() == table_field.id())?; + let literals = devolve_literals(data_type, data_field.data_type(), literals)?; + Some(Predicate::Leaf { + column: data_field.name().to_string(), + index: data_index, + data_type: data_field.data_type().clone(), + op: *op, + literals, + }) + } + Predicate::And(children) => { + let remapped = children + .iter() + .filter_map(|child| remap_predicate(table_fields, data_fields, child)) + .collect::>(); + (!remapped.is_empty()).then(|| Predicate::and(remapped)) + } + Predicate::Or(children) => { + let remapped = children + .iter() + .map(|child| remap_predicate(table_fields, data_fields, child)) + .collect::>>()?; + Some(Predicate::or(remapped)) + } + Predicate::Not(inner) => { + remap_predicate_exact(table_fields, data_fields, inner).map(Predicate::negate) + } + Predicate::AlwaysTrue => Some(Predicate::AlwaysTrue), + Predicate::AlwaysFalse => Some(Predicate::AlwaysFalse), + } +} + +/// Remap only complete subtrees, as required below negation where widening is unsafe. +fn remap_predicate_exact( + table_fields: &[DataField], + data_fields: &[DataField], + predicate: &Predicate, +) -> Option { + match predicate { + Predicate::And(children) => children + .iter() + .map(|child| remap_predicate_exact(table_fields, data_fields, child)) + .collect::>>() + .map(Predicate::and), + Predicate::Or(children) => children + .iter() + .map(|child| remap_predicate_exact(table_fields, data_fields, child)) + .collect::>>() + .map(Predicate::or), + Predicate::Not(inner) => { + remap_predicate_exact(table_fields, data_fields, inner).map(Predicate::negate) + } + Predicate::Leaf { .. } | Predicate::AlwaysTrue | Predicate::AlwaysFalse => { + remap_predicate(table_fields, data_fields, predicate) + } + } +} + +fn devolve_literals( + table_type: &DataType, + data_type: &DataType, + literals: &[Datum], +) -> Option> { + if same_type_ignoring_nullability(table_type, data_type) { + return Some(literals.to_vec()); + } + if !is_integer_type(table_type) || !is_integer_type(data_type) { + return None; + } + literals + .iter() + .map(|literal| { + let value = integer_value(table_type, literal)?; + integer_datum(data_type, value) + }) + .collect() +} + +fn same_type_ignoring_nullability(left: &DataType, right: &DataType) -> bool { + match ( + left.copy_with_nullable(true), + right.copy_with_nullable(true), + ) { + (Ok(left), Ok(right)) => left == right, + _ => false, + } +} + +fn is_integer_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::TinyInt(_) | DataType::SmallInt(_) | DataType::Int(_) | DataType::BigInt(_) + ) +} + +fn integer_value(data_type: &DataType, datum: &Datum) -> Option { + match (data_type, datum) { + (DataType::TinyInt(_), Datum::TinyInt(value)) => Some(i64::from(*value)), + (DataType::SmallInt(_), Datum::SmallInt(value)) => Some(i64::from(*value)), + (DataType::Int(_), Datum::Int(value)) => Some(i64::from(*value)), + (DataType::BigInt(_), Datum::Long(value)) => Some(*value), + _ => None, + } +} + +fn integer_datum(data_type: &DataType, value: i64) -> Option { + match data_type { + DataType::TinyInt(_) => i8::try_from(value).ok().map(Datum::TinyInt), + DataType::SmallInt(_) => i16::try_from(value).ok().map(Datum::SmallInt), + DataType::Int(_) => i32::try_from(value).ok().map(Datum::Int), + DataType::BigInt(_) => Some(Datum::Long(value)), + _ => None, + } +} + +fn collect_required_columns(predicate: &Predicate, columns: &mut HashSet) { + match predicate { + Predicate::Leaf { column, .. } => { + columns.insert(column.clone()); + } + Predicate::And(children) | Predicate::Or(children) => { + for child in children { + collect_required_columns(child, columns); + } + } + Predicate::Not(inner) => collect_required_columns(inner, columns), + Predicate::AlwaysTrue | Predicate::AlwaysFalse => {} + } +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::common::Options; + use crate::file_index::file_index_format::write_column_indexes; + use crate::file_index::file_index_result::FileIndexResult; + use crate::file_index::file_indexer_factory::{FileIndexerFactory, BITMAP_INDEX}; + use crate::io::FileIOBuilder; + use crate::spec::stats::BinaryTableStats; + use crate::spec::{ + BigIntType, FloatType, IntType, PredicateBuilder, PredicateOperator, VarCharType, + }; + + fn field(id: i32, name: &str, data_type: DataType) -> DataField { + DataField::new(id, name.to_string(), data_type) + } + + fn data_file(row_count: i64) -> DataFileMeta { + DataFileMeta { + file_name: "part-0.parquet".to_string(), + file_size: 1, + row_count, + min_key: Vec::new(), + max_key: Vec::new(), + key_stats: BinaryTableStats::empty(), + value_stats: BinaryTableStats::empty(), + 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, + file_source: None, + value_stats_cols: None, + external_path: None, + first_row_id: None, + write_cols: None, + column_max_sequence_numbers: None, + } + } + + async fn bitmap_index_bytes( + path: &str, + column: &str, + data_type: DataType, + values: &[Datum], + ) -> crate::Result { + let mut writer = + FileIndexerFactory::create_writer(BITMAP_INDEX, data_type, &Options::new())?; + for value in values { + writer.write(Some(value))?; + } + let indexes = HashMap::from([( + column.to_string(), + HashMap::from([(BITMAP_INDEX.to_string(), Some(writer.serialized_bytes()?))]), + )]); + write_column_indexes(path, indexes) + .await? + .to_input_file() + .read() + .await + } + + fn assert_int_leaf( + predicate: Predicate, + expected_column: &str, + expected_index: usize, + expected_literal: i32, + ) { + assert!(matches!( + predicate, + Predicate::Leaf { + column, + index, + data_type: DataType::Int(_), + op: PredicateOperator::Eq, + literals, + } if column == expected_column + && index == expected_index + && literals == vec![Datum::Int(expected_literal)] + )); + } + + #[test] + fn test_remap_predicate_uses_field_id_for_rename_reorder_and_integer_devolution() { + let table_fields = vec![ + field(0, "new_id", DataType::BigInt(BigIntType::new())), + field(1, "name", DataType::VarChar(VarCharType::new(20).unwrap())), + ]; + let data_fields = vec![ + table_fields[1].clone(), + field(0, "old_id", DataType::Int(IntType::new())), + ]; + let predicate = PredicateBuilder::new(&table_fields) + .equal("new_id", Datum::Long(42)) + .unwrap(); + + let remapped = remap_predicate(&table_fields, &data_fields, &predicate).unwrap(); + + assert_int_leaf(remapped, "old_id", 1, 42); + } + + #[test] + fn test_remap_predicate_falls_back_for_unsafe_schema_changes() { + let table_fields = vec![ + field(0, "id", DataType::BigInt(BigIntType::new())), + field(1, "added", DataType::Int(IntType::new())), + ]; + let data_fields = vec![field(0, "id", DataType::Int(IntType::new()))]; + let builder = PredicateBuilder::new(&table_fields); + + assert!(remap_predicate( + &table_fields, + &data_fields, + &builder + .equal("id", Datum::Long(i64::from(i32::MAX) + 1)) + .unwrap(), + ) + .is_none()); + assert!(remap_predicate( + &table_fields, + &data_fields, + &builder.equal("added", Datum::Int(1)).unwrap(), + ) + .is_none()); + + let promoted_table = vec![field(0, "id", DataType::Float(FloatType::new()))]; + let promoted_predicate = PredicateBuilder::new(&promoted_table) + .equal("id", Datum::Float(1.0)) + .unwrap(); + assert!(remap_predicate(&promoted_table, &data_fields, &promoted_predicate).is_none()); + } + + #[test] + fn test_remap_predicate_keeps_safe_and_child_but_requires_complete_or_and_not() { + let table_fields = vec![ + field(0, "id", DataType::Int(IntType::new())), + field(1, "added", DataType::Int(IntType::new())), + ]; + let data_fields = vec![table_fields[0].clone()]; + let builder = PredicateBuilder::new(&table_fields); + let safe = builder.equal("id", Datum::Int(1)).unwrap(); + let unsafe_predicate = builder.equal("added", Datum::Int(2)).unwrap(); + + let remapped_and = remap_predicate( + &table_fields, + &data_fields, + &Predicate::and(vec![safe.clone(), unsafe_predicate.clone()]), + ) + .unwrap(); + assert_int_leaf(remapped_and, "id", 0, 1); + assert!(remap_predicate( + &table_fields, + &data_fields, + &Predicate::or(vec![safe.clone(), unsafe_predicate.clone()]), + ) + .is_none()); + assert!(remap_predicate( + &table_fields, + &data_fields, + &Predicate::negate(unsafe_predicate.clone()), + ) + .is_none()); + + let exact_double_not = Predicate::Not(Box::new(Predicate::Not(Box::new(safe.clone())))); + assert_int_leaf( + remap_predicate(&table_fields, &data_fields, &exact_double_not).unwrap(), + "id", + 0, + 1, + ); + + let nested_not = Predicate::negate(Predicate::and(vec![ + Predicate::negate(safe), + unsafe_predicate, + ])); + assert!(remap_predicate(&table_fields, &data_fields, &nested_not).is_none()); + } + + #[tokio::test] + async fn test_evaluator_prefers_embedded_index_over_ambiguous_sidecars() -> crate::Result<()> { + let fields = vec![field(0, "id", DataType::Int(IntType::new()))]; + let bytes = bitmap_index_bytes( + "memory:/evaluator_embedded_source", + "id", + fields[0].data_type().clone(), + &[Datum::Int(1), Datum::Int(2)], + ) + .await?; + let mut file = data_file(2); + file.embedded_index = Some(bytes.to_vec()); + file.extra_files = vec!["first.index".to_string(), "second.index".to_string()]; + let predicate = PredicateBuilder::new(&fields).equal("id", Datum::Int(2))?; + let file_io = FileIOBuilder::new("memory").build()?; + + let result = evaluate_file_index( + &file_io, + "memory:/unused-bucket", + &file, + &fields, + &fields, + &[predicate], + ) + .await?; + + assert_eq!( + result, + FileIndexResult::Selection([1_u32].into_iter().collect()) + ); + Ok(()) + } + + #[tokio::test] + async fn test_evaluator_resolves_regular_and_external_sidecars() -> crate::Result<()> { + let fields = vec![field(0, "id", DataType::Int(IntType::new()))]; + let bytes = bitmap_index_bytes( + "memory:/evaluator_sidecar_source", + "id", + fields[0].data_type().clone(), + &[Datum::Int(1), Datum::Int(2)], + ) + .await?; + let predicate = PredicateBuilder::new(&fields).equal("id", Datum::Int(1))?; + let file_io = FileIOBuilder::new("memory").build()?; + + for (bucket_path, external_path, sidecar_path) in [ + ( + "memory:/regular/bucket-0", + None, + "memory:/regular/bucket-0/part-0.parquet.index", + ), + ( + "memory:/ignored/bucket-0", + Some("memory:/external/data/part-0.parquet".to_string()), + "memory:/external/data/part-0.parquet.index", + ), + ] { + file_io + .new_output(sidecar_path)? + .write(bytes.clone()) + .await?; + let mut file = data_file(2); + file.extra_files = vec!["part-0.parquet.index".to_string()]; + file.external_path = external_path; + + let result = evaluate_file_index( + &file_io, + bucket_path, + &file, + &fields, + &fields, + std::slice::from_ref(&predicate), + ) + .await?; + assert_eq!( + result, + FileIndexResult::Selection([0_u32].into_iter().collect()) + ); + } + Ok(()) + } + + #[tokio::test] + async fn test_evaluator_rejects_ambiguous_sidecars_without_embedded_index() -> crate::Result<()> + { + let fields = vec![field(0, "id", DataType::Int(IntType::new()))]; + let predicate = PredicateBuilder::new(&fields).equal("id", Datum::Int(1))?; + let mut file = data_file(1); + file.extra_files = vec![ + "first.index".to_string(), + "notes.txt".to_string(), + "second.index".to_string(), + ]; + let file_io = FileIOBuilder::new("memory").build()?; + + let error = evaluate_file_index( + &file_io, + "memory:/bucket-0", + &file, + &fields, + &fields, + &[predicate], + ) + .await + .unwrap_err(); + + assert!(matches!( + error, + Error::DataInvalid { message, .. } + if message.contains("first.index") && message.contains("second.index") + )); + Ok(()) + } + + #[tokio::test] + async fn test_evaluator_absent_or_unsupported_index_remains() -> crate::Result<()> { + let fields = vec![field(0, "id", DataType::Int(IntType::new()))]; + let predicate = PredicateBuilder::new(&fields).equal("id", Datum::Int(1))?; + let file_io = FileIOBuilder::new("memory").build()?; + let file = data_file(1); + assert_eq!( + evaluate_file_index( + &file_io, + "memory:/bucket-0", + &file, + &fields, + &fields, + std::slice::from_ref(&predicate), + ) + .await?, + FileIndexResult::Remain + ); + + let indexes = HashMap::from([( + "id".to_string(), + HashMap::from([( + "range-bitmap".to_string(), + Some(Bytes::from_static(b"unsupported payload")), + )]), + )]); + let bytes = write_column_indexes("memory:/evaluator_unsupported_source", indexes) + .await? + .to_input_file() + .read() + .await?; + let mut file = data_file(1); + file.embedded_index = Some(bytes.to_vec()); + assert_eq!( + evaluate_file_index( + &file_io, + "memory:/bucket-0", + &file, + &fields, + &fields, + &[predicate], + ) + .await?, + FileIndexResult::Remain + ); + Ok(()) + } + + #[tokio::test] + async fn test_evaluator_row_count_boundaries_are_conservative() -> crate::Result<()> { + let fields = vec![field(0, "id", DataType::Int(IntType::new()))]; + let predicate = PredicateBuilder::new(&fields).equal("id", Datum::Int(1))?; + let file_io = FileIOBuilder::new("memory").build()?; + + for row_count in [-1, i64::from(i32::MAX) + 1] { + let mut file = data_file(row_count); + file.embedded_index = Some(vec![0]); + assert_eq!( + evaluate_file_index( + &file_io, + "memory:/bucket-0", + &file, + &fields, + &fields, + std::slice::from_ref(&predicate), + ) + .await?, + FileIndexResult::Remain + ); + } + + for row_count in [0, i64::from(i32::MAX)] { + let mut file = data_file(row_count); + file.embedded_index = Some(vec![0]); + assert!(matches!( + evaluate_file_index( + &file_io, + "memory:/bucket-0", + &file, + &fields, + &fields, + std::slice::from_ref(&predicate), + ) + .await, + Err(Error::FileIndexFormatInvalid { .. }) + )); + } + Ok(()) + } +} diff --git a/crates/paimon/src/file_index/file_index_format.rs b/crates/paimon/src/file_index/file_index_format.rs index f5fcb22a9..1e24319bf 100644 --- a/crates/paimon/src/file_index/file_index_format.rs +++ b/crates/paimon/src/file_index/file_index_format.rs @@ -24,7 +24,7 @@ use crate::{ file_index_reader::{EmptyFileIndexReader, FileIndexReader}, file_indexer_factory::FileIndexerFactory, }, - io::{FileIO, FileRead, FileStatus, InputFile, OutputFile}, + io::{FileIO, FileRead, InputFile, OutputFile}, spec::{DataField, DataType}, Error, }; @@ -364,7 +364,6 @@ pub struct FileIndex { impl FileIndex { /// Constructs readers for the required columns described by this outer-format file. - #[allow(dead_code)] pub(crate) async fn create_index_readers( &self, fields: &[DataField], @@ -382,6 +381,9 @@ impl FileIndex { }; let mut column_readers = Vec::with_capacity(index_info.len()); for (identifier, info) in index_info { + if !FileIndexerFactory::is_supported(identifier) { + continue; + } if info.start_pos == EMPTY_INDEX_FLAG { column_readers.push(Box::new(EmptyFileIndexReader) as Box); continue; @@ -468,16 +470,27 @@ impl FileIndex { pub struct FileIndexFormatReader { reader: Box, - stat: FileStatus, + file_size: u64, } impl FileIndexFormatReader { pub async fn get_file_index(input_file: InputFile) -> crate::Result { let reader = input_file.reader().await?; - let mut file_reader = Self { - reader: Box::new(reader), - stat: input_file.metadata().await?, - }; + let file_size = input_file.metadata().await?.size; + Self::get_file_index_from_reader(Box::new(reader), file_size).await + } + + pub(crate) async fn get_file_index_from_bytes(bytes: Bytes) -> crate::Result { + let file_size = u64::try_from(bytes.len()) + .map_err(|_| format_invalid("embedded file index is too large"))?; + Self::get_file_index_from_reader(Box::new(BytesFileRead(bytes)), file_size).await + } + + async fn get_file_index_from_reader( + reader: Box, + file_size: u64, + ) -> crate::Result { + let mut file_reader = Self { reader, file_size }; let header = file_reader.read_header().await?; Ok(FileIndex { header, @@ -486,10 +499,10 @@ impl FileIndexFormatReader { } async fn read_header(&mut self) -> crate::Result>> { - if self.stat.size < FIXED_HEADER_LENGTH as u64 { + if self.file_size < FIXED_HEADER_LENGTH as u64 { return Err(format_invalid(format!( "truncated fixed header: need {FIXED_HEADER_LENGTH} bytes, but file has {}", - self.stat.size + self.file_size ))); } @@ -521,10 +534,10 @@ impl FileIndexFormatReader { "header length {head_length} is smaller than the minimum {MIN_HEADER_LENGTH}" ))); } - if head_length as u64 > self.stat.size { + if head_length as u64 > self.file_size { return Err(format_invalid(format!( "header length {head_length} exceeds file size {}", - self.stat.size + self.file_size ))); } @@ -548,7 +561,7 @@ impl FileIndexFormatReader { let index_name = read_java_utf(&mut buffer, "index name")?; let start_pos = read_i32(&mut buffer, "index start position")?; let length = read_i32(&mut buffer, "index length")?; - Self::validate_index_range(start_pos, length, head_length as u64, self.stat.size)?; + Self::validate_index_range(start_pos, length, head_length as u64, self.file_size)?; index_info_map.insert(index_name, IndexInfo { start_pos, length }); } @@ -624,6 +637,27 @@ impl FileIndexFormatReader { } } +struct BytesFileRead(Bytes); + +#[async_trait::async_trait] +impl FileRead for BytesFileRead { + async fn read(&self, range: std::ops::Range) -> crate::Result { + let start = usize::try_from(range.start) + .map_err(|_| format_invalid("embedded file index range start is too large"))?; + let end = usize::try_from(range.end) + .map_err(|_| format_invalid("embedded file index range end is too large"))?; + if start > end || end > self.0.len() { + return Err(format_invalid(format!( + "embedded file index range {}..{} exceeds byte length {}", + range.start, + range.end, + self.0.len() + ))); + } + Ok(self.0.slice(start..end)) + } +} + #[cfg(test)] mod file_index_format_tests { @@ -784,6 +818,53 @@ mod file_index_format_tests { Ok(()) } + #[tokio::test] + async fn test_embedded_bytes_reader_matches_sidecar_reader() -> crate::Result<()> { + let fields = [DataField::new( + 0, + "a".to_string(), + DataType::Int(IntType::new()), + )]; + let mut writer = FileIndexerFactory::create_writer( + BITMAP_INDEX, + fields[0].data_type().clone(), + &Options::new(), + )?; + for value in [Datum::Int(1), Datum::Int(2), Datum::Int(1)] { + writer.write(Some(&value))?; + } + let indexes = HashMap::from([( + "a".to_string(), + HashMap::from([(BITMAP_INDEX.to_string(), Some(writer.serialized_bytes()?))]), + )]); + let output = write_column_indexes("memory:/tmp/embedded_file_index", indexes).await?; + let bytes = output.clone().to_input_file().read().await?; + let sidecar = FileIndexFormatReader::get_file_index(output.to_input_file()).await?; + let embedded = FileIndexFormatReader::get_file_index_from_bytes(bytes).await?; + let required_columns = HashSet::from(["a".to_string()]); + let predicate = PredicateBuilder::new(&fields).equal("a", Datum::Int(1))?; + + let sidecar_result = FileIndexPredicate::new( + sidecar + .create_index_readers(&fields, &required_columns) + .await?, + ) + .evaluate(&predicate); + let embedded_result = FileIndexPredicate::new( + embedded + .create_index_readers(&fields, &required_columns) + .await?, + ) + .evaluate(&predicate); + + assert_eq!(embedded_result, sidecar_result); + assert_eq!( + embedded_result, + FileIndexResult::Selection([0_u32, 2].into_iter().collect()) + ); + Ok(()) + } + #[tokio::test] async fn test_writer_matches_java_modified_utf8_bytes() -> crate::Result<()> { let indexes = HashMap::from([( @@ -1098,7 +1179,7 @@ mod file_index_format_tests { assert_eq!(readers.len(), 3); assert_eq!(readers["a"].len(), 2); assert_eq!(readers["b"].len(), 1); - assert_eq!(readers["empty"].len(), 1); + assert!(readers["empty"].is_empty()); let predicate = FileIndexPredicate::new(readers); let builder = PredicateBuilder::new(&fields); @@ -1112,7 +1193,7 @@ mod file_index_format_tests { ); assert_eq!( predicate.evaluate(&builder.equal("empty", Datum::Int(1))?), - FileIndexResult::Skip + FileIndexResult::Remain ); Ok(()) @@ -1177,28 +1258,52 @@ mod file_index_format_tests { } #[tokio::test] - async fn test_composition_rejects_unknown_non_empty_identifier() -> crate::Result<()> { - let indexes = HashMap::from([( - "a".to_string(), - HashMap::from([("unknown".to_string(), Some(Bytes::from_static(b"payload")))]), - )]); - let output = write_column_indexes("memory:/tmp/unknown_file_index", indexes).await?; - let file_index = FileIndexFormatReader::get_file_index(output.to_input_file()).await?; - let fields = [DataField::new( - 0, - "a".to_string(), - DataType::Int(IntType::new()), - )]; - - let required_columns = HashSet::from(["a".to_string()]); - let error = match file_index - .create_index_readers(&fields, &required_columns) - .await - { - Ok(_) => panic!("unknown identifier must fail"), - Err(error) => error, + async fn test_composition_skips_unknown_payload_and_keeps_supported_reader() -> crate::Result<()> + { + let data_type = DataType::Int(IntType::new()); + let mut writer = + FileIndexerFactory::create_writer(BITMAP_INDEX, data_type.clone(), &Options::new())?; + writer.write(Some(&Datum::Int(1)))?; + let supported_payload = writer.serialized_bytes()?; + let unknown_payload = Bytes::from_static(b"must not be read"); + let supported_end = supported_payload.len() as u64; + let mut data = BytesMut::with_capacity(supported_payload.len() + unknown_payload.len()); + data.extend_from_slice(&supported_payload); + data.extend_from_slice(&unknown_payload); + let ranges = Arc::new(Mutex::new(Vec::new())); + let file_index = FileIndex { + reader: Box::new(TrackingFileRead { + data: data.freeze(), + ranges: Arc::clone(&ranges), + }), + header: HashMap::from([( + "a".to_string(), + HashMap::from([ + ( + BITMAP_INDEX.to_string(), + IndexInfo { + start_pos: 0, + length: supported_end as i32, + }, + ), + ( + "unknown".to_string(), + IndexInfo { + start_pos: supported_end as i32, + length: unknown_payload.len() as i32, + }, + ), + ]), + )]), }; - assert!(matches!(error, Error::Unsupported { .. })); + let fields = [DataField::new(0, "a".to_string(), data_type)]; + + let readers = file_index + .create_index_readers(&fields, &HashSet::from(["a".to_string()])) + .await?; + + assert_eq!(readers["a"].len(), 1); + assert_eq!(*ranges.lock().unwrap(), vec![0..supported_end]); Ok(()) } diff --git a/crates/paimon/src/file_index/file_index_reader.rs b/crates/paimon/src/file_index/file_index_reader.rs index be9860995..f2e722264 100644 --- a/crates/paimon/src/file_index/file_index_reader.rs +++ b/crates/paimon/src/file_index/file_index_reader.rs @@ -19,7 +19,7 @@ use crate::file_index::file_index_result::FileIndexResult; use crate::spec::{DataType, Datum, PredicateOperator}; /// Evaluates leaf predicates against one concrete file index. -pub(crate) trait FileIndexReader { +pub(crate) trait FileIndexReader: Send + Sync { /// Evaluates the fields carried by [`crate::spec::Predicate::Leaf`]. /// /// Readers must return [`FileIndexResult::Remain`] for unsupported operators. diff --git a/crates/paimon/src/file_index/file_indexer_factory.rs b/crates/paimon/src/file_index/file_indexer_factory.rs index c6fb4b4bb..7750e2aed 100644 --- a/crates/paimon/src/file_index/file_indexer_factory.rs +++ b/crates/paimon/src/file_index/file_indexer_factory.rs @@ -51,6 +51,10 @@ impl BuiltinFileIndexer { pub(crate) struct FileIndexerFactory; impl FileIndexerFactory { + pub(crate) fn is_supported(identifier: &str) -> bool { + matches!(identifier, BITMAP_INDEX | BLOOM_FILTER_INDEX) + } + pub(crate) fn create_writer( identifier: &str, data_type: DataType, diff --git a/crates/paimon/src/file_index/mod.rs b/crates/paimon/src/file_index/mod.rs index c070ebdc8..ba62ab8b2 100644 --- a/crates/paimon/src/file_index/mod.rs +++ b/crates/paimon/src/file_index/mod.rs @@ -16,11 +16,12 @@ // under the License. // Concrete readers/writers and predicate plumbing stay crate-private until -// data-writer and scan integration land in later changes. +// writer integration and the API surface are ready to stabilize. #[allow(dead_code)] pub(crate) mod bitmap; #[allow(dead_code)] pub(crate) mod bloom_filter; +pub(crate) mod evaluator; mod file_index_format; #[allow(dead_code)] pub(crate) mod file_index_predicate; diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index 7afff6f67..7c0b74a66 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -23,6 +23,7 @@ const DELETION_VECTORS_ENABLED_OPTION: &str = "deletion-vectors.enabled"; const DELETION_VECTORS_MERGE_ON_READ_OPTION: &str = "deletion-vectors.merge-on-read"; pub(crate) const QUERY_AUTH_ENABLED_OPTION: &str = "query-auth.enabled"; const DATA_EVOLUTION_ENABLED_OPTION: &str = "data-evolution.enabled"; +const FILE_INDEX_READ_ENABLED_OPTION: &str = "file-index.read.enabled"; const GLOBAL_INDEX_ENABLED_OPTION: &str = "global-index.enabled"; const GLOBAL_INDEX_SEARCH_MODE_OPTION: &str = "global-index.search-mode"; const SCALAR_INDEX_SEARCH_MODE_OPTION: &str = "scalar-index.search-mode"; @@ -656,6 +657,14 @@ impl<'a> CoreOptions<'a> { .unwrap_or(false) } + /// Whether raw data-file reads use FileIndex pruning. Default is true. + pub fn file_index_read_enabled(&self) -> bool { + self.options + .get(FILE_INDEX_READ_ENABLED_OPTION) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(true) + } + /// The declared [`TableType`], defaulting to [`TableType::Table`]. /// Fails on a value this client does not know. pub fn table_type(&self) -> crate::Result { @@ -1577,6 +1586,18 @@ mod tests { } } + #[test] + fn test_file_index_read_enabled() { + let options = HashMap::new(); + assert!(CoreOptions::new(&options).file_index_read_enabled()); + + let options = HashMap::from([( + FILE_INDEX_READ_ENABLED_OPTION.to_string(), + "false".to_string(), + )]); + assert!(!CoreOptions::new(&options).file_index_read_enabled()); + } + #[test] fn test_parquet_row_group_read_budget_options() { let options = HashMap::new(); diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 8958b3827..a6314a5ec 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -20,6 +20,8 @@ use crate::arrow::format::create_format_reader_with_budget; use crate::arrow::schema_evolution::{create_index_mapping, NULL_FIELD_INDEX}; use crate::arrow::ParquetReadBudget; use crate::deletion_vector::{DeletionVector, DeletionVectorFactory}; +use crate::file_index::evaluator::evaluate_file_index; +use crate::file_index::file_index_result::FileIndexResult; use crate::io::{FileIO, FileRead}; use crate::spec::{ is_variant_extraction_row_type, DataField, DataFileMeta, DataType, Predicate, ROW_ID_FIELD_NAME, @@ -33,6 +35,7 @@ use arrow_cast::cast; use async_stream::try_stream; use futures::StreamExt; +use roaring::RoaringBitmap; use std::ops::Range; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; @@ -112,6 +115,7 @@ pub(crate) struct DataFileReader { table_fields: Vec, read_type: Vec, predicates: Vec, + file_index_read_enabled: bool, row_filter_factory: Option>, blob_as_descriptor: bool, batch_size: Option, @@ -135,6 +139,7 @@ impl DataFileReader { table_fields, read_type, predicates, + file_index_read_enabled: false, row_filter_factory: None, blob_as_descriptor: false, batch_size: None, @@ -148,6 +153,11 @@ impl DataFileReader { self } + pub(crate) fn with_file_index_read_enabled(mut self, enabled: bool) -> Self { + self.file_index_read_enabled = enabled; + self + } + pub(crate) fn with_batch_size(mut self, batch_size: Option) -> Self { self.batch_size = batch_size; self @@ -257,12 +267,59 @@ impl DataFileReader { timing.add_file_schema_open(start.elapsed()); } - let mut stream = reader.read_single_file_stream( + let file_fields = data_fields + .as_deref() + .unwrap_or(reader.table_fields.as_slice()); + let file_index_result = if reader.file_index_read_enabled { + evaluate_file_index( + &reader.file_io, + split.bucket_path(), + &file_meta, + &reader.table_fields, + file_fields, + &reader.predicates, + ) + .await? + } else { + FileIndexResult::Remain + }; + + let split_ranges = split.row_ranges().map(|ranges| { + to_local_row_ranges( + ranges, + file_meta.first_row_id.unwrap_or(0), + file_meta.row_count, + ) + }); + let selected_ranges = match file_index_result { + FileIndexResult::Remain => split_ranges, + FileIndexResult::Skip => Some(Vec::new()), + FileIndexResult::Selection(selection) => { + match file_index_selection_to_local_ranges( + &selection, + file_meta.row_count, + )? { + Some(index_ranges) => Some(match split_ranges { + Some(split_ranges) => { + intersect_sorted_ranges(&index_ranges, &split_ranges) + } + None => index_ranges, + }), + None => split_ranges, + } + } + }; + let row_selection = merge_row_selection( + file_meta.row_count, + dv.as_deref(), + selected_ranges.as_deref(), + ); + + let mut stream = reader.read_single_file_stream_with_selection( &split, file_meta, data_fields, - dv, - split.row_ranges().map(|ranges| ranges.to_vec()), + row_selection, )?; while let Some(batch) = stream.next().await { yield batch?; @@ -339,6 +396,25 @@ impl DataFileReader { data_fields: Option>, dv: Option>, row_ranges: Option>, + ) -> crate::Result { + 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()); + self.read_single_file_stream_with_selection(split, file_meta, data_fields, row_selection) + } + + fn read_single_file_stream_with_selection( + &self, + split: &DataSplit, + file_meta: DataFileMeta, + data_fields: Option>, + row_selection: Option>, ) -> crate::Result { // Guard at the true risk site: `_ROW_ID` is materialized positionally from // each batch's row count (see `row_id_column_for_batch`), assuming the @@ -351,6 +427,9 @@ impl DataFileReader { // data-evolution readers; both strip/omit `_ROW_ID` from the read_type // they pass, so this guard does not affect them. Self::reject_row_id_with_predicates(&self.read_type, &self.predicates)?; + if row_selection.as_ref().is_some_and(Vec::is_empty) { + return Ok(futures::stream::empty().boxed()); + } let read_type = self.read_type.clone(); let table_fields = self.table_fields.clone(); @@ -440,21 +519,11 @@ impl DataFileReader { None => Box::new(file_reader), }; let is_parquet = path_to_read.to_ascii_lowercase().ends_with(".parquet"); - 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 selected_row_ids = selected_row_ids_for_read( + projects_row_id, + file_meta.first_row_id, + row_selection.as_deref(), ); - let selected_row_ids = match (file_meta.first_row_id, row_selection.as_ref()) { - (Some(first_row_id), Some(ranges)) => { - Some(expand_local_selected_row_ids(first_row_id, ranges)) - } - _ => None, - }; let mut row_id_cursor = file_meta.first_row_id.unwrap_or(0); let mut row_id_offset = 0usize; @@ -891,24 +960,26 @@ fn is_row_file(file_meta: &DataFileMeta) -> bool { .is_some_and(|path| path.to_ascii_lowercase().ends_with(".row")) } -/// Convert absolute RowRanges to file-local 0-based ranges. +/// Convert absolute RowRanges to normalized file-local 0-based ranges. fn to_local_row_ranges( row_ranges: &[RowRange], first_row_id: i64, row_count: i64, ) -> Vec { let file_end = first_row_id + row_count - 1; - row_ranges - .iter() - .filter_map(|r| { - if r.to() < first_row_id || r.from() > file_end { - return None; - } - let local_from = (r.from() - first_row_id).max(0); - let local_to = (r.to() - first_row_id).min(row_count - 1); - Some(RowRange::new(local_from, local_to)) - }) - .collect() + crate::table::merge_row_ranges( + row_ranges + .iter() + .filter_map(|r| { + if r.to() < first_row_id || r.from() > file_end { + return None; + } + let local_from = (r.from() - first_row_id).max(0); + let local_to = (r.to() - first_row_id).min(row_count - 1); + Some(RowRange::new(local_from, local_to)) + }) + .collect(), + ) } /// Coalesce sorted, de-duplicated 0-based physical positions into contiguous @@ -935,6 +1006,39 @@ fn coalesce_positions_to_local_ranges(sorted_positions: &[i64]) -> Vec ranges } +const MAX_FILE_INDEX_ROW_RANGES: usize = 65_536; + +/// Convert a bitmap into contiguous ranges without visiting every selected row. +/// `None` means the bitmap is too fragmented to materialize safely and callers +/// must preserve other restrictions and rely on the residual predicate. +fn file_index_selection_to_local_ranges( + selection: &RoaringBitmap, + row_count: i64, +) -> crate::Result>> { + if let Some(position) = selection.max() { + if i64::from(position) >= row_count { + return Err(Error::FileIndexFormatInvalid { + message: format!( + "FileIndex selected row position {position} outside data file row count {row_count}" + ), + }); + } + } + + let mut ranges = Vec::new(); + let mut positions = selection.iter(); + while let Some(range) = positions.next_range() { + if ranges.len() == MAX_FILE_INDEX_ROW_RANGES { + return Ok(None); + } + ranges.push(RowRange::new( + i64::from(*range.start()), + i64::from(*range.end()), + )); + } + Ok(Some(ranges)) +} + /// Merge DV and row_ranges into a unified list of 0-based inclusive RowRanges. /// Returns `None` if no filtering is needed (no DV and no ranges). /// @@ -1053,6 +1157,22 @@ fn expand_local_selected_row_ids(first_row_id: i64, local_ranges: &[RowRange]) - ids } +fn selected_row_ids_for_read( + projects_row_id: bool, + first_row_id: Option, + row_selection: Option<&[RowRange]>, +) -> Option> { + if !projects_row_id { + return None; + } + match (first_row_id, row_selection) { + (Some(first_row_id), Some(ranges)) => { + Some(expand_local_selected_row_ids(first_row_id, ranges)) + } + _ => None, + } +} + fn row_id_column_for_batch( first_row_id: Option, num_rows: usize, @@ -1486,10 +1606,17 @@ mod row_tests { mod tests { use super::*; use crate::arrow::build_target_arrow_schema; + use crate::common::Options; + use crate::file_index::file_index_result::FileIndexResult; + use crate::file_index::file_indexer_factory::{ + FileIndexerFactory, BITMAP_INDEX, BLOOM_FILTER_INDEX, + }; + use crate::file_index::write_column_indexes; use crate::io::FileIOBuilder; use crate::spec::stats::BinaryTableStats; use crate::spec::{ - ArrayType, DataFileMeta, DataType, Datum, IntType, Predicate, PredicateBuilder, VarCharType, + ArrayType, BigIntType, DataFileMeta, DataType, Datum, IntType, Predicate, PredicateBuilder, + PredicateOperator, Schema, SchemaChange, TableSchema, VarCharType, }; use crate::table::source::{DataSplitBuilder, DeletionFile}; use arrow_array::{Int32Array, StringArray}; @@ -1544,6 +1671,56 @@ mod tests { ); } + #[test] + fn file_index_selection_coalesces_positions_and_rejects_out_of_range_values() { + let selection = [0_u32, 1, 3].into_iter().collect(); + assert_eq!( + file_index_selection_to_local_ranges(&selection, 4).unwrap(), + Some(vec![RowRange::new(0, 1), RowRange::new(3, 3)]) + ); + + let out_of_range = [4_u32].into_iter().collect(); + assert!(matches!( + file_index_selection_to_local_ranges(&out_of_range, 4), + Err(Error::FileIndexFormatInvalid { .. }) + )); + } + + #[test] + fn dense_file_index_selection_stays_compact_and_fragmented_selection_falls_back() { + let mut dense = RoaringBitmap::new(); + dense.insert_range(0..=10_000_000); + assert_eq!( + file_index_selection_to_local_ranges(&dense, 10_000_001).unwrap(), + Some(vec![RowRange::new(0, 10_000_000)]) + ); + + let fragmented = (0..=MAX_FILE_INDEX_ROW_RANGES as u32) + .map(|position| position * 2) + .collect(); + assert_eq!( + file_index_selection_to_local_ranges(&fragmented, 200_000).unwrap(), + None + ); + } + + #[test] + fn selected_row_ids_are_built_only_when_projected() { + let huge_selection = [RowRange::new(0, i64::from(i32::MAX))]; + assert_eq!( + selected_row_ids_for_read(false, Some(10), Some(&huge_selection)), + None + ); + assert_eq!( + selected_row_ids_for_read( + true, + Some(10), + Some(&[RowRange::new(1, 2), RowRange::new(4, 4)]), + ), + Some(vec![11, 12, 14]) + ); + } + #[test] fn test_accessors_expose_read_type_and_row_filtering_predicate() { use crate::spec::{DataField, DataType, IntType}; @@ -1696,6 +1873,34 @@ mod tests { Bytes::from(writer.output().data.to_vec()) } + async fn file_index_bytes( + path: &str, + column: &str, + identifier: &str, + data_type: DataType, + options: &Options, + values: &[Datum], + ) -> Bytes { + let mut writer = FileIndexerFactory::create_writer(identifier, data_type, options).unwrap(); + for value in values { + writer.write(Some(value)).unwrap(); + } + let indexes = std::collections::HashMap::from([( + column.to_string(), + std::collections::HashMap::from([( + identifier.to_string(), + Some(writer.serialized_bytes().unwrap()), + )]), + )]); + write_column_indexes(path, indexes) + .await + .unwrap() + .to_input_file() + .read() + .await + .unwrap() + } + #[tokio::test] async fn test_mosaic_physical_missing_column_is_null_filled() { let physical_fields = vec![ @@ -1895,6 +2100,406 @@ mod tests { .collect() } + #[tokio::test] + async fn test_file_index_bitmap_skip_avoids_opening_data_file_and_disabled_falls_back() { + let fields = pk_fields(); + let index = file_index_bytes( + "memory:/file_index_skip_source", + "id", + BITMAP_INDEX, + fields[0].data_type().clone(), + &Options::new(), + &[Datum::Int(1), Datum::Int(2)], + ) + .await; + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/file_index_skip"; + let bucket_path = format!("{table_path}/bucket-0"); + let mut file = data_file("missing.mosaic", 1, 2, 0); + file.embedded_index = Some(index.to_vec()); + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![file]) + .build() + .unwrap(); + let predicate = PredicateBuilder::new(&fields) + .equal("id", Datum::Int(99)) + .unwrap(); + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + + let enabled = DataFileReader::new( + file_io.clone(), + schema_manager.clone(), + 0, + fields.clone(), + fields.clone(), + vec![predicate.clone()], + ) + .with_file_index_read_enabled(true) + .read(std::slice::from_ref(&split)) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(enabled.is_empty()); + + let disabled = DataFileReader::new( + file_io, + schema_manager, + 0, + fields.clone(), + fields, + vec![predicate], + ) + .with_file_index_read_enabled(false) + .read(&[split]) + .unwrap() + .try_collect::>() + .await; + assert!( + disabled.is_err(), + "disabled reads must preserve the data path" + ); + } + + #[tokio::test] + async fn test_file_index_bitmap_selection_intersects_unordered_split_ranges_and_deletion_vector( + ) { + let fields = pk_fields(); + let data = write_mosaic(&pk_batch( + vec![1, 2, 3, 4, 5, 6], + vec!["a", "b", "c", "d", "e", "f"], + )); + let index = file_index_bytes( + "memory:/file_index_selection_source", + "id", + BITMAP_INDEX, + fields[0].data_type().clone(), + &Options::new(), + &[ + Datum::Int(1), + Datum::Int(2), + Datum::Int(3), + Datum::Int(4), + Datum::Int(5), + Datum::Int(6), + ], + ) + .await; + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/file_index_selection"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_name = "part-0.mosaic"; + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(data.clone()) + .await + .unwrap(); + let sidecar_name = format!("{file_name}.index"); + file_io + .new_output(&format!("{bucket_path}/{sidecar_name}")) + .unwrap() + .write(index) + .await + .unwrap(); + let dv = write_deletion_file(&file_io, &format!("{table_path}/index/dv-0"), &[3]).await; + let mut file = data_file(file_name, data.len() as i64, 6, 0); + file.extra_files = vec![sidecar_name]; + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![file]) + .with_data_deletion_files(vec![Some(dv)]) + .with_row_ranges(vec![ + RowRange::new(4, 5), + RowRange::new(1, 3), + RowRange::new(2, 4), + ]) + .build() + .unwrap(); + let predicate = PredicateBuilder::new(&fields) + .is_in( + "id", + vec![Datum::Int(2), Datum::Int(3), Datum::Int(4), Datum::Int(5)], + ) + .unwrap(); + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + + for enabled in [true, false] { + let batches = DataFileReader::new( + file_io.clone(), + schema_manager.clone(), + 0, + fields.clone(), + fields.clone(), + vec![predicate.clone()], + ) + .with_file_index_read_enabled(enabled) + .read(std::slice::from_ref(&split)) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_ids(&batches), vec![2, 3, 5], "enabled={enabled}"); + } + } + + #[tokio::test] + async fn test_file_index_bloom_false_positive_is_removed_by_residual_filter() { + let fields = pk_fields(); + let data_type = fields[0].data_type().clone(); + let mut bloom_options = Options::new(); + bloom_options.set("items", "1"); + bloom_options.set("fpp", "0.99"); + let mut writer = FileIndexerFactory::create_writer( + BLOOM_FILTER_INDEX, + data_type.clone(), + &bloom_options, + ) + .unwrap(); + writer.write(Some(&Datum::Int(1))).unwrap(); + let payload = writer.serialized_bytes().unwrap(); + let bloom_reader = FileIndexerFactory::create_reader( + BLOOM_FILTER_INDEX, + data_type.clone(), + payload.clone(), + ) + .unwrap(); + let false_positive = (2..10_000) + .find(|candidate| { + bloom_reader.evaluate( + "id", + 0, + &data_type, + PredicateOperator::Eq, + &[Datum::Int(*candidate)], + ) == FileIndexResult::Remain + }) + .expect("high-FPP Bloom filter should have a false positive"); + let indexes = std::collections::HashMap::from([( + "id".to_string(), + std::collections::HashMap::from([(BLOOM_FILTER_INDEX.to_string(), Some(payload))]), + )]); + let index = write_column_indexes("memory:/file_index_bloom_source", indexes) + .await + .unwrap() + .to_input_file() + .read() + .await + .unwrap(); + + let data = write_mosaic(&pk_batch(vec![1], vec!["a"])); + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/file_index_bloom"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_name = "part-0.mosaic"; + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(data.clone()) + .await + .unwrap(); + let mut file = data_file(file_name, data.len() as i64, 1, 0); + file.embedded_index = Some(index.to_vec()); + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![file]) + .build() + .unwrap(); + let predicate = PredicateBuilder::new(&fields) + .equal("id", Datum::Int(false_positive)) + .unwrap(); + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let batches = DataFileReader::new( + file_io, + schema_manager, + 0, + fields.clone(), + fields, + vec![predicate], + ) + .with_file_index_read_enabled(true) + .read(&[split]) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert!(collect_ids(&batches).is_empty()); + } + + #[tokio::test] + async fn test_file_index_uses_schema_evolved_file_fields_and_remapped_predicate() { + let old_schema = TableSchema::new( + 0, + &Schema::builder() + .column("old_id", DataType::Int(IntType::new())) + .build() + .unwrap(), + ); + let current_schema = old_schema + .apply_changes(vec![ + SchemaChange::rename_column("old_id".to_string(), "new_id".to_string()), + SchemaChange::update_column_type( + "new_id".to_string(), + DataType::BigInt(BigIntType::new()), + ), + ]) + .unwrap(); + let index = file_index_bytes( + "memory:/file_index_schema_evolution_source", + "old_id", + BITMAP_INDEX, + old_schema.fields()[0].data_type().clone(), + &Options::new(), + &[Datum::Int(1), Datum::Int(2)], + ) + .await; + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/file_index_schema_evolution"; + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let schema_path = schema_manager.schema_path(old_schema.id()); + let schema_dir = schema_path.rsplit_once('/').unwrap().0; + file_io.mkdirs(schema_dir).await.unwrap(); + file_io + .new_output(&schema_path) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&old_schema).unwrap())) + .await + .unwrap(); + let mut file = data_file("missing.mosaic", 1, 2, old_schema.id()); + file.embedded_index = Some(index.to_vec()); + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(format!("{table_path}/bucket-0")) + .with_total_buckets(1) + .with_data_files(vec![file]) + .build() + .unwrap(); + let predicate = PredicateBuilder::new(current_schema.fields()) + .equal("new_id", Datum::Long(99)) + .unwrap(); + + let batches = DataFileReader::new( + file_io, + schema_manager, + current_schema.id(), + current_schema.fields().to_vec(), + current_schema.fields().to_vec(), + vec![predicate], + ) + .with_file_index_read_enabled(true) + .read(&[split]) + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert!(batches.is_empty()); + } + + #[tokio::test] + async fn test_file_index_nested_not_with_added_column_falls_back() { + let old_schema = TableSchema::new( + 0, + &Schema::builder() + .column("id", DataType::Int(IntType::new())) + .build() + .unwrap(), + ); + let current_schema = old_schema + .apply_changes(vec![SchemaChange::add_column( + "added".to_string(), + DataType::Int(IntType::new()), + )]) + .unwrap(); + let batch = RecordBatch::try_new( + build_target_arrow_schema(old_schema.fields()).unwrap(), + vec![Arc::new(Int32Array::from(vec![1, 2]))], + ) + .unwrap(); + let data = write_mosaic(&batch); + let index = file_index_bytes( + "memory:/file_index_nested_not_source", + "id", + BITMAP_INDEX, + old_schema.fields()[0].data_type().clone(), + &Options::new(), + &[Datum::Int(1), Datum::Int(2)], + ) + .await; + + let file_io = FileIOBuilder::new("memory").build().unwrap(); + let table_path = "memory:/file_index_nested_not"; + let bucket_path = format!("{table_path}/bucket-0"); + let file_name = "part-0.mosaic"; + file_io + .new_output(&format!("{bucket_path}/{file_name}")) + .unwrap() + .write(data.clone()) + .await + .unwrap(); + let schema_manager = SchemaManager::new(file_io.clone(), table_path.to_string()); + let schema_path = schema_manager.schema_path(old_schema.id()); + let schema_dir = schema_path.rsplit_once('/').unwrap().0; + file_io.mkdirs(schema_dir).await.unwrap(); + file_io + .new_output(&schema_path) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&old_schema).unwrap())) + .await + .unwrap(); + + let mut file = data_file(file_name, data.len() as i64, 2, old_schema.id()); + file.embedded_index = Some(index.to_vec()); + let split = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(crate::spec::BinaryRow::new(0)) + .with_bucket(0) + .with_bucket_path(bucket_path) + .with_total_buckets(1) + .with_data_files(vec![file]) + .build() + .unwrap(); + let builder = PredicateBuilder::new(current_schema.fields()); + let predicate = Predicate::negate(Predicate::and(vec![ + Predicate::negate(builder.equal("id", Datum::Int(1)).unwrap()), + builder.equal("added", Datum::Int(2)).unwrap(), + ])); + + for enabled in [true, false] { + let batches = DataFileReader::new( + file_io.clone(), + schema_manager.clone(), + current_schema.id(), + current_schema.fields().to_vec(), + vec![current_schema.fields()[0].clone()], + vec![predicate.clone()], + ) + .with_file_index_read_enabled(enabled) + .read(std::slice::from_ref(&split)) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(collect_ids(&batches), vec![1, 2], "enabled={enabled}"); + } + } + /// Deletion vectors are applied format-agnostically by `DataFileReader`; verify a /// Mosaic file honors deleted rows end to end. #[tokio::test] diff --git a/crates/paimon/src/table/table_read.rs b/crates/paimon/src/table/table_read.rs index 99c01f7d9..acaebeeb1 100644 --- a/crates/paimon/src/table/table_read.rs +++ b/crates/paimon/src/table/table_read.rs @@ -382,6 +382,7 @@ impl<'a> PaimonTableRead<'a> { has_value_kind: bool, ) -> crate::Result { plan.validate()?; + let core_options = self.table.schema().core_options(); let data_splits = plan.data_splits(); let user_read_type = self.read_type.clone(); let include_sequence = audit_sequence_number_enabled(self.table); @@ -414,7 +415,8 @@ impl<'a> PaimonTableRead<'a> { read_type, self.data_predicates.clone(), ) - .with_batch_size(Some(self.table.schema().core_options().read_batch_size()?)) + .with_file_index_read_enabled(core_options.file_index_read_enabled()) + .with_batch_size(Some(core_options.read_batch_size()?)) .with_parquet_read_budget(Some(self.parquet_read_budget()?)); let raw_stream = reader.read(&data_splits)?; @@ -883,6 +885,7 @@ impl<'a> PaimonTableRead<'a> { } fn new_data_file_reader(&self) -> crate::Result { + let core_options = self.table.schema().core_options(); let mut reader = DataFileReader::new( self.table.file_io.clone(), self.table.schema_manager().clone(), @@ -891,7 +894,8 @@ impl<'a> PaimonTableRead<'a> { self.read_type().to_vec(), self.data_predicates.clone(), ) - .with_batch_size(Some(self.table.schema().core_options().read_batch_size()?)) + .with_file_index_read_enabled(core_options.file_index_read_enabled()) + .with_batch_size(Some(core_options.read_batch_size()?)) .with_parquet_read_budget(Some(self.parquet_read_budget()?)) .with_read_timing(self.data_file_read_timing.clone()); // The engine decoder filter is safe only on the plain append/raw path. @@ -1481,11 +1485,17 @@ fn pk_split_needs_merge(split: &DataSplit, dv_enabled: bool) -> bool { mod tests { use super::*; use crate::catalog::Identifier; + use crate::common::Options; + use crate::file_index::file_indexer_factory::{FileIndexerFactory, BITMAP_INDEX}; + use crate::file_index::write_column_indexes; use crate::io::FileIOBuilder; use crate::spec::stats::BinaryTableStats; - use crate::spec::{BinaryRow, DataFileMeta, DataType, IntType, Schema, TableSchema}; + use crate::spec::{ + BinaryRow, DataFileMeta, DataType, Datum, IntType, PredicateBuilder, Schema, TableSchema, + }; use crate::table::query_auth_table; use crate::table::source::DataSplitBuilder; + use futures::TryStreamExt; fn file(name: &str, level: i32, delete_row_count: Option) -> DataFileMeta { DataFileMeta { @@ -1542,6 +1552,101 @@ mod tests { ) } + async fn embedded_bitmap_index() -> Vec { + let mut writer = FileIndexerFactory::create_writer( + BITMAP_INDEX, + DataType::Int(IntType::new()), + &Options::new(), + ) + .unwrap(); + writer.write(Some(&Datum::Int(1))).unwrap(); + let indexes = std::collections::HashMap::from([( + "id".to_string(), + std::collections::HashMap::from([( + BITMAP_INDEX.to_string(), + Some(writer.serialized_bytes().unwrap()), + )]), + )]); + write_column_indexes("memory:/table_read_file_index_source", indexes) + .await + .unwrap() + .to_input_file() + .read() + .await + .unwrap() + .to_vec() + } + + fn file_index_table(path: &str, enabled: Option) -> Table { + let mut builder = Schema::builder().column("id", DataType::Int(IntType::new())); + if let Some(enabled) = enabled { + builder = builder.option("file-index.read.enabled", enabled.to_string()); + } + Table::new( + FileIOBuilder::new("memory").build().unwrap(), + Identifier::new("default", "file_index_t"), + path.to_string(), + TableSchema::new(0, &builder.build().unwrap()), + None, + ) + } + + #[tokio::test] + async fn test_raw_table_read_paths_honor_file_index_read_option() { + let mut indexed_file = file("missing.mosaic", 5, Some(0)); + indexed_file.row_count = 1; + indexed_file.embedded_index = Some(embedded_bitmap_index().await); + let split = split(vec![indexed_file], true); + let table = file_index_table("memory:/table_read_file_index", None); + let fields = table.schema().fields().to_vec(); + let predicate = PredicateBuilder::new(&fields) + .equal("id", Datum::Int(99)) + .unwrap(); + let read = TableRead::new(&table, fields, vec![predicate]); + + let normal = read + .to_arrow(std::slice::from_ref(&split)) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(normal.is_empty()); + + let plan = IncrementalPlan::new( + IncrementalScanMode::Delta, + vec![IncrementalSplit::Data(split.clone())], + ); + let incremental = read + .to_incremental_arrow(&plan) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(incremental.is_empty()); + let audit = read + .to_audit_log_arrow(&plan) + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert!(audit.is_empty()); + + let disabled_table = + file_index_table("memory:/table_read_file_index_disabled", Some(false)); + let disabled_fields = disabled_table.schema().fields().to_vec(); + let disabled_predicate = PredicateBuilder::new(&disabled_fields) + .equal("id", Datum::Int(99)) + .unwrap(); + let disabled_read = + TableRead::new(&disabled_table, disabled_fields, vec![disabled_predicate]); + let disabled = disabled_read + .to_arrow(&[split]) + .unwrap() + .try_collect::>() + .await; + assert!(disabled.is_err()); + } + #[test] fn test_pk_split_needs_merge_routing() { // Raw convertible with known delete counts: raw read.