diff --git a/crates/paimon/src/file_index/file_index_format.rs b/crates/paimon/src/file_index/file_index_format.rs index 1e24319bf..484739468 100644 --- a/crates/paimon/src/file_index/file_index_format.rs +++ b/crates/paimon/src/file_index/file_index_format.rs @@ -237,6 +237,14 @@ pub async fn write_column_indexes( ) -> crate::Result { let file_io = FileIO::from_path(path)?.build()?; let output = file_io.new_output(path)?; + output.write(serialize_column_indexes(indexes)?).await?; + Ok(output) +} + +/// Serialize the complete container, including its header, for either storage form. +pub(crate) fn serialize_column_indexes( + indexes: HashMap>>, +) -> crate::Result { let mut body_info: HashMap> = HashMap::new(); let mut total_data_size = 0usize; @@ -304,11 +312,8 @@ pub async fn write_column_indexes( head_buffer.put_i32(0); debug_assert_eq!(head_buffer.len(), head_length); - let mut writer = output.writer().await?; - writer.write(head_buffer.freeze()).await?; - writer.write(body.freeze()).await?; - writer.close().await?; - Ok(output) + head_buffer.extend_from_slice(&body); + Ok(head_buffer.freeze()) } fn calculate_head_length( diff --git a/crates/paimon/src/file_index/file_index_writer.rs b/crates/paimon/src/file_index/file_index_writer.rs index 3fce9bb81..9ad942ce2 100644 --- a/crates/paimon/src/file_index/file_index_writer.rs +++ b/crates/paimon/src/file_index/file_index_writer.rs @@ -21,7 +21,7 @@ use crate::spec::Datum; use crate::Result; /// Writes one concrete file index payload. -pub(crate) trait FileIndexWriter { +pub(crate) trait FileIndexWriter: Send { /// Adds one row to the index. `None` represents a null value. fn write(&mut self, datum: Option<&Datum>) -> Result<()>; diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index d3114fd4c..1e32a2fb1 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -674,6 +674,16 @@ impl<'a> CoreOptions<'a> { .unwrap_or(false) } + /// Maximum complete FileIndex size stored in the manifest. Default is 500 bytes. + pub(crate) fn file_index_in_manifest_threshold(&self) -> crate::Result { + match self.options.get("file-index.in-manifest-threshold") { + None => Ok(500), + Some(raw) => parse_memory_size(raw).ok_or_else(|| crate::Error::ConfigInvalid { + message: format!("Invalid file-index.in-manifest-threshold: {raw}"), + }), + } + } + /// Whether raw data-file reads use FileIndex pruning. Default is true. pub fn file_index_read_enabled(&self) -> bool { self.options diff --git a/crates/paimon/src/table/data_file_index_writer.rs b/crates/paimon/src/table/data_file_index_writer.rs new file mode 100644 index 000000000..14444b407 --- /dev/null +++ b/crates/paimon/src/table/data_file_index_writer.rs @@ -0,0 +1,232 @@ +// 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::{BTreeMap, HashMap}; + +use arrow_array::RecordBatch; +use bytes::Bytes; + +use crate::common::Options; +use crate::file_index::file_index_writer::FileIndexWriter; +use crate::file_index::file_indexer_factory::FileIndexerFactory; +use crate::file_index::serialize_column_indexes; +use crate::spec::{extract_datum_from_arrow, CoreOptions, DataField}; +use crate::{Error, Result}; + +#[derive(Clone)] +struct IndexColumnOptions { + field: DataField, + position: usize, + indexes: BTreeMap, +} + +/// Validated top-level column indexes, enabled explicitly by ordinary append writes. +#[derive(Clone)] +pub(super) struct FileIndexOptions { + columns: Vec, + pub(super) in_manifest_threshold: i64, +} + +impl FileIndexOptions { + pub(super) fn parse( + options: &HashMap, + fields: &[DataField], + ) -> Result> { + let mut columns: BTreeMap> = BTreeMap::new(); + for (key, value) in options { + let Some(identifier) = key + .strip_prefix("file-index.") + .and_then(|key| key.strip_suffix(".columns")) + else { + continue; + }; + if !FileIndexerFactory::is_supported(identifier) { + return Err(Error::Unsupported { + message: format!("Unsupported file index in {key}: {identifier}"), + }); + } + for column in value.split(',').map(str::trim) { + if column.is_empty() { + return Err(Error::ConfigInvalid { + message: format!("Empty column in {key}"), + }); + } + columns + .entry(column.to_string()) + .or_default() + .entry(identifier.to_string()) + .or_default(); + } + } + + for (key, value) in options { + let Some(suffix) = key.strip_prefix("file-index.") else { + continue; + }; + if suffix == "read.enabled" + || suffix == "in-manifest-threshold" + || suffix.ends_with(".columns") + { + continue; + } + let parts = suffix.split_once('.').and_then(|(identifier, rest)| { + rest.rsplit_once('.') + .map(|(column, option)| (identifier, column, option)) + }); + let Some((identifier, column, option)) = parts else { + return Err(Error::ConfigInvalid { + message: format!("Invalid file index option: {key}"), + }); + }; + let Some(index_options) = columns.get_mut(column).and_then(|c| c.get_mut(identifier)) + else { + return Err(Error::ConfigInvalid { + message: format!( + "{key} requires column '{column}' in file-index.{identifier}.columns" + ), + }); + }; + if !matches!( + (identifier, option), + ("bitmap", "version" | "index-block-size") | ("bloom-filter", "items" | "fpp") + ) { + return Err(Error::ConfigInvalid { + message: format!("Unknown file index option: {key}"), + }); + } + index_options.set(option, value); + } + + let in_manifest_threshold = CoreOptions::new(options).file_index_in_manifest_threshold()?; + let columns = columns + .into_iter() + .map(|(name, indexes)| { + let (position, field) = fields + .iter() + .enumerate() + .find(|(_, field)| field.name() == name) + .ok_or_else(|| Error::ConfigInvalid { + message: format!( + "File index column '{name}' does not exist as a top-level field" + ), + })?; + for (identifier, options) in &indexes { + FileIndexerFactory::create_writer( + identifier, + field.data_type().clone(), + options, + )?; + } + Ok(IndexColumnOptions { + field: field.clone(), + position, + indexes, + }) + }) + .collect::>>()?; + Ok((!columns.is_empty()).then_some(Self { + columns, + in_manifest_threshold, + })) + } + + pub(super) fn create_writer(&self) -> Result { + let columns = self + .columns + .iter() + .map(|column| { + let writers = column + .indexes + .iter() + .map(|(identifier, options)| { + Ok(( + identifier.clone(), + FileIndexerFactory::create_writer( + identifier, + column.field.data_type().clone(), + options, + )?, + )) + }) + .collect::>>()?; + Ok(IndexColumn { + field: column.field.clone(), + position: column.position, + writers, + }) + }) + .collect::>>()?; + Ok(DataFileIndexWriter { columns }) + } +} + +struct IndexColumn { + field: DataField, + position: usize, + writers: Vec<(String, Box)>, +} + +pub(super) struct DataFileIndexWriter { + columns: Vec, +} + +impl DataFileIndexWriter { + pub(super) fn write(&mut self, batch: &RecordBatch) -> Result<()> { + for column in &mut self.columns { + for row in 0..batch.num_rows() { + let datum = extract_datum_from_arrow( + batch, + row, + column.position, + column.field.data_type(), + )?; + for (_, writer) in &mut column.writers { + writer.write(datum.as_ref())?; + } + } + } + Ok(()) + } + + pub(super) fn serialize(mut self) -> Result { + let indexes = self + .columns + .iter_mut() + .map(|column| { + let indexes = column + .writers + .iter_mut() + .map(|(identifier, writer)| { + Ok(( + identifier.clone(), + if writer.empty() { + None + } else { + Some(writer.serialized_bytes()?) + }, + )) + }) + .collect::>>()?; + Ok((column.field.name().to_string(), indexes)) + }) + .collect::>>()?; + serialize_column_indexes(indexes) + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/paimon/src/table/data_file_index_writer/tests.rs b/crates/paimon/src/table/data_file_index_writer/tests.rs new file mode 100644 index 000000000..65155a744 --- /dev/null +++ b/crates/paimon/src/table/data_file_index_writer/tests.rs @@ -0,0 +1,773 @@ +// 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 std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use arrow_array::Int32Array; +use arrow_schema::{DataType as ArrowType, Field, Schema as ArrowSchema}; +use futures::TryStreamExt; +use opendal::{services::MemoryConfig, Operator}; + +use crate::catalog::Identifier; +use crate::file_index::evaluator::evaluate_file_index; +use crate::file_index::file_index_result::FileIndexResult; +use crate::io::{FileIO, FileIOBuilder, FileIOProvider}; +use crate::spec::{ + BooleanType, DataFileMeta, Datum, IntType, Predicate, PredicateBuilder, Schema, TableSchema, +}; +use crate::table::{DataSplitBuilder, SchemaManager, Table}; + +fn schema(options: &[(&str, &str)]) -> Schema { + let mut builder = Schema::builder() + .column("id", crate::spec::DataType::Int(IntType::new())) + .column("value", crate::spec::DataType::Int(IntType::new())); + for (key, value) in options { + builder = builder.option(*key, *value); + } + builder.build().unwrap() +} + +async fn table(io: FileIO, schema: Schema) -> Table { + let path = format!("memory:/append-index-test-{}", uuid::Uuid::new_v4()); + let schema = TableSchema::new(0, &schema); + for dir in ["schema", "snapshot", "manifest"] { + io.mkdirs(&format!("{path}/{dir}")).await.unwrap(); + } + io.new_output(&format!("{path}/schema/schema-0")) + .unwrap() + .write(Bytes::from(serde_json::to_vec(&schema).unwrap())) + .await + .unwrap(); + Table::new( + io, + Identifier::new("default", "indexed_append"), + path, + schema, + None, + ) +} + +fn memory_io() -> FileIO { + FileIOBuilder::new("memory").build().unwrap() +} + +fn batch(ids: Vec>, values: Vec>) -> RecordBatch { + RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + Field::new("id", ArrowType::Int32, true), + Field::new("value", ArrowType::Int32, true), + ])), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap() +} + +fn rows(batches: &[RecordBatch]) -> Vec<(Option, Option)> { + let mut result: Vec<_> = batches + .iter() + .flat_map(|batch| { + let ids = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + ids.iter().zip(values.iter()) + }) + .collect(); + result.sort_unstable(); + result +} + +async fn query( + table: &Table, + enabled: bool, + predicate: Option, +) -> Vec<(Option, Option)> { + let table = table.copy_with_options(HashMap::from([( + "file-index.read.enabled".to_string(), + enabled.to_string(), + )])); + let mut builder = table.new_read_builder(); + if let Some(predicate) = predicate { + builder.with_filter(predicate); + } + let plan = builder.new_scan().plan().await.unwrap(); + let batches: Vec<_> = builder + .new_read() + .unwrap() + .to_arrow(plan.splits()) + .unwrap() + .try_collect() + .await + .unwrap(); + rows(&batches) +} + +async fn evaluate( + table: &Table, + bucket_path: &str, + file: &DataFileMeta, + predicate: Predicate, +) -> FileIndexResult { + evaluate_file_index( + table.file_io(), + bucket_path, + file, + table.schema().fields(), + table.schema().fields(), + &[predicate], + ) + .await + .unwrap() +} + +#[tokio::test] +async fn test_file_index_append_commit_reload_and_rolling() { + for identifier in ["bitmap", "bloom-filter", "both"] { + for rolling in [false, true] { + for threshold in ["0 B", "1 MB"] { + let mut options = vec![ + ("target-file-size", if rolling { "1 B" } else { "128 MB" }), + ("file-index.in-manifest-threshold", threshold), + ("file-index.read.enabled", "false"), + ]; + if identifier != "bloom-filter" { + options.push(("file-index.bitmap.columns", " id, value, id ")); + } + if identifier != "bitmap" { + options.extend([ + ("file-index.bloom-filter.columns", "id"), + ("file-index.bloom-filter.id.items", "10"), + ]); + } + let table = table(memory_io(), schema(&options)).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .write_arrow_batch(&batch(vec![], vec![])) + .await + .unwrap(); + assert!(writer.prepare_commit().await.unwrap().is_empty()); + let sliced = batch( + vec![Some(99), Some(1), None, Some(99)], + vec![Some(99), Some(10), Some(20), Some(99)], + ) + .slice(1, 2); + writer.write_arrow_batch(&sliced).await.unwrap(); + writer + .write_arrow_batch(&batch(vec![Some(3)], vec![None])) + .await + .unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + let expected_files = if rolling { 2 } else { 1 }; + let files: Vec<_> = messages.iter().flat_map(|m| &m.new_files).collect(); + assert_eq!(files.len(), expected_files); + assert_eq!(files.iter().map(|f| f.row_count).sum::(), 3); + for file in files { + assert_eq!(file.embedded_index.is_some(), threshold != "0 B"); + assert_eq!(file.extra_files.len(), usize::from(threshold == "0 B")); + } + builder.new_commit().commit(messages).await.unwrap(); + let persisted = + SchemaManager::new(table.file_io().clone(), table.location().to_string()) + .latest() + .await + .unwrap() + .unwrap(); + let reloaded = Table::new( + table.file_io().clone(), + Identifier::new("default", "reloaded"), + table.location().to_string(), + persisted.as_ref().clone(), + None, + ); + let read_builder = reloaded.new_read_builder(); + let plan = read_builder.new_scan().plan().await.unwrap(); + let predicates = PredicateBuilder::new(reloaded.schema().fields()); + for split in plan.splits() { + for file in split.data_files() { + assert_eq!(file.embedded_index.is_some(), threshold != "0 B"); + if identifier != "bloom-filter" { + let expected = FileIndexResult::Selection( + if file.row_count == 1 { vec![] } else { vec![1] } + .into_iter() + .collect(), + ); + assert_eq!( + evaluate( + &reloaded, + split.bucket_path(), + file, + predicates.is_null("id").unwrap() + ) + .await, + expected + ); + let expected = FileIndexResult::Selection( + if rolling && file.row_count == 2 { + vec![] + } else { + vec![if rolling { 0 } else { 2 }] + } + .into_iter() + .collect(), + ); + assert_eq!( + evaluate( + &reloaded, + split.bucket_path(), + file, + predicates.is_null("value").unwrap() + ) + .await, + expected + ); + } + } + } + for enabled in [false, true] { + assert_eq!( + query(&reloaded, enabled, None).await, + vec![(None, Some(20)), (Some(1), Some(10)), (Some(3), None)] + ); + assert_eq!( + query( + &reloaded, + enabled, + Some(predicates.equal("id", Datum::Int(3)).unwrap()) + ) + .await, + vec![(Some(3), None)] + ); + assert_eq!( + query(&reloaded, enabled, Some(predicates.is_null("id").unwrap())).await, + vec![(None, Some(20))] + ); + } + } + } + } +} + +#[tokio::test] +async fn test_file_index_unconfigured_and_read_flag_only() { + for options in [vec![], vec![("file-index.read.enabled", "true")]] { + let table = table(memory_io(), schema(&options)).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .write_arrow_batch(&batch(vec![Some(1)], vec![None])) + .await + .unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert!(file.embedded_index.is_none()); + assert!(file.extra_files.is_empty()); + builder.new_commit().commit(messages).await.unwrap(); + assert_eq!(query(&table, true, None).await, vec![(Some(1), None)]); + } +} + +#[tokio::test] +async fn test_file_index_threshold_boundary_and_abort() { + let data = batch(vec![Some(1), None, Some(3)], vec![None, None, None]); + let table_schema = schema(&[("file-index.bitmap.columns", "id")]); + let config = FileIndexOptions::parse(table_schema.options(), table_schema.fields()) + .unwrap() + .unwrap(); + assert_eq!(config.in_manifest_threshold, 500); + let mut index = config.create_writer().unwrap(); + index.write(&data).unwrap(); + let size = index.serialize().unwrap().len(); + for threshold in [size - 1, size, size + 1] { + let raw = format!("{threshold} B"); + let table = table( + memory_io(), + schema(&[ + ("file-index.bitmap.columns", "id"), + ("file-index.in-manifest-threshold", &raw), + ]), + ) + .await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer.write_arrow_batch(&data).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!(file.embedded_index.is_some(), size <= threshold); + let bucket_path = format!("{}/bucket-0", table.location()); + let paths = file.collect_files(&bucket_path); + for path in &paths { + assert!(table.file_io().exists(path).await.unwrap()); + } + let bytes = if let Some(bytes) = &file.embedded_index { + Bytes::copy_from_slice(bytes) + } else { + table + .file_io() + .new_input(&paths[1]) + .unwrap() + .read() + .await + .unwrap() + }; + assert_eq!(bytes.len(), size); + builder.new_commit().abort(&messages).await.unwrap(); + for path in paths { + assert!(!table.file_io().exists(&path).await.unwrap()); + } + } +} + +#[tokio::test] +async fn test_file_index_invalid_configuration_fails_before_writing() { + let cases = vec![ + vec![("file-index.bitmap.columns", "missing")], + vec![("file-index.unknown.columns", "id")], + vec![("file-index.bitmap.columns", "")], + vec![("file-index.bitmap.columns", "id,")], + vec![("file-index.bitmap.columns", "id[nested]")], + vec![ + ("file-index.bitmap.columns", "id"), + ("file-index.bitmap.id.version", "1"), + ], + vec![ + ("file-index.bitmap.columns", "id"), + ("file-index.bitmap.id.index-block-size", "0 B"), + ], + vec![ + ("file-index.bloom-filter.columns", "id"), + ("file-index.bloom-filter.id.items", "0"), + ], + vec![ + ("file-index.bloom-filter.columns", "id"), + ("file-index.bloom-filter.id.fpp", "1.5"), + ], + vec![("file-index.bloom-filter.id.items", "10")], + vec![ + ("file-index.bitmap.columns", "id"), + ("file-index.bitmap.id.typo", "1"), + ], + vec![("file-index.in-manifest-threshold", "invalid")], + vec![("file-index.in-manifest-threshold", "-1")], + vec![("file-index.in-manifest-threshold", "9223372036854775807 TB")], + ]; + for options in cases { + let table = table(memory_io(), schema(&options)).await; + assert!( + table.new_write_builder().new_write().is_err(), + "{options:?}" + ); + assert!(!table + .file_io() + .exists(&format!("{}/bucket-0", table.location())) + .await + .unwrap()); + } + let schema = Schema::builder() + .column("flag", crate::spec::DataType::Boolean(BooleanType::new())) + .option("file-index.bloom-filter.columns", "flag") + .build() + .unwrap(); + assert!(matches!( + FileIndexOptions::parse(schema.options(), schema.fields()), + Err(Error::Unsupported { .. }) + )); +} + +#[tokio::test] +async fn test_file_index_rejects_unsupported_table_write_modes() { + for schema in [ + Schema::builder() + .column("id", crate::spec::DataType::Int(IntType::new())) + .primary_key(["id"]) + .option("bucket", "1") + .option("file-index.bitmap.columns", "id") + .build() + .unwrap(), + Schema::builder() + .column("id", crate::spec::DataType::Int(IntType::new())) + .option("data-evolution.enabled", "true") + .option("row-tracking.enabled", "true") + .option("file-index.bitmap.columns", "id") + .build() + .unwrap(), + ] { + let table = table(memory_io(), schema).await; + let error = match table.new_write_builder().new_write() { + Ok(_) => panic!("unsupported write mode must reject index generation"), + Err(error) => error, + }; + assert!( + error + .to_string() + .contains("FileIndex generation supports ordinary append writes only"), + "{error}" + ); + } +} + +#[tokio::test] +async fn test_file_index_uses_partition_bucket_file_row_order() { + let schema = Schema::builder() + .column("id", crate::spec::DataType::Int(IntType::new())) + .column("value", crate::spec::DataType::Int(IntType::new())) + .partition_keys(["value"]) + .option("bucket", "2") + .option("bucket-key", "id") + .option("file-index.bitmap.columns", "id,value") + .option("file-index.in-manifest-threshold", "0 B") + .build() + .unwrap(); + let table = table(memory_io(), schema).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + let data = batch( + vec![Some(3), Some(1), Some(1), Some(3)], + vec![Some(1), Some(2), Some(1), Some(2)], + ); + writer.write_arrow_batch(&data).await.unwrap(); + let messages = writer.prepare_commit().await.unwrap(); + assert_eq!( + messages + .iter() + .map(|m| &m.partition) + .collect::>() + .len(), + 2 + ); + builder.new_commit().commit(messages).await.unwrap(); + let read_builder = table.new_read_builder(); + let plan = read_builder.new_scan().plan().await.unwrap(); + for split in plan.splits() { + for file in split.data_files() { + let single = DataSplitBuilder::new() + .with_snapshot(1) + .with_partition(split.partition().clone()) + .with_bucket(split.bucket()) + .with_bucket_path(split.bucket_path().to_string()) + .with_total_buckets(2) + .with_data_files(vec![file.clone()]) + .build() + .unwrap(); + let batches: Vec = read_builder + .new_read() + .unwrap() + .to_arrow(&[single]) + .unwrap() + .try_collect() + .await + .unwrap(); + let expected = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .iter() + }) + .enumerate() + .filter_map(|(row, id)| (id == Some(3)).then_some(row as u32)) + .collect(); + let actual = evaluate( + &table, + split.bucket_path(), + file, + PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(3)) + .unwrap(), + ) + .await; + match actual { + FileIndexResult::Selection(rows) => assert_eq!(rows, expected), + FileIndexResult::Skip => assert!(roaring::RoaringBitmap::is_empty(&expected)), + FileIndexResult::Remain => panic!("bitmap equality must select physical rows"), + } + } + } + assert_eq!(query(&table, true, None).await, rows(&[data])); +} + +#[tokio::test] +async fn test_file_index_other_append_formats() { + for format in ["row"] + .into_iter() + .chain(cfg!(feature = "vortex").then_some("vortex")) + { + let table = table( + memory_io(), + schema(&[ + ("file.format", format), + ("file-index.bitmap.columns", "id"), + ("file-index.in-manifest-threshold", "0 B"), + ]), + ) + .await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .write_arrow_batch(&batch( + vec![Some(1), None, Some(3)], + vec![None, Some(2), Some(3)], + )) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(3)) + .unwrap(); + for enabled in [false, true] { + assert_eq!( + query(&table, enabled, Some(predicate.clone())).await, + vec![(Some(3), Some(3))] + ); + } + } +} + +#[derive(Debug)] +struct StorageProbe { + op: Operator, + data_accesses: AtomicUsize, + index_accesses: AtomicUsize, + fail_index_at: usize, +} + +impl StorageProbe { + fn new(fail_index_at: usize) -> Arc { + Arc::new(Self { + op: Operator::from_config(MemoryConfig::default()).unwrap(), + data_accesses: AtomicUsize::new(0), + index_accesses: AtomicUsize::new(0), + fail_index_at, + }) + } + fn io(self: &Arc) -> FileIO { + memory_io().with_provider(self.clone()) + } +} + +#[async_trait::async_trait] +impl FileIOProvider for StorageProbe { + async fn create(&self, path: &str) -> Result<(Operator, String)> { + let relative = path.strip_prefix("memory:/").unwrap().to_string(); + if path.ends_with(".parquet") { + self.data_accesses.fetch_add(1, Ordering::SeqCst); + } + if path.ends_with(".index") + && self.index_accesses.fetch_add(1, Ordering::SeqCst) + 1 == self.fail_index_at + { + self.op + .write(&relative, Bytes::from_static(b"partial index")) + .await + .unwrap(); + return Err(Error::DataInvalid { + message: "Injected sidecar write failure".to_string(), + source: None, + }); + } + Ok((self.op.clone(), relative)) + } +} + +#[tokio::test] +async fn test_file_index_prunes_without_opening_data_file() { + for threshold in ["0 B", "1 MB"] { + let storage = StorageProbe::new(0); + let table = table( + storage.io(), + schema(&[ + ("file-index.bitmap.columns", "id"), + ("file-index.in-manifest-threshold", threshold), + ]), + ) + .await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .write_arrow_batch(&batch(vec![Some(1), Some(3)], vec![None, None])) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let predicate = PredicateBuilder::new(table.schema().fields()) + .equal("id", Datum::Int(2)) + .unwrap(); + let mut read_builder = table.new_read_builder(); + read_builder.with_filter(predicate.clone()); + let (_, trace) = read_builder.new_scan().plan_with_trace().await.unwrap(); + assert_eq!(trace.final_files, 1, "statistics must retain the file"); + storage.data_accesses.store(0, Ordering::SeqCst); + assert!(query(&table, true, Some(predicate.clone())) + .await + .is_empty()); + assert_eq!(storage.data_accesses.load(Ordering::SeqCst), 0); + assert!(query(&table, false, Some(predicate)).await.is_empty()); + assert!(storage.data_accesses.load(Ordering::SeqCst) > 0); + } +} + +#[tokio::test] +async fn test_file_index_bloom_false_positive_keeps_residual_filter() { + let table = table( + memory_io(), + schema(&[ + ("file-index.bloom-filter.columns", "id"), + ("file-index.bloom-filter.id.items", "1"), + ("file-index.bloom-filter.id.fpp", "0.99"), + ]), + ) + .await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer + .write_arrow_batch(&batch(vec![Some(1), Some(10_000)], vec![None, None])) + .await + .unwrap(); + builder + .new_commit() + .commit(writer.prepare_commit().await.unwrap()) + .await + .unwrap(); + let read_builder = table.new_read_builder(); + let plan = read_builder.new_scan().plan().await.unwrap(); + let split = &plan.splits()[0]; + let predicates = PredicateBuilder::new(table.schema().fields()); + let mut false_positive = None; + for candidate in 2..10_000 { + let predicate = predicates.equal("id", Datum::Int(candidate)).unwrap(); + if evaluate( + &table, + split.bucket_path(), + &split.data_files()[0], + predicate.clone(), + ) + .await + == FileIndexResult::Remain + { + false_positive = Some(predicate); + break; + } + } + let predicate = false_positive.expect("high-FPP filter should have an in-range false positive"); + for enabled in [false, true] { + assert!(query(&table, enabled, Some(predicate.clone())) + .await + .is_empty()); + } +} + +#[tokio::test] +async fn test_file_index_sidecar_failure_cleans_all_partitions_and_rolled_files() { + for partitioned in [false, true] { + let storage = StorageProbe::new(2); + let mut schema = schema(&[ + ("file-index.bitmap.columns", "id"), + ("target-file-size", "1 B"), + ("file-index.in-manifest-threshold", "0 B"), + ]); + if partitioned { + schema = Schema::builder() + .column("id", crate::spec::DataType::Int(IntType::new())) + .column("value", crate::spec::DataType::Int(IntType::new())) + .partition_keys(["value"]) + .option("file-index.bitmap.columns", "id") + .option("file-index.in-manifest-threshold", "0 B") + .build() + .unwrap(); + } + let table = table(storage.io(), schema).await; + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + for value in 1..=3 { + writer + .write_arrow_batch(&batch(vec![Some(value)], vec![Some(value)])) + .await + .unwrap(); + } + let error = writer.prepare_commit().await.unwrap_err(); + assert!(error.to_string().contains("Injected sidecar"), "{error}"); + let files = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap(); + assert!( + !files + .iter() + .any(|f| f.path.ends_with(".parquet") || f.path.ends_with(".index")), + "{files:?}" + ); + } +} + +#[tokio::test] +async fn test_file_index_serialization_failure_cleans_data() { + for rolling in [false, true] { + let name = "a".repeat(65536); + let schema = Schema::builder() + .column(&name, crate::spec::DataType::Int(IntType::new())) + .option("file-index.bitmap.columns", &name) + .option("target-file-size", if rolling { "1 B" } else { "128 MB" }) + .build() + .unwrap(); + let table = table(memory_io(), schema).await; + let data = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![Field::new( + &name, + ArrowType::Int32, + true, + )])), + vec![Arc::new(Int32Array::from(vec![Some(1)]))], + ) + .unwrap(); + let builder = table.new_write_builder(); + let mut writer = builder.new_write().unwrap(); + writer.write_arrow_batch(&data).await.unwrap(); + assert!(matches!( + writer.prepare_commit().await, + Err(Error::FileIndexFormatInvalid { .. }) + )); + let files = table + .file_io() + .list_status_recursive(table.location()) + .await + .unwrap(); + assert!(!files + .iter() + .any(|f| f.path.ends_with(".parquet") || f.path.ends_with(".index"))); + } +} diff --git a/crates/paimon/src/table/data_file_writer.rs b/crates/paimon/src/table/data_file_writer.rs index 76f49a0d0..cc5969f7e 100644 --- a/crates/paimon/src/table/data_file_writer.rs +++ b/crates/paimon/src/table/data_file_writer.rs @@ -22,14 +22,17 @@ //! handles file rolling when `target_file_size` is reached, and collects //! [`DataFileMeta`] for the commit path. +use super::data_file_index_writer::{DataFileIndexWriter, FileIndexOptions}; use crate::arrow::format::{create_format_writer, FormatFileWriter, FormatValueStats}; use crate::io::FileIO; +use crate::spec::data_file_to_file_index_file_name; use crate::spec::stats::BinaryTableStats; use crate::spec::{bucket_dir_name, DataField, DataFileMeta, EMPTY_SERIALIZED_ROW}; use crate::Result; use arrow_array::RecordBatch; use chrono::Utc; use std::collections::HashMap; +use std::sync::Arc; use tokio::task::JoinSet; /// Low-level writer that produces Parquet data files for a single (partition, bucket). @@ -62,6 +65,10 @@ pub(crate) struct DataFileWriter { current_writer: Option>, current_file_name: Option, current_row_count: i64, + index_options: Option>, + current_index: Option, + /// Paths owned by this indexed write until prepare_commit hands them to the caller. + created_paths: Vec, } impl DataFileWriter { @@ -104,11 +111,27 @@ impl DataFileWriter { current_writer: None, current_file_name: None, current_row_count: 0, + index_options: None, + current_index: None, + created_paths: Vec::new(), } } + pub(super) fn with_file_index(mut self, options: Option>) -> Self { + self.index_options = options; + self + } + /// Write a RecordBatch. Rolls to a new file when target size is reached. pub(crate) async fn write(&mut self, batch: &RecordBatch) -> Result<()> { + let result = self.write_batch(batch).await; + if result.is_err() && self.index_options.is_some() { + self.abort().await; + } + result + } + + async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> { if batch.num_rows() == 0 { return Ok(()); } @@ -118,6 +141,9 @@ impl DataFileWriter { } self.current_writer.as_mut().unwrap().write(batch).await?; + if let Some(index) = &mut self.current_index { + index.write(batch)?; + } self.current_row_count += batch.num_rows() as i64; // Roll to a new file if target size is reached — close in background @@ -136,25 +162,28 @@ impl DataFileWriter { } async fn open_new_file(&mut self, schema: arrow_schema::SchemaRef) -> Result<()> { + let index = self + .index_options + .as_ref() + .map(|options| options.create_writer()) + .transpose()?; let file_name = format!( "data-{}-{}.{}", uuid::Uuid::new_v4(), self.written_files.len(), self.file_format, ); - let bucket_dir = if self.partition_path.is_empty() { - format!("{}/{}", self.table_location, bucket_dir_name(self.bucket)) - } else { - format!( - "{}/{}/{}", - self.table_location, - self.partition_path, - bucket_dir_name(self.bucket) - ) - }; + let bucket_dir = self.bucket_dir(); self.file_io.mkdirs(&format!("{bucket_dir}/")).await?; let file_path = format!("{bucket_dir}/{file_name}"); + if self.index_options.is_some() { + self.created_paths.push(file_path.clone()); + self.created_paths.push(format!( + "{bucket_dir}/{}", + data_file_to_file_index_file_name(&file_name) + )); + } let output = self.file_io.new_output(&file_path)?; let writer = create_format_writer( &output, @@ -167,6 +196,7 @@ impl DataFileWriter { ) .await?; self.current_writer = Some(writer); + self.current_index = index; self.current_file_name = Some(file_name); self.current_row_count = 0; Ok(()) @@ -174,36 +204,30 @@ impl DataFileWriter { /// Close the current file writer and record the file metadata. pub(crate) async fn close_current_file(&mut self) -> Result<()> { - let writer = match self.current_writer.take() { - Some(w) => w, - None => return Ok(()), - }; - let file_name = self.current_file_name.take().unwrap(); - - let row_count = self.current_row_count; - self.current_row_count = 0; - let write_result = writer.close().await?; - - let meta = Self::build_meta( - file_name, - write_result.file_size as i64, - row_count, - self.schema_id, - self.file_source, - self.first_row_id, - self.write_cols.clone(), - write_result.value_stats, - ); - self.written_files.push(meta); + if let Some(close) = self.take_close() { + self.written_files.push(close.await?); + } Ok(()) } /// Spawn the current writer's close in the background for non-blocking rolling. fn roll_file(&mut self) { - let writer = match self.current_writer.take() { - Some(w) => w, - None => return, - }; + if let Some(close) = self.take_close() { + self.in_flight_closes.spawn(close); + } + } + + fn take_close( + &mut self, + ) -> Option> + Send + 'static> { + let writer = self.current_writer.take()?; + let index = self.current_index.take(); + let file_io = self.file_io.clone(); + let bucket_dir = self.bucket_dir(); + let threshold = self + .index_options + .as_ref() + .map(|options| options.in_manifest_threshold); let file_name = self.current_file_name.take().unwrap(); let row_count = self.current_row_count; self.current_row_count = 0; @@ -212,9 +236,9 @@ impl DataFileWriter { let first_row_id = self.first_row_id; let write_cols = self.write_cols.clone(); - self.in_flight_closes.spawn(async move { + Some(async move { let write_result = writer.close().await?; - Ok(Self::build_meta( + let mut meta = Self::build_meta( file_name, write_result.file_size as i64, row_count, @@ -223,12 +247,34 @@ impl DataFileWriter { first_row_id, write_cols, write_result.value_stats, - )) - }); + ); + if let Some(index) = index { + let bytes = index.serialize()?; + if bytes.len() as u64 > threshold.unwrap() as u64 { + let name = data_file_to_file_index_file_name(&meta.file_name); + file_io + .new_output(&format!("{bucket_dir}/{name}"))? + .write(bytes) + .await?; + meta.extra_files.push(name); + } else { + meta.embedded_index = Some(bytes.to_vec()); + } + } + Ok(meta) + }) } /// Close the current writer and return all written file metadata. pub(crate) async fn prepare_commit(&mut self) -> Result> { + let result = self.finish().await; + if result.is_err() && self.index_options.is_some() { + self.abort().await; + } + result + } + + async fn finish(&mut self) -> Result> { self.close_current_file().await?; while let Some(result) = self.in_flight_closes.join_next().await { let meta = result.map_err(|e| crate::Error::DataInvalid { @@ -237,9 +283,36 @@ impl DataFileWriter { })??; self.written_files.push(meta); } + self.created_paths.clear(); Ok(std::mem::take(&mut self.written_files)) } + pub(super) async fn abort(&mut self) { + if let Some(writer) = self.current_writer.take() { + let _ = writer.close().await; + } + self.current_index = None; + self.current_file_name = None; + while self.in_flight_closes.join_next().await.is_some() {} + for path in self.created_paths.drain(..) { + let _ = self.file_io.delete_file(&path).await; + } + self.written_files.clear(); + } + + fn bucket_dir(&self) -> String { + if self.partition_path.is_empty() { + format!("{}/{}", self.table_location, bucket_dir_name(self.bucket)) + } else { + format!( + "{}/{}/{}", + self.table_location, + self.partition_path, + bucket_dir_name(self.bucket) + ) + } + } + #[allow(clippy::too_many_arguments)] fn build_meta( file_name: String, diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 708aa6b34..ed3e0b065 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -37,6 +37,7 @@ mod consumer_manager; pub(crate) mod cow_writer; mod data_evolution_reader; pub mod data_evolution_writer; +mod data_file_index_writer; mod data_file_reader; mod data_file_writer; mod dedicated_format_file_writer; diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 1a0c39938..04153ec8a 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -34,6 +34,7 @@ use crate::table::bucket_assigner_dynamic::DynamicBucketAssigner; use crate::table::bucket_assigner_fixed::FixedBucketAssigner; use crate::table::bucket_function::validate_bucket_function; use crate::table::commit_message::CommitMessage; +use crate::table::data_file_index_writer::FileIndexOptions; use crate::table::data_file_writer::DataFileWriter; use crate::table::dedicated_format_file_writer::AppendDedicatedFormatFileWriter; use crate::table::kv_file_writer::{KeyValueFileWriter, KeyValueWriteConfig}; @@ -148,6 +149,7 @@ pub struct TableWrite { has_dedicated_vector_fields: bool, row_kind_generator: Option, row_kind_filter: Option, + file_index_options: Option>, } impl TableWrite { @@ -374,6 +376,19 @@ impl TableWrite { .iter() .any(|f| matches!(f.data_type(), DataType::Vector(_))); + let file_index_options = FileIndexOptions::parse(schema.options(), schema.fields())?; + if file_index_options.is_some() + && (has_primary_keys + || has_blob_fields + || has_dedicated_vector_fields + || !blob_view_fields.is_empty() + || core_options.data_evolution_enabled()) + { + return Err(crate::Error::Unsupported { + message: "FileIndex generation supports ordinary append writes only; primary-key, data-evolution and dedicated Blob/Vector writes are not supported".to_string(), + }); + } + Ok(Self { table: table.clone(), write_schema, @@ -408,6 +423,7 @@ impl TableWrite { has_dedicated_vector_fields, row_kind_generator, row_kind_filter, + file_index_options: file_index_options.map(Arc::new), }) } @@ -801,12 +817,26 @@ impl TableWrite { bucket: i32, batch: RecordBatch, ) -> Result<()> { - let key = (partition_bytes, bucket); - if !self.partition_writers.contains_key(&key) { - self.create_writer(key.0.clone(), key.1).await?; + let result = async { + let key = (partition_bytes, bucket); + if !self.partition_writers.contains_key(&key) { + self.create_writer(key.0.clone(), key.1).await?; + } + self.partition_writers + .get_mut(&key) + .unwrap() + .write(&batch) + .await + } + .await; + if result.is_err() && self.file_index_options.is_some() { + for (_, writer) in self.partition_writers.drain() { + if let FileWriter::Append(mut writer) = writer { + writer.abort().await; + } + } } - let writer = self.partition_writers.get_mut(&key).unwrap(); - writer.write(&batch).await + result } /// Write multiple Arrow RecordBatches. @@ -820,6 +850,9 @@ impl TableWrite { /// Close all writers and collect CommitMessages for use with TableCommit. /// Writers are cleared after this call, allowing the TableWrite to be reused. pub async fn prepare_commit(&mut self) -> Result> { + if self.file_index_options.is_some() { + return self.prepare_indexed_append_commit().await; + } let writers: Vec<(PartitionBucketKey, FileWriter)> = self.partition_writers.drain().collect(); @@ -863,6 +896,37 @@ impl TableWrite { Ok(messages) } + async fn prepare_indexed_append_commit(&mut self) -> Result> { + let closes = + self.partition_writers + .drain() + .map(|((partition, bucket), writer)| async move { + (partition, bucket, writer.prepare_commit().await) + }); + // Do not cancel another partition's close when one fails: its completed + // files must remain reachable for abort cleanup. + let results = futures::future::join_all(closes).await; + let mut messages = Vec::new(); + let mut error = None; + for (partition, bucket, result) in results { + match result { + Ok(files) if !files.data_files.is_empty() => { + messages.push(CommitMessage::new(partition, bucket, files.data_files)); + } + Ok(_) => {} + Err(err) => { + error.get_or_insert(err); + } + } + } + if let Some(error) = error { + let commit = super::TableCommit::new(self.table.clone(), self.commit_user.clone()); + let _ = commit.abort(&messages).await; + return Err(error); + } + Ok(messages) + } + async fn create_writer(&mut self, partition_bytes: Vec, bucket: i32) -> Result<()> { let partition_path = self.resolve_partition_path(&partition_bytes)?; @@ -919,23 +983,26 @@ impl TableWrite { ), ))) } else { - Ok(FileWriter::Append(DataFileWriter::new( - self.table.file_io().clone(), - self.table.location().to_string(), - partition_path, - bucket, - self.schema_id, - self.target_file_size, - self.file_compression.clone(), - self.file_compression_zstd_level, - self.write_buffer_size, - self.file_format.clone(), - self.table.schema().fields().to_vec(), - self.table.schema().options().clone(), - Some(0), - None, - None, - ))) + Ok(FileWriter::Append( + DataFileWriter::new( + self.table.file_io().clone(), + self.table.location().to_string(), + partition_path, + bucket, + self.schema_id, + self.target_file_size, + self.file_compression.clone(), + self.file_compression_zstd_level, + self.write_buffer_size, + self.file_format.clone(), + self.table.schema().fields().to_vec(), + self.table.schema().options().clone(), + Some(0), + None, + None, + ) + .with_file_index(self.file_index_options.clone()), + )) } } diff --git a/docs/src/getting-started.md b/docs/src/getting-started.md index 5a8725006..8c8a9592d 100644 --- a/docs/src/getting-started.md +++ b/docs/src/getting-started.md @@ -55,6 +55,39 @@ Available storage features: Mosaic data file reads are always available. The current Mosaic support is read-only: Paimon Rust can read existing `.mosaic` data files, including array and map columns, in a Paimon table, but it does not write Mosaic data files yet. +## FileIndexes for Append Writes + +Ordinary append writes can generate Bitmap and Bloom Filter indexes for supported +top-level columns using table options: + +```text +file-index.bitmap.columns = category +file-index.bloom-filter.columns = id +file-index.bloom-filter.id.items = 100000 +file-index.bloom-filter.id.fpp = 0.01 +file-index.in-manifest-threshold = 500 B +``` + +Column lists are comma-separated. Bitmap supports `version` (currently `2` only) +and `index-block-size` per column. Bloom Filter supports `items` and `fpp`. +Invalid columns, unsupported index/data types, and invalid index options fail +when creating the writer. + +Each data file gets its own index containing all configured columns and index +types. The complete serialized index is embedded in the manifest when its size +is at most `file-index.in-manifest-threshold` (default `500 B`); larger indexes +are stored beside the data file as a `.index` sidecar. Indexed write failures +return an error and clean up newly created files on a best-effort basis. +Use commit `abort` to clean up files after a successful `prepare_commit` +when the prepared write will not be committed. + +`file-index.read.enabled` controls only reading, independently of index creation. +Existing files are not backfilled. Index generation is not supported for +primary-key writes, data-evolution writes, or dedicated Blob/Vector paths; +ordinary writer creation rejects index configuration on these paths. COW and +partial DataEvolution rewrites do not generate indexes. Nested indexes and +additional index types are not supported. + ## Catalog Management Paimon supports multiple catalog types. The `CatalogFactory` provides a unified way to create catalogs based on configuration options.