This is an automated email from the ASF dual-hosted git repository. lzljs3620320 pushed a commit to branch feat/schema-evolution-read in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
commit 3afdfd6aceb002e85198cd0c1f70dc5c31c7a93b Author: JingsongLi <[email protected]> AuthorDate: Fri Apr 3 10:03:20 2026 +0800 feat(arrow): support schema evolution read with SchemaManager and field-ID-based index mapping Introduce SchemaManager to load versioned schema files (schema/schema-{id}) and use field-ID-based index mapping (aligned with Java Paimon's SchemaEvolutionUtil) for reading data files written under older schemas. This supports add-column (null fill) and type promotion (cast) scenarios. Also align paimon_type_to_arrow with Python's from_paimon_type: - TIME always maps to Time32(Millisecond) - TIMESTAMP precision 0 maps to Second - Add nested type support: Array, Map, Multiset, Row Co-Authored-By: Claude Opus 4.6 <[email protected]> --- Cargo.toml | 1 + crates/integration_tests/tests/read_tables.rs | 179 +++++++++++++- crates/paimon/Cargo.toml | 1 + crates/paimon/src/arrow/mod.rs | 108 +++++++++ crates/paimon/src/arrow/reader.rs | 324 +++++++++++++++++++------- crates/paimon/src/arrow/schema_evolution.rs | 151 ++++++++++++ crates/paimon/src/spec/types.rs | 16 ++ crates/paimon/src/table/mod.rs | 10 + crates/paimon/src/table/read_builder.rs | 8 +- crates/paimon/src/table/schema_manager.rs | 81 +++++++ dev/spark/provision.py | 141 +++++++++++ 11 files changed, 936 insertions(+), 84 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index d7ebd00..f1f7295 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ rust-version = "1.86.0" [workspace.dependencies] arrow-array = { version = "57.0", features = ["ffi"] } arrow-schema = "57.0" +arrow-cast = "57.0" arrow-select = "57.0" parquet = "57.0" tokio = "1.39.2" \ No newline at end of file diff --git a/crates/integration_tests/tests/read_tables.rs b/crates/integration_tests/tests/read_tables.rs index 090197d..8a4ad8f 100644 --- a/crates/integration_tests/tests/read_tables.rs +++ b/crates/integration_tests/tests/read_tables.rs @@ -17,7 +17,7 @@ //! Integration tests for reading Paimon tables provisioned by Spark. -use arrow_array::{Int32Array, RecordBatch, StringArray}; +use arrow_array::{Array, ArrowPrimitiveType, Int32Array, Int64Array, RecordBatch, StringArray}; use futures::TryStreamExt; use paimon::api::ConfigResponse; use paimon::catalog::{Identifier, RESTCatalog}; @@ -947,3 +947,180 @@ async fn test_read_data_evolution_table_with_projection() { "Projected data evolution read should return correct values" ); } + +// --------------------------------------------------------------------------- +// Schema Evolution integration tests +// --------------------------------------------------------------------------- + +/// Test reading a table after ALTER TABLE ADD COLUMNS. +/// Old data files lack the new column; reader should fill nulls. +#[tokio::test] +async fn test_read_schema_evolution_add_column() { + let (_, batches) = scan_and_read_with_fs_catalog("schema_evolution_add_column", None).await; + + let mut rows: Vec<(i32, String, Option<i32>)> = Vec::new(); + for batch in &batches { + let id = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::<Int32Array>()) + .expect("id"); + let name = batch + .column_by_name("name") + .and_then(|c| c.as_any().downcast_ref::<StringArray>()) + .expect("name"); + let age = batch + .column_by_name("age") + .and_then(|c| c.as_any().downcast_ref::<Int32Array>()) + .expect("age"); + for i in 0..batch.num_rows() { + let age_val = if age.is_null(i) { + None + } else { + Some(age.value(i)) + }; + rows.push((id.value(i), name.value(i).to_string(), age_val)); + } + } + rows.sort_by_key(|(id, _, _)| *id); + + assert_eq!( + rows, + vec![ + (1, "alice".into(), None), + (2, "bob".into(), None), + (3, "carol".into(), Some(30)), + (4, "dave".into(), Some(40)), + ], + "Old rows should have null for added column 'age'" + ); +} + +/// Test reading a table after ALTER TABLE ALTER COLUMN TYPE (INT -> BIGINT). +/// Old data files have INT; reader should cast to BIGINT. +#[tokio::test] +async fn test_read_schema_evolution_type_promotion() { + let (_, batches) = scan_and_read_with_fs_catalog("schema_evolution_type_promotion", None).await; + + // Verify the value column is Int64 (BIGINT) in all batches + for batch in &batches { + let value_col = batch.column_by_name("value").expect("value column"); + assert_eq!( + value_col.data_type(), + &arrow_array::types::Int64Type::DATA_TYPE, + "value column should be Int64 (BIGINT) after type promotion" + ); + } + + let mut rows: Vec<(i32, i64)> = Vec::new(); + for batch in &batches { + let id = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::<Int32Array>()) + .expect("id"); + let value = batch + .column_by_name("value") + .and_then(|c| c.as_any().downcast_ref::<Int64Array>()) + .expect("value as Int64Array"); + for i in 0..batch.num_rows() { + rows.push((id.value(i), value.value(i))); + } + } + rows.sort_by_key(|(id, _)| *id); + + assert_eq!( + rows, + vec![(1, 100i64), (2, 200i64), (3, 3_000_000_000i64)], + "INT values should be promoted to BIGINT, including values > INT_MAX" + ); +} + +/// Test reading a data-evolution table after ALTER TABLE ADD COLUMNS. +/// Old files lack the new column; reader should fill nulls even in data evolution mode. +#[tokio::test] +async fn test_read_data_evolution_add_column() { + let (_, batches) = scan_and_read_with_fs_catalog("data_evolution_add_column", None).await; + + let mut rows: Vec<(i32, String, i32, Option<String>)> = Vec::new(); + for batch in &batches { + let id = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::<Int32Array>()) + .expect("id"); + let name = batch + .column_by_name("name") + .and_then(|c| c.as_any().downcast_ref::<StringArray>()) + .expect("name"); + let value = batch + .column_by_name("value") + .and_then(|c| c.as_any().downcast_ref::<Int32Array>()) + .expect("value"); + let extra = batch + .column_by_name("extra") + .and_then(|c| c.as_any().downcast_ref::<StringArray>()) + .expect("extra"); + for i in 0..batch.num_rows() { + let extra_val = if extra.is_null(i) { + None + } else { + Some(extra.value(i).to_string()) + }; + rows.push(( + id.value(i), + name.value(i).to_string(), + value.value(i), + extra_val, + )); + } + } + rows.sort_by_key(|(id, _, _, _)| *id); + + assert_eq!( + rows, + vec![ + (1, "alice-v2".into(), 100, None), + (2, "bob".into(), 200, None), + (3, "carol".into(), 300, Some("new".into())), + (4, "dave".into(), 400, Some("new".into())), + ], + "Data evolution + add column: old rows should have null for 'extra', MERGE INTO updates name" + ); +} + +/// Test reading a data-evolution table after ALTER TABLE ALTER COLUMN TYPE (INT -> BIGINT). +/// Old files have INT; reader should cast to BIGINT in data evolution mode. +#[tokio::test] +async fn test_read_data_evolution_type_promotion() { + let (_, batches) = scan_and_read_with_fs_catalog("data_evolution_type_promotion", None).await; + + // Verify the value column is Int64 (BIGINT) in all batches + for batch in &batches { + let value_col = batch.column_by_name("value").expect("value column"); + assert_eq!( + value_col.data_type(), + &arrow_array::types::Int64Type::DATA_TYPE, + "value column should be Int64 (BIGINT) after type promotion in data evolution mode" + ); + } + + let mut rows: Vec<(i32, i64)> = Vec::new(); + for batch in &batches { + let id = batch + .column_by_name("id") + .and_then(|c| c.as_any().downcast_ref::<Int32Array>()) + .expect("id"); + let value = batch + .column_by_name("value") + .and_then(|c| c.as_any().downcast_ref::<Int64Array>()) + .expect("value as Int64Array"); + for i in 0..batch.num_rows() { + rows.push((id.value(i), value.value(i))); + } + } + rows.sort_by_key(|(id, _)| *id); + + assert_eq!( + rows, + vec![(1, 999i64), (2, 200i64), (3, 3_000_000_000i64)], + "Data evolution + type promotion: INT should be cast to BIGINT, MERGE INTO updates value" + ); +} diff --git a/crates/paimon/Cargo.toml b/crates/paimon/Cargo.toml index 878d1c5..b1677cf 100644 --- a/crates/paimon/Cargo.toml +++ b/crates/paimon/Cargo.toml @@ -55,6 +55,7 @@ apache-avro = { version = "0.17", features = ["snappy", "zstandard"] } indexmap = "2.5.0" roaring = "0.11" arrow-array = { workspace = true } +arrow-cast = { workspace = true } arrow-schema = { workspace = true } arrow-select = { workspace = true } futures = "0.3" diff --git a/crates/paimon/src/arrow/mod.rs b/crates/paimon/src/arrow/mod.rs index 0ed18b8..a46c779 100644 --- a/crates/paimon/src/arrow/mod.rs +++ b/crates/paimon/src/arrow/mod.rs @@ -16,5 +16,113 @@ // under the License. mod reader; +pub mod schema_evolution; pub use crate::arrow::reader::ArrowReaderBuilder; + +use crate::spec::DataType as PaimonDataType; +use arrow_schema::DataType as ArrowDataType; +use arrow_schema::{Field as ArrowField, TimeUnit}; +use std::sync::Arc; + +/// Converts a Paimon [`DataType`](PaimonDataType) to an Arrow [`DataType`](ArrowDataType). +pub fn paimon_type_to_arrow(dt: &PaimonDataType) -> crate::Result<ArrowDataType> { + Ok(match dt { + PaimonDataType::Boolean(_) => ArrowDataType::Boolean, + PaimonDataType::TinyInt(_) => ArrowDataType::Int8, + PaimonDataType::SmallInt(_) => ArrowDataType::Int16, + PaimonDataType::Int(_) => ArrowDataType::Int32, + PaimonDataType::BigInt(_) => ArrowDataType::Int64, + PaimonDataType::Float(_) => ArrowDataType::Float32, + PaimonDataType::Double(_) => ArrowDataType::Float64, + PaimonDataType::VarChar(_) | PaimonDataType::Char(_) => ArrowDataType::Utf8, + PaimonDataType::Binary(_) | PaimonDataType::VarBinary(_) => ArrowDataType::Binary, + PaimonDataType::Date(_) => ArrowDataType::Date32, + PaimonDataType::Time(_) => ArrowDataType::Time32(TimeUnit::Millisecond), + PaimonDataType::Timestamp(t) => { + ArrowDataType::Timestamp(timestamp_time_unit(t.precision())?, None) + } + PaimonDataType::LocalZonedTimestamp(t) => { + ArrowDataType::Timestamp(timestamp_time_unit(t.precision())?, Some("UTC".into())) + } + PaimonDataType::Decimal(d) => { + let p = u8::try_from(d.precision()).map_err(|_| crate::Error::Unsupported { + message: "Decimal precision exceeds u8".to_string(), + })?; + let s = i8::try_from(d.scale() as i32).map_err(|_| crate::Error::Unsupported { + message: "Decimal scale out of i8 range".to_string(), + })?; + ArrowDataType::Decimal128(p, s) + } + PaimonDataType::Array(a) => { + let element_type = paimon_type_to_arrow(a.element_type())?; + ArrowDataType::List(Arc::new(ArrowField::new( + "element", + element_type, + a.element_type().is_nullable(), + ))) + } + PaimonDataType::Map(m) => { + let key_type = paimon_type_to_arrow(m.key_type())?; + let value_type = paimon_type_to_arrow(m.value_type())?; + ArrowDataType::Map( + Arc::new(ArrowField::new( + "entries", + ArrowDataType::Struct( + vec![ + ArrowField::new("key", key_type, false), + ArrowField::new("value", value_type, m.value_type().is_nullable()), + ] + .into(), + ), + false, + )), + false, + ) + } + PaimonDataType::Multiset(m) => { + let element_type = paimon_type_to_arrow(m.element_type())?; + ArrowDataType::Map( + Arc::new(ArrowField::new( + "entries", + ArrowDataType::Struct( + vec![ + ArrowField::new("key", element_type, m.element_type().is_nullable()), + ArrowField::new("value", ArrowDataType::Int32, false), + ] + .into(), + ), + false, + )), + false, + ) + } + PaimonDataType::Row(r) => { + let fields: Vec<ArrowField> = r + .fields() + .iter() + .map(|f| { + let arrow_type = paimon_type_to_arrow(f.data_type())?; + Ok(ArrowField::new( + f.name(), + arrow_type, + f.data_type().is_nullable(), + )) + }) + .collect::<crate::Result<Vec<_>>>()?; + ArrowDataType::Struct(fields.into()) + } + }) +} + +fn timestamp_time_unit(precision: u32) -> crate::Result<TimeUnit> { + match precision { + 0 => Ok(TimeUnit::Second), + 1..=3 => Ok(TimeUnit::Millisecond), + 4..=6 => Ok(TimeUnit::Microsecond), + 7..=9 => Ok(TimeUnit::Nanosecond), + _ => Err(crate::Error::Unsupported { + message: format!("Unsupported TIMESTAMP precision {precision}"), + }), + } +} diff --git a/crates/paimon/src/arrow/reader.rs b/crates/paimon/src/arrow/reader.rs index 7d57243..561f969 100644 --- a/crates/paimon/src/arrow/reader.rs +++ b/crates/paimon/src/arrow/reader.rs @@ -15,12 +15,16 @@ // specific language governing permissions and limitations // under the License. +use crate::arrow::paimon_type_to_arrow; +use crate::arrow::schema_evolution::{create_index_mapping, NULL_FIELD_INDEX}; use crate::deletion_vector::{DeletionVector, DeletionVectorFactory}; use crate::io::{FileIO, FileRead, FileStatus}; use crate::spec::{DataField, DataFileMeta}; +use crate::table::schema_manager::SchemaManager; use crate::table::ArrowRecordBatchStream; use crate::{DataSplit, Error}; use arrow_array::RecordBatch; +use arrow_cast::cast; use arrow_schema::{Field as ArrowField, Schema as ArrowSchema}; use async_stream::try_stream; @@ -41,14 +45,22 @@ use tokio::try_join; pub struct ArrowReaderBuilder { batch_size: Option<usize>, file_io: FileIO, + schema_manager: SchemaManager, + table_schema_id: i64, } impl ArrowReaderBuilder { /// Create a new ArrowReaderBuilder - pub(crate) fn new(file_io: FileIO) -> Self { + pub(crate) fn new( + file_io: FileIO, + schema_manager: SchemaManager, + table_schema_id: i64, + ) -> Self { ArrowReaderBuilder { batch_size: None, file_io, + schema_manager, + table_schema_id, } } @@ -58,6 +70,8 @@ impl ArrowReaderBuilder { ArrowReader { batch_size: self.batch_size, file_io: self.file_io, + schema_manager: self.schema_manager, + table_schema_id: self.table_schema_id, read_type, } } @@ -68,33 +82,30 @@ impl ArrowReaderBuilder { pub struct ArrowReader { batch_size: Option<usize>, file_io: FileIO, + schema_manager: SchemaManager, + table_schema_id: i64, read_type: Vec<DataField>, } impl ArrowReader { /// Take a stream of DataSplits and read every data file in each split. /// Returns a stream of Arrow RecordBatches from all files. - /// When a split has deletion files (see [DataSplit::data_deletion_files]), the corresponding - /// deletion vectors are loaded and applied so that deleted rows are filtered out from the stream. - /// Row positions are 0-based within each data file, matching Java's ApplyDeletionVectorReader. /// - /// Matches [RawFileSplitRead.createReader](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java): - /// one DV factory per DataSplit (created from that split's data files and deletion files). + /// Uses SchemaManager to load the data file's schema (via `DataFileMeta.schema_id`) + /// and computes field-ID-based index mapping for schema evolution (added columns, + /// type promotion, column reordering). /// - /// Parquet schema is clipped to this reader's read type (column names from [DataField]s). - /// File-only columns are not read. See [ParquetReaderFactory.clipParquetSchema](https://github.com/apache/paimon/blob/master/paimon-format/paimon-format-common/src/main/java/org/apache/paimon/format/FormatReaderFactory.java). + /// Matches [RawFileSplitRead.createReader](https://github.com/apache/paimon/blob/master/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java). pub fn read(self, data_splits: &[DataSplit]) -> crate::Result<ArrowRecordBatchStream> { let file_io = self.file_io.clone(); let batch_size = self.batch_size; let splits: Vec<DataSplit> = data_splits.to_vec(); let read_type = self.read_type; - let projected_column_names: Vec<String> = read_type - .iter() - .map(|field| field.name().to_string()) - .collect(); + let mut schema_manager = self.schema_manager; + let table_schema_id = self.table_schema_id; Ok(try_stream! { for split in splits { - // Create DV factory for this split only (like Java createReader(partition, bucket, files, deletionFiles)). + // Create DV factory for this split only. let dv_factory = if split .data_deletion_files() .is_some_and(|files| files.iter().any(Option::is_some)) @@ -117,11 +128,20 @@ impl ArrowReader { .and_then(|factory| factory.get_deletion_vector(&file_meta.file_name)) .cloned(); + // Load data file's schema if it differs from the table schema. + let data_fields: Option<Vec<DataField>> = if file_meta.schema_id != table_schema_id { + let data_schema = schema_manager.schema(file_meta.schema_id).await?; + Some(data_schema.fields().to_vec()) + } else { + None + }; + let mut stream = read_single_file_stream( file_io.clone(), split.clone(), file_meta, - projected_column_names.clone(), + read_type.clone(), + data_fields, batch_size, dv, )?; @@ -157,26 +177,38 @@ impl ArrowReader { .iter() .map(|field| field.name().to_string()) .collect(); + let mut schema_manager = self.schema_manager; + let table_schema_id = self.table_schema_id; Ok(try_stream! { for split in splits { if split.raw_convertible() || split.data_files().len() == 1 { - // Single file or raw convertible — stream lazily without loading all into memory. for file_meta in split.data_files().to_vec() { + let data_fields: Option<Vec<DataField>> = if file_meta.schema_id != table_schema_id { + let data_schema = schema_manager.schema(file_meta.schema_id).await?; + Some(data_schema.fields().to_vec()) + } else { + None + }; + let mut stream = read_single_file_stream( - file_io.clone(), split.clone(), file_meta, projected_column_names.clone(), batch_size, None, + file_io.clone(), split.clone(), file_meta, read_type.clone(), + data_fields, batch_size, None, )?; while let Some(batch) = stream.next().await { yield batch?; } } } else { - // Multiple files need column-wise merge — also streamed lazily. + // Multiple files need column-wise merge. let mut merge_stream = merge_files_by_columns( &file_io, &split, + &read_type, &projected_column_names, &table_field_names, + schema_manager.clone(), + table_schema_id, batch_size, )?; while let Some(batch) = merge_stream.next().await { @@ -189,16 +221,72 @@ impl ArrowReader { } } +/// Build the target Arrow schema from the read type (Paimon DataFields). +fn build_target_arrow_schema(read_type: &[DataField]) -> crate::Result<Arc<ArrowSchema>> { + let fields: Vec<ArrowField> = read_type + .iter() + .map(|f| { + let arrow_type = paimon_type_to_arrow(f.data_type())?; + Ok(ArrowField::new( + f.name(), + arrow_type, + f.data_type().is_nullable(), + )) + }) + .collect::<crate::Result<Vec<_>>>()?; + Ok(Arc::new(ArrowSchema::new(fields))) +} + /// Read a single parquet file from a split, returning a lazy stream of batches. /// Optionally applies a deletion vector. +/// +/// Handles schema evolution using field-ID-based index mapping: +/// - `data_fields`: if `Some`, the fields from the data file's schema (loaded via SchemaManager). +/// Used to compute index mapping between `read_type` and data fields by field ID. +/// - Columns missing from the file are filled with null arrays. +/// - Columns whose Arrow type differs from the target type are cast (type promotion). +/// +/// Reference: [RawFileSplitRead.createFileReader](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/operation/RawFileSplitRead.java) fn read_single_file_stream( file_io: FileIO, split: DataSplit, file_meta: DataFileMeta, - projected_column_names: Vec<String>, + read_type: Vec<DataField>, + data_fields: Option<Vec<DataField>>, batch_size: Option<usize>, dv: Option<Arc<DeletionVector>>, ) -> crate::Result<ArrowRecordBatchStream> { + let target_schema = build_target_arrow_schema(&read_type)?; + + // Compute index mapping and determine which columns to read from the parquet file. + // If data_fields is provided, use field-ID-based mapping; otherwise use read_type names directly. + let (parquet_read_fields, index_mapping) = if let Some(ref df) = data_fields { + let mapping = create_index_mapping(&read_type, df); + match mapping { + Some(ref idx_map) => { + // Only read data fields that are referenced by the index mapping. + let fields_to_read: Vec<DataField> = idx_map + .iter() + .filter(|&&idx| idx != NULL_FIELD_INDEX) + .map(|&idx| df[idx as usize].clone()) + .collect(); + (fields_to_read, Some(idx_map.clone())) + } + None => { + // Identity mapping — read data fields in order. + (df.clone(), None) + } + } + } else { + // No schema evolution — read by read_type names. + (read_type.clone(), None) + }; + + let parquet_column_names: Vec<String> = parquet_read_fields + .iter() + .map(|f| f.name().to_string()) + .collect(); + Ok(try_stream! { let path_to_read = split.data_file_path(&file_meta); if !path_to_read.to_ascii_lowercase().ends_with(".parquet") { @@ -219,7 +307,7 @@ fn read_single_file_stream( // Only project columns that exist in this file. let parquet_schema = batch_stream_builder.parquet_schema().clone(); let file_column_names: Vec<&str> = parquet_schema.columns().iter().map(|c| c.name()).collect(); - let available_columns: Vec<&str> = projected_column_names + let available_columns: Vec<&str> = parquet_column_names .iter() .filter(|name| file_column_names.contains(&name.as_str())) .map(String::as_str) @@ -242,24 +330,74 @@ fn read_single_file_stream( let mut batch_stream = batch_stream_builder.build()?; while let Some(batch) = batch_stream.next().await { let batch = batch?; - // Reorder columns from parquet-schema order to projected_column_names order, - // consistent with the normal read() path. - let reorder_indices: Vec<usize> = projected_column_names - .iter() - .filter_map(|name| batch.schema().index_of(name).ok()) - .collect(); - if reorder_indices.len() == batch.num_columns() { - yield batch.project(&reorder_indices).map_err(|e| { - Error::UnexpectedError { - message: "Failed to reorder projected columns".to_string(), - source: Some(Box::new(e)), + let num_rows = batch.num_rows(); + let batch_schema = batch.schema(); + + // Build output columns using index mapping (field-ID-based) or by name. + let mut columns: Vec<Arc<dyn arrow_array::Array>> = Vec::with_capacity(target_schema.fields().len()); + for (i, target_field) in target_schema.fields().iter().enumerate() { + let source_col = if let Some(ref idx_map) = index_mapping { + let data_idx = idx_map[i]; + if data_idx == NULL_FIELD_INDEX { + None + } else { + // Find the column in the batch by the data field's name. + let data_field = &data_fields.as_ref().unwrap()[data_idx as usize]; + batch_schema + .index_of(data_field.name()) + .ok() + .map(|col_idx| batch.column(col_idx)) } - })?; - } else { - // Not all projected columns exist in this file (data evolution case), - // return as-is; the caller (merge_files_by_columns) handles missing columns. - yield batch; + } else if let Some(ref df) = data_fields { + // Identity mapping with data_fields present (e.g. renamed column). + // Use data field name (old name in parquet) at the same position. + batch_schema + .index_of(df[i].name()) + .ok() + .map(|col_idx| batch.column(col_idx)) + } else { + // No schema evolution — look up by target field name. + batch_schema + .index_of(target_field.name()) + .ok() + .map(|col_idx| batch.column(col_idx)) + }; + + match source_col { + Some(col) => { + if col.data_type() == target_field.data_type() { + columns.push(col.clone()); + } else { + // Type promotion: cast to target type. + let casted = cast(col, target_field.data_type()).map_err(|e| { + Error::UnexpectedError { + message: format!( + "Failed to cast column '{}' from {:?} to {:?}: {e}", + target_field.name(), + col.data_type(), + target_field.data_type() + ), + source: Some(Box::new(e)), + } + })?; + columns.push(casted); + } + } + None => { + // Column missing from file: fill with nulls. + let null_array = arrow_array::new_null_array(target_field.data_type(), num_rows); + columns.push(null_array); + } + } } + + let result = RecordBatch::try_new(target_schema.clone(), columns).map_err(|e| { + Error::UnexpectedError { + message: format!("Failed to build schema-evolved RecordBatch: {e}"), + source: Some(Box::new(e)), + } + })?; + yield result; } } .boxed()) @@ -274,8 +412,11 @@ fn read_single_file_stream( fn merge_files_by_columns( file_io: &FileIO, split: &DataSplit, + read_type: &[DataField], projected_column_names: &[String], table_field_names: &[String], + mut schema_manager: SchemaManager, + table_schema_id: i64, batch_size: Option<usize>, ) -> crate::Result<ArrowRecordBatchStream> { let data_files = split.data_files(); @@ -283,67 +424,88 @@ fn merge_files_by_columns( return Ok(futures::stream::empty().boxed()); } - // Determine which columns each file provides and resolve conflicts by max_sequence_number. - // column_name -> (file_index, max_sequence_number) - let mut column_source: HashMap<String, (usize, i64)> = HashMap::new(); - - for (file_idx, file_meta) in data_files.iter().enumerate() { - let file_columns: Vec<String> = if let Some(ref wc) = file_meta.write_cols { - wc.clone() - } else { - table_field_names.to_vec() - }; - - for col in &file_columns { - let entry = column_source - .entry(col.clone()) - .or_insert((file_idx, i64::MIN)); - if file_meta.max_sequence_number > entry.1 { - *entry = (file_idx, file_meta.max_sequence_number); - } - } - } - - // For each file, determine which projected columns to read from it. - // file_index -> Vec<column_name> - let mut file_read_columns: HashMap<usize, Vec<String>> = HashMap::new(); - for col_name in projected_column_names { - if let Some(&(file_idx, _)) = column_source.get(col_name) { - file_read_columns - .entry(file_idx) - .or_default() - .push(col_name.clone()); - } - } - - // For each projected column, record (file_index, column_name) for assembly. - let column_plan: Vec<(Option<usize>, String)> = projected_column_names - .iter() - .map(|col_name| { - let file_idx = column_source.get(col_name).map(|&(idx, _)| idx); - (file_idx, col_name.clone()) - }) - .collect(); - - // Collect which file indices we need to open streams for. - let active_file_indices: Vec<usize> = file_read_columns.keys().copied().collect(); - // Build owned data for the stream closure. let file_io = file_io.clone(); let split = split.clone(); let data_files: Vec<DataFileMeta> = data_files.to_vec(); + let read_type = read_type.to_vec(); let projected_column_names = projected_column_names.to_vec(); + let table_field_names = table_field_names.to_vec(); let output_batch_size = batch_size.unwrap_or(1024); Ok(try_stream! { + // Determine which columns each file provides and resolve conflicts by max_sequence_number. + // column_name -> (file_index, max_sequence_number) + let mut column_source: HashMap<String, (usize, i64)> = HashMap::new(); + + for (file_idx, file_meta) in data_files.iter().enumerate() { + let file_columns: Vec<String> = if let Some(ref wc) = file_meta.write_cols { + wc.clone() + } else if file_meta.schema_id != table_schema_id { + let file_schema = schema_manager.schema(file_meta.schema_id).await?; + file_schema.fields().iter().map(|f| f.name().to_string()).collect() + } else { + table_field_names.clone() + }; + + for col in &file_columns { + let entry = column_source + .entry(col.clone()) + .or_insert((file_idx, i64::MIN)); + if file_meta.max_sequence_number > entry.1 { + *entry = (file_idx, file_meta.max_sequence_number); + } + } + } + + // For each file, determine which projected columns to read from it. + // file_index -> Vec<column_name> + let mut file_read_columns: HashMap<usize, Vec<String>> = HashMap::new(); + for col_name in &projected_column_names { + if let Some(&(file_idx, _)) = column_source.get(col_name) { + file_read_columns + .entry(file_idx) + .or_default() + .push(col_name.clone()); + } + } + + // For each projected column, record (file_index, column_name) for assembly. + let column_plan: Vec<(Option<usize>, String)> = projected_column_names + .iter() + .map(|col_name| { + let file_idx = column_source.get(col_name).map(|&(idx, _)| idx); + (file_idx, col_name.clone()) + }) + .collect(); + + // Collect which file indices we need to open streams for. + let active_file_indices: Vec<usize> = file_read_columns.keys().copied().collect(); + // Open a stream for each active file. + // Build per-file read_type: only the DataFields this file is responsible for. let mut file_streams: HashMap<usize, ArrowRecordBatchStream> = HashMap::new(); for &file_idx in &active_file_indices { + let file_cols = file_read_columns.get(&file_idx).cloned().unwrap_or_default(); + let file_read_type: Vec<DataField> = file_cols + .iter() + .filter_map(|col_name| read_type.iter().find(|f| f.name() == col_name).cloned()) + .collect(); + + let file_meta = &data_files[file_idx]; + let data_fields: Option<Vec<DataField>> = if file_meta.schema_id != table_schema_id { + let data_schema = schema_manager.schema(file_meta.schema_id).await?; + Some(data_schema.fields().to_vec()) + } else { + None + }; + let stream = read_single_file_stream( file_io.clone(), split.clone(), data_files[file_idx].clone(), - projected_column_names.clone(), + file_read_type, + data_fields, batch_size, None, )?; diff --git a/crates/paimon/src/arrow/schema_evolution.rs b/crates/paimon/src/arrow/schema_evolution.rs new file mode 100644 index 0000000..ea9450b --- /dev/null +++ b/crates/paimon/src/arrow/schema_evolution.rs @@ -0,0 +1,151 @@ +// 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. + +//! Schema evolution utilities for mapping between table schema and data file schema. +//! +//! Reference: [org.apache.paimon.schema.SchemaEvolutionUtil](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/schema/SchemaEvolutionUtil.java) + +use crate::spec::DataField; +use std::collections::HashMap; + +/// Sentinel value indicating a field does not exist in the data schema. +pub const NULL_FIELD_INDEX: i32 = -1; + +/// Create index mapping from table fields to underlying data fields using field IDs. +/// +/// For example, the table and data fields are as follows: +/// - table fields: `1->c, 6->b, 3->a` +/// - data fields: `1->a, 3->c` +/// +/// We get the index mapping `[0, -1, 1]`, where: +/// - `0` is the index of table field `1->c` in data fields +/// - `-1` means field `6->b` does not exist in data fields +/// - `1` is the index of table field `3->a` in data fields +/// +/// Returns `None` if the mapping is identity (no evolution needed). +/// +/// Reference: [SchemaEvolutionUtil.createIndexMapping](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/schema/SchemaEvolutionUtil.java) +pub fn create_index_mapping( + table_fields: &[DataField], + data_fields: &[DataField], +) -> Option<Vec<i32>> { + let mut field_id_to_index: HashMap<i32, i32> = HashMap::with_capacity(data_fields.len()); + for (i, field) in data_fields.iter().enumerate() { + field_id_to_index.insert(field.id(), i as i32); + } + + let mut index_mapping = Vec::with_capacity(table_fields.len()); + for field in table_fields { + let data_index = field_id_to_index + .get(&field.id()) + .copied() + .unwrap_or(NULL_FIELD_INDEX); + index_mapping.push(data_index); + } + + // Check if mapping is identity (no evolution needed). + let is_identity = index_mapping.len() == data_fields.len() + && index_mapping + .iter() + .enumerate() + .all(|(i, &idx)| idx == i as i32); + + if is_identity { + None + } else { + Some(index_mapping) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spec::{DataType, IntType, VarCharType}; + + fn field(id: i32, name: &str) -> DataField { + DataField::new(id, name.to_string(), DataType::Int(IntType::new())) + } + + #[test] + fn test_identity_mapping() { + let table_fields = vec![field(0, "a"), field(1, "b"), field(2, "c")]; + let data_fields = vec![field(0, "a"), field(1, "b"), field(2, "c")]; + assert_eq!(create_index_mapping(&table_fields, &data_fields), None); + } + + #[test] + fn test_added_column() { + // Table has 3 fields, data file only has the first 2 + let table_fields = vec![field(0, "a"), field(1, "b"), field(2, "c")]; + let data_fields = vec![field(0, "a"), field(1, "b")]; + assert_eq!( + create_index_mapping(&table_fields, &data_fields), + Some(vec![0, 1, -1]) + ); + } + + #[test] + fn test_reordered_fields() { + let table_fields = vec![field(1, "c"), field(6, "b"), field(3, "a")]; + let data_fields = vec![field(1, "a"), field(3, "c")]; + assert_eq!( + create_index_mapping(&table_fields, &data_fields), + Some(vec![0, -1, 1]) + ); + } + + #[test] + fn test_renamed_column() { + // Field ID stays the same even if name changed + let table_fields = vec![field(0, "id"), field(1, "new_name")]; + let data_fields = vec![field(0, "id"), field(1, "old_name")]; + // Identity mapping since field IDs match positionally + assert_eq!(create_index_mapping(&table_fields, &data_fields), None); + } + + #[test] + fn test_empty_data_fields() { + let table_fields = vec![field(0, "a"), field(1, "b")]; + let data_fields: Vec<DataField> = vec![]; + assert_eq!( + create_index_mapping(&table_fields, &data_fields), + Some(vec![-1, -1]) + ); + } + + #[test] + fn test_type_promotion_same_mapping() { + // Type promotion doesn't affect index mapping — only field IDs matter + let table_fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "name".to_string(), + DataType::VarChar(VarCharType::string_type()), + ), + ]; + let data_fields = vec![ + DataField::new(0, "id".to_string(), DataType::Int(IntType::new())), + DataField::new( + 1, + "name".to_string(), + DataType::VarChar(VarCharType::string_type()), + ), + ]; + assert_eq!(create_index_mapping(&table_fields, &data_fields), None); + } +} diff --git a/crates/paimon/src/spec/types.rs b/crates/paimon/src/spec/types.rs index e0ba04a..1bb4566 100644 --- a/crates/paimon/src/spec/types.rs +++ b/crates/paimon/src/spec/types.rs @@ -227,6 +227,10 @@ impl ArrayType { pub fn family(&self) -> DataTypeFamily { DataTypeFamily::CONSTRUCTED | DataTypeFamily::COLLECTION } + + pub fn element_type(&self) -> &DataType { + &self.element_type + } } /// BigIntType for paimon. @@ -1348,6 +1352,14 @@ impl MapType { pub fn family(&self) -> DataTypeFamily { DataTypeFamily::CONSTRUCTED | DataTypeFamily::COLLECTION } + + pub fn key_type(&self) -> &DataType { + &self.key_type + } + + pub fn value_type(&self) -> &DataType { + &self.value_type + } } /// MultisetType for paimon. @@ -1381,6 +1393,10 @@ impl MultisetType { pub fn family(&self) -> DataTypeFamily { DataTypeFamily::CONSTRUCTED | DataTypeFamily::COLLECTION } + + pub fn element_type(&self) -> &DataType { + &self.element_type + } } /// RowType for paimon. diff --git a/crates/paimon/src/table/mod.rs b/crates/paimon/src/table/mod.rs index 7c3f9a5..e98276c 100644 --- a/crates/paimon/src/table/mod.rs +++ b/crates/paimon/src/table/mod.rs @@ -19,6 +19,7 @@ pub(crate) mod bin_pack; mod read_builder; +pub(crate) mod schema_manager; mod snapshot_manager; mod source; mod table_scan; @@ -27,6 +28,7 @@ use crate::Result; use arrow_array::RecordBatch; use futures::stream::BoxStream; pub use read_builder::{ReadBuilder, TableRead}; +pub use schema_manager::SchemaManager; pub use snapshot_manager::SnapshotManager; pub use source::{DataSplit, DataSplitBuilder, DeletionFile, PartitionBucket, Plan}; pub use table_scan::TableScan; @@ -42,6 +44,7 @@ pub struct Table { identifier: Identifier, location: String, schema: TableSchema, + schema_manager: SchemaManager, } #[allow(dead_code)] @@ -53,11 +56,13 @@ impl Table { location: String, schema: TableSchema, ) -> Self { + let schema_manager = SchemaManager::new(file_io.clone(), location.clone()); Self { file_io, identifier, location, schema, + schema_manager, } } @@ -81,6 +86,11 @@ impl Table { &self.file_io } + /// Get the SchemaManager for this table. + pub fn schema_manager(&self) -> &SchemaManager { + &self.schema_manager + } + /// Create a read builder for scan/read. /// /// Reference: [pypaimon FileStoreTable.new_read_builder](https://github.com/apache/paimon/blob/release-1.3/paimon-python/pypaimon/table/file_store_table.py). diff --git a/crates/paimon/src/table/read_builder.rs b/crates/paimon/src/table/read_builder.rs index 65c89b8..eabfbbf 100644 --- a/crates/paimon/src/table/read_builder.rs +++ b/crates/paimon/src/table/read_builder.rs @@ -166,8 +166,12 @@ impl<'a> TableRead<'a> { }); } - let reader = - ArrowReaderBuilder::new(self.table.file_io.clone()).build(self.read_type().to_vec()); + let reader = ArrowReaderBuilder::new( + self.table.file_io.clone(), + self.table.schema_manager().clone(), + self.table.schema().id(), + ) + .build(self.read_type().to_vec()); if data_evolution { reader.read_data_evolution(data_splits, self.table.schema.fields()) diff --git a/crates/paimon/src/table/schema_manager.rs b/crates/paimon/src/table/schema_manager.rs new file mode 100644 index 0000000..7b937e0 --- /dev/null +++ b/crates/paimon/src/table/schema_manager.rs @@ -0,0 +1,81 @@ +// 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. + +//! Schema manager for reading versioned table schemas. +//! +//! Reference: [org.apache.paimon.schema.SchemaManager](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java) + +use crate::io::FileIO; +use crate::spec::TableSchema; +use std::collections::HashMap; + +const SCHEMA_DIR: &str = "schema"; +const SCHEMA_PREFIX: &str = "schema-"; + +/// Manager for versioned table schema files. +/// +/// Each table stores schema versions as JSON files under `{table_path}/schema/schema-{id}`. +/// When a schema evolution occurs (e.g. ADD COLUMN, ALTER COLUMN TYPE), a new schema file +/// is written with an incremented ID. Data files record which schema they were written with +/// via `DataFileMeta.schema_id`. +/// +/// Reference: [org.apache.paimon.schema.SchemaManager](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java) +#[derive(Debug, Clone)] +pub struct SchemaManager { + file_io: FileIO, + table_path: String, + /// Cache of loaded schemas by ID. + cache: HashMap<i64, TableSchema>, +} + +impl SchemaManager { + pub fn new(file_io: FileIO, table_path: String) -> Self { + Self { + file_io, + table_path, + cache: HashMap::new(), + } + } + + /// Path to the schema directory (e.g. `{table_path}/schema`). + fn schema_directory(&self) -> String { + format!("{}/{}", self.table_path.trim_end_matches('/'), SCHEMA_DIR) + } + + /// Path to a specific schema file (e.g. `{table_path}/schema/schema-0`). + fn schema_path(&self, schema_id: i64) -> String { + format!("{}/{}{}", self.schema_directory(), SCHEMA_PREFIX, schema_id) + } + + /// Load a schema by ID. Returns cached version if available. + /// + /// Reference: [SchemaManager.schema(long)](https://github.com/apache/paimon/blob/release-1.3/paimon-core/src/main/java/org/apache/paimon/schema/SchemaManager.java) + pub async fn schema(&mut self, schema_id: i64) -> crate::Result<&TableSchema> { + if !self.cache.contains_key(&schema_id) { + let path = self.schema_path(schema_id); + let input = self.file_io.new_input(&path)?; + let bytes = input.read().await?; + let schema: TableSchema = + serde_json::from_slice(&bytes).map_err(|e| crate::Error::DataInvalid { + message: format!("Failed to parse schema file: {path}"), + source: Some(Box::new(e)), + })?; + self.cache.insert(schema_id, schema); + } + Ok(self.cache.get(&schema_id).unwrap()) + } +} diff --git a/dev/spark/provision.py b/dev/spark/provision.py index cdc538a..0277255 100644 --- a/dev/spark/provision.py +++ b/dev/spark/provision.py @@ -260,6 +260,147 @@ def main(): ) spark.sql("DROP TABLE data_evolution_updates") + # ===== Schema Evolution: Add Column ===== + # Old files have (id, name); after ALTER TABLE ADD COLUMNS, new files have (id, name, age). + # Reader must fill nulls for 'age' when reading old files. + spark.sql( + """ + CREATE TABLE IF NOT EXISTS schema_evolution_add_column ( + id INT, + name STRING + ) USING paimon + """ + ) + spark.sql( + "INSERT INTO schema_evolution_add_column VALUES (1, 'alice'), (2, 'bob')" + ) + spark.sql("ALTER TABLE schema_evolution_add_column ADD COLUMNS (age INT)") + spark.sql( + "INSERT INTO schema_evolution_add_column VALUES (3, 'carol', 30), (4, 'dave', 40)" + ) + + # ===== Schema Evolution: Type Promotion (INT -> BIGINT) ===== + # Old files have value as INT; after ALTER TABLE, new files have value as BIGINT. + # Reader must cast INT to BIGINT when reading old files. + spark.sql( + """ + CREATE TABLE IF NOT EXISTS schema_evolution_type_promotion ( + id INT, + value INT + ) USING paimon + """ + ) + spark.sql( + "INSERT INTO schema_evolution_type_promotion VALUES (1, 100), (2, 200)" + ) + spark.sql( + "ALTER TABLE schema_evolution_type_promotion ALTER COLUMN value TYPE BIGINT" + ) + spark.sql( + "INSERT INTO schema_evolution_type_promotion VALUES (3, 3000000000)" + ) + + # ===== Data Evolution + Schema Evolution: Add Column ===== + # Combines data-evolution (row-tracking + MERGE INTO) with ALTER TABLE ADD COLUMNS. + # Old files lack the new column; MERGE INTO produces partial-column files. + # Reader must fill nulls for missing columns AND merge columns across files. + spark.sql( + """ + CREATE TABLE IF NOT EXISTS data_evolution_add_column ( + id INT, + name STRING, + value INT + ) USING paimon + TBLPROPERTIES ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true' + ) + """ + ) + spark.sql( + """ + INSERT INTO data_evolution_add_column VALUES + (1, 'alice', 100), + (2, 'bob', 200) + """ + ) + spark.sql("ALTER TABLE data_evolution_add_column ADD COLUMNS (extra STRING)") + spark.sql( + """ + INSERT INTO data_evolution_add_column VALUES + (3, 'carol', 300, 'new'), + (4, 'dave', 400, 'new') + """ + ) + # MERGE INTO to trigger merge_files_by_columns with schema evolution. + spark.sql( + """ + CREATE TABLE IF NOT EXISTS data_evolution_add_column_updates ( + id INT, + name STRING + ) USING paimon + """ + ) + spark.sql( + "INSERT INTO data_evolution_add_column_updates VALUES (1, 'alice-v2')" + ) + spark.sql( + """ + MERGE INTO data_evolution_add_column t + USING data_evolution_add_column_updates s + ON t.id = s.id + WHEN MATCHED THEN UPDATE SET t.name = s.name + """ + ) + spark.sql("DROP TABLE data_evolution_add_column_updates") + + # ===== Data Evolution + Schema Evolution: Type Promotion ===== + # Combines data-evolution with ALTER TABLE ALTER COLUMN TYPE (INT -> BIGINT). + # Old files have INT; new files have BIGINT. MERGE INTO updates some rows. + # Reader must cast old INT columns to BIGINT AND merge columns across files. + spark.sql( + """ + CREATE TABLE IF NOT EXISTS data_evolution_type_promotion ( + id INT, + value INT + ) USING paimon + TBLPROPERTIES ( + 'row-tracking.enabled' = 'true', + 'data-evolution.enabled' = 'true' + ) + """ + ) + spark.sql( + "INSERT INTO data_evolution_type_promotion VALUES (1, 100), (2, 200)" + ) + spark.sql( + "ALTER TABLE data_evolution_type_promotion ALTER COLUMN value TYPE BIGINT" + ) + spark.sql( + "INSERT INTO data_evolution_type_promotion VALUES (3, 3000000000)" + ) + # MERGE INTO to trigger merge_files_by_columns with type promotion. + spark.sql( + """ + CREATE TABLE IF NOT EXISTS data_evolution_type_promotion_updates ( + id INT, + value BIGINT + ) USING paimon + """ + ) + spark.sql( + "INSERT INTO data_evolution_type_promotion_updates VALUES (1, 999)" + ) + spark.sql( + """ + MERGE INTO data_evolution_type_promotion t + USING data_evolution_type_promotion_updates s + ON t.id = s.id + WHEN MATCHED THEN UPDATE SET t.value = s.value + """ + ) + spark.sql("DROP TABLE data_evolution_type_promotion_updates") + if __name__ == "__main__": main()
