diff --git a/crates/integrations/datafusion/src/merge_into.rs b/crates/integrations/datafusion/src/merge_into.rs index b73b1555..2620d9d4 100644 --- a/crates/integrations/datafusion/src/merge_into.rs +++ b/crates/integrations/datafusion/src/merge_into.rs @@ -26,7 +26,9 @@ use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; -use datafusion::arrow::array::{Array, Int32Array, RecordBatch, UInt32Array, UInt64Array}; +use datafusion::arrow::array::{ + Array, ArrayRef, Int32Array, RecordBatch, UInt32Array, UInt64Array, +}; use datafusion::arrow::compute; use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema}; use datafusion::error::{DataFusionError, Result as DFResult}; @@ -37,7 +39,7 @@ use datafusion::sql::sqlparser::ast::{ }; use futures::TryStreamExt; -use paimon::spec::{datums_to_binary_row, extract_datum_from_arrow, CoreOptions}; +use paimon::spec::{datums_to_binary_row, extract_datum_from_arrow, CoreOptions, DataField}; use paimon::table::{CopyOnWriteMergeWriter, DataSplitBuilder, Table, WriteBuilder}; use crate::error::to_datafusion_error; @@ -569,13 +571,6 @@ async fn execute_cow_merge_inner( // Handle NOT MATCHED → INSERT if !clauses.inserts.is_empty() { - let table_fields: Vec = table - .schema() - .fields() - .iter() - .map(|f| f.name().to_string()) - .collect(); - let insert_sql = if has_target_data { format!( "SELECT {s_alias}.* FROM {source_ref} AS {s_alias} \ @@ -595,7 +590,7 @@ async fn execute_cow_merge_inner( &clauses.inserts, s_alias, &[], - &table_fields, + table.schema().fields(), temp_tracker, ) .await?; @@ -734,13 +729,6 @@ async fn execute_merge_into_once( injected_columns.push(format!("__upd_{col}")); } } - // Table schema field names for reordering INSERT columns - let table_fields: Vec = table - .schema() - .fields() - .iter() - .map(|f| f.name().to_string()) - .collect(); let mut temp_tracker = TempTableTracker::new(ctx); let insert_batches = build_insert_batches( ctx, @@ -748,7 +736,7 @@ async fn execute_merge_into_once( &parsed.inserts, s_alias, &injected_columns, - &table_fields, + table.schema().fields(), &mut temp_tracker, ) .await?; @@ -854,7 +842,7 @@ async fn build_insert_batches( inserts: &[MergeInsertClause], s_alias: &str, injected_columns: &[String], - table_fields: &[String], + table_fields: &[DataField], temp_tracker: &mut TempTableTracker<'_>, ) -> DFResult> { if not_matched_batches.is_empty() || not_matched_batches.iter().all(|b| b.num_rows() == 0) { @@ -883,7 +871,7 @@ async fn build_insert_batches_inner( inserts: &[MergeInsertClause], s_alias: &str, tmp_name: &str, - table_fields: &[String], + table_fields: &[DataField], ) -> DFResult> { let mut all_batches = Vec::new(); let mut consumed_predicates: Vec = Vec::new(); @@ -910,12 +898,61 @@ async fn build_insert_batches_inner( let sql = format!("SELECT {select_clause} FROM {tmp_name} AS {s_alias}{where_clause}"); let batches = ctx.ctx().sql(&sql).await?.collect().await?; - all_batches.extend(batches); + for batch in batches { + all_batches.push(normalize_insert_batch_to_table_schema( + &batch, + table_fields, + )?); + } } Ok(all_batches) } +fn normalize_insert_batch_to_table_schema( + batch: &RecordBatch, + table_fields: &[DataField], +) -> DFResult { + if batch.num_columns() != table_fields.len() { + return Err(DataFusionError::Plan(format!( + "MERGE INSERT output has {} columns but target table has {}", + batch.num_columns(), + table_fields.len() + ))); + } + + let target_schema = + paimon::arrow::build_target_arrow_schema(table_fields).map_err(to_datafusion_error)?; + let mut columns = Vec::with_capacity(table_fields.len()); + + for (target_idx, field) in table_fields.iter().enumerate() { + let column = batch.column(target_idx).clone(); + let target_type = target_schema.field(target_idx).data_type(); + let column = cast_insert_column(field.name(), column, target_type)?; + columns.push(column); + } + + RecordBatch::try_new(target_schema, columns).map_err(DataFusionError::from) +} + +fn cast_insert_column( + name: &str, + column: ArrayRef, + target_type: &ArrowDataType, +) -> DFResult { + if column.data_type() == target_type { + return Ok(column); + } + + compute::cast(column.as_ref(), target_type).map_err(|e| { + DataFusionError::Plan(format!( + "Cannot cast MERGE INSERT column '{name}' from {:?} to {:?}: {e}", + column.data_type(), + target_type + )) + }) +} + /// Remove injected columns from batches, keeping only source columns. fn strip_non_source_columns( batches: &[RecordBatch], @@ -946,7 +983,7 @@ fn strip_non_source_columns( /// When the INSERT specifies explicit columns (`INSERT (col2, col1) VALUES (expr2, expr1)`), /// the output must be reordered to match the table schema so that `write_arrow_batch` /// (which reads columns by positional index) maps them correctly. -fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[String]) -> String { +fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[DataField]) -> String { if ins.columns.is_empty() && ins.value_exprs.is_empty() { "*".to_string() } else { @@ -962,11 +999,11 @@ fn insert_select_clause(ins: &MergeInsertClause, table_fields: &[String]) -> Str table_fields .iter() .map(|field| { - let key = field.to_lowercase(); + let key = field.name().to_lowercase(); match col_expr_map.get(&key) { - Some(expr) => format!("{expr} AS {}", quote_identifier(field)), + Some(expr) => format!("{expr} AS {}", quote_identifier(field.name())), // Column not in INSERT list — fill with NULL - None => format!("NULL AS {}", quote_identifier(field)), + None => format!("NULL AS {}", quote_identifier(field.name())), } }) .collect::>() @@ -1546,7 +1583,7 @@ mod tests { use datafusion::sql::sqlparser::parser::Parser; use paimon::catalog::{Catalog, Identifier}; use paimon::io::FileIOBuilder; - use paimon::spec::{DataType, IntType, Schema as PaimonSchema, TableSchema}; + use paimon::spec::{DataField, DataType, IntType, Schema as PaimonSchema, TableSchema}; use paimon::{CatalogOptions, FileSystemCatalog, Options}; use tempfile::TempDir; @@ -1613,6 +1650,42 @@ mod tests { } } + #[test] + fn test_normalize_merge_insert_batch_uses_position() { + let table_fields = vec![ + DataField::new(0, "a".to_string(), DataType::Int(IntType::new())), + DataField::new(1, "b".to_string(), DataType::Int(IntType::new())), + ]; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("b", ArrowDataType::Int32, false), + Field::new("x", ArrowDataType::Int32, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![100])), + Arc::new(Int32Array::from(vec![7])), + ], + ) + .unwrap(); + + let normalized = normalize_insert_batch_to_table_schema(&batch, &table_fields).unwrap(); + let first = normalized + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let second = normalized + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + + assert_eq!(normalized.schema().field(0).name(), "a"); + assert_eq!(normalized.schema().field(1).name(), "b"); + assert_eq!(first.value(0), 100); + assert_eq!(second.value(0), 7); + } + #[test] fn test_source_partition_pruning_requires_partition_equality() { let merge = parse_merge( diff --git a/crates/paimon/src/arrow/format/blob.rs b/crates/paimon/src/arrow/format/blob.rs index 65bd2f3d..136170ec 100644 --- a/crates/paimon/src/arrow/format/blob.rs +++ b/crates/paimon/src/arrow/format/blob.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{FilePredicates, FormatFileReader, FormatFileWriter}; +use super::{FilePredicates, FormatFileReader, FormatFileWriter, FormatWriteResult}; use crate::arrow::build_target_arrow_schema; use crate::io::{FileRead, FileWrite}; use crate::spec::{BlobDescriptor, DataField, DataType}; @@ -725,7 +725,7 @@ impl FormatFileWriter for BlobFormatWriter { Ok(()) } - async fn close(mut self: Box) -> crate::Result { + async fn close(mut self: Box) -> crate::Result { let index_bytes = encode_delta_varints_write(&self.lengths); let index_length = index_bytes.len() as i32; @@ -739,7 +739,7 @@ impl FormatFileWriter for BlobFormatWriter { let total = self.bytes_written + index_length as u64 + BLOB_FOOTER_SIZE; self.writer.close().await?; - Ok(total) + Ok(FormatWriteResult::new(total)) } } diff --git a/crates/paimon/src/arrow/format/mod.rs b/crates/paimon/src/arrow/format/mod.rs index 2a2fcbde..cd981b86 100644 --- a/crates/paimon/src/arrow/format/mod.rs +++ b/crates/paimon/src/arrow/format/mod.rs @@ -19,7 +19,7 @@ mod avro; pub(crate) mod blob; mod mosaic; mod orc; -mod parquet; +pub(crate) mod parquet; mod row; mod shredding; #[cfg(feature = "vortex")] @@ -29,6 +29,7 @@ mod vortex; pub(crate) use parquet::ParquetFormatWriter; use crate::io::{FileRead, OutputFile}; +use crate::spec::stats::BinaryTableStats; use crate::spec::{DataField, Predicate}; use crate::table::{ArrowRecordBatchStream, RowRange}; use crate::Error; @@ -102,8 +103,40 @@ pub(crate) trait FormatFileWriter: Send { async fn flush(&mut self) -> crate::Result<()>; /// Flush and close the writer, finalizing the file on storage. - /// Returns the total number of bytes written. - async fn close(self: Box) -> crate::Result; + async fn close(self: Box) -> crate::Result; +} + +pub(crate) struct FormatWriteResult { + pub(crate) file_size: u64, + pub(crate) value_stats: Option, +} + +pub(crate) struct FormatValueStats { + pub(crate) stats: BinaryTableStats, + pub(crate) columns: Option>, +} + +impl FormatWriteResult { + pub(crate) fn new(file_size: u64) -> Self { + Self { + file_size, + value_stats: None, + } + } + + pub(crate) fn with_value_stats( + file_size: u64, + value_stats: BinaryTableStats, + columns: Option>, + ) -> Self { + Self { + file_size, + value_stats: Some(FormatValueStats { + stats: value_stats, + columns, + }), + } + } } /// Create a format reader based on the file extension. @@ -180,6 +213,7 @@ pub(crate) async fn create_format_writer( output, compression, zstd_level, + format_options.cloned().unwrap_or_default(), )); shredding::ShreddingFormatWriter::create( writer_factory, diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 5da5eeea..cf55e0bc 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -16,10 +16,14 @@ // under the License. use super::shredding::PhysicalFormatWriterFactory; -use super::{FilePredicates, FormatFileReader, FormatFileWriter}; +use super::{FilePredicates, FormatFileReader, FormatFileWriter, FormatWriteResult}; use crate::arrow::filtering::{predicates_may_match_with_schema, StatsAccessor}; use crate::io::{FileRead, OutputFile}; -use crate::spec::{DataField, DataType, Datum, Predicate, PredicateOperator}; +use crate::spec::stats::BinaryTableStats; +use crate::spec::{ + BinaryRowBuilder, CoreOptions, DataField, DataType, Datum, MetadataStatsMode, Predicate, + PredicateOperator, +}; use crate::table::{ArrowRecordBatchStream, RowRange}; use crate::Error; use arrow_array::{BooleanArray, RecordBatch}; @@ -39,6 +43,7 @@ use parquet::file::metadata::{ use parquet::file::page_index::column_index::ColumnIndexMetaData; use parquet::file::properties::WriterProperties; use parquet::file::statistics::Statistics as ParquetStatistics; +use std::cmp::Ordering; use std::collections::HashMap; use std::ops::Range; use std::sync::Arc; @@ -49,20 +54,30 @@ pub(crate) struct ParquetFormatReader; /// Streams data directly to storage via `AsyncArrowWriter` + opendal. pub(crate) struct ParquetFormatWriter { inner: AsyncArrowWriter>, + write_fields: Option>, + stats_modes: Option>, + stats_dense_store: bool, } pub(crate) struct ParquetPhysicalWriterFactory { output: OutputFile, compression: String, zstd_level: i32, + format_options: HashMap, } impl ParquetPhysicalWriterFactory { - pub(crate) fn new(output: &OutputFile, compression: &str, zstd_level: i32) -> Self { + pub(crate) fn new( + output: &OutputFile, + compression: &str, + zstd_level: i32, + format_options: HashMap, + ) -> Self { Self { output: output.clone(), compression: compression.to_string(), zstd_level, + format_options, } } } @@ -72,11 +87,18 @@ impl PhysicalFormatWriterFactory for ParquetPhysicalWriterFactory { async fn create_writer( &mut self, schema: arrow_schema::SchemaRef, - _write_fields: Option<&[DataField]>, + write_fields: Option<&[DataField]>, ) -> crate::Result> { Ok(Box::new( - ParquetFormatWriter::new(&self.output, schema, &self.compression, self.zstd_level) - .await?, + ParquetFormatWriter::new( + &self.output, + schema, + &self.compression, + self.zstd_level, + write_fields, + &self.format_options, + ) + .await?, )) } } @@ -87,11 +109,22 @@ impl ParquetFormatWriter { schema: arrow_schema::SchemaRef, compression: &str, zstd_level: i32, + write_fields: Option<&[DataField]>, + format_options: &HashMap, ) -> crate::Result { let async_write = output.async_writer().await?; let codec = parse_compression(compression, zstd_level); let inner = create_parquet_arrow_writer(async_write, schema, codec)?; - Ok(Self { inner }) + let core_options = CoreOptions::new(format_options); + let stats_modes = write_fields + .map(|fields| core_options.metadata_stats_modes(fields.iter().map(DataField::name))) + .transpose()?; + Ok(Self { + inner, + write_fields: write_fields.map(|fields| fields.to_vec()), + stats_modes, + stats_dense_store: core_options.metadata_stats_dense_store(), + }) } } @@ -154,15 +187,27 @@ impl FormatFileWriter for ParquetFormatWriter { }) } - async fn close(mut self: Box) -> crate::Result { - self.inner + async fn close(mut self: Box) -> crate::Result { + let metadata = self + .inner .finish() .await .map_err(|e| crate::Error::DataInvalid { message: format!("Failed to close parquet writer: {e}"), source: None, })?; - Ok(self.inner.bytes_written() as u64) + let file_size = self.inner.bytes_written() as u64; + if let (Some(write_fields), Some(stats_modes)) = (&self.write_fields, &self.stats_modes) { + let (value_stats, value_stats_cols) = + extract_value_stats(&metadata, write_fields, stats_modes, self.stats_dense_store); + Ok(FormatWriteResult::with_value_stats( + file_size, + value_stats, + value_stats_cols, + )) + } else { + Ok(FormatWriteResult::new(file_size)) + } } } @@ -625,6 +670,272 @@ fn build_row_group_column_indices( .collect() } +fn extract_value_stats( + metadata: &ParquetMetaData, + write_fields: &[DataField], + stats_modes: &[MetadataStatsMode], + stats_dense_store: bool, +) -> (BinaryTableStats, Option>) { + debug_assert_eq!(write_fields.len(), stats_modes.len()); + let row_groups = metadata.row_groups(); + let column_indices = row_groups + .first() + .map(|row_group| build_row_group_column_indices(row_group.columns(), write_fields)) + .unwrap_or_else(|| vec![None; write_fields.len()]); + let mut column_names = Vec::new(); + let mut column_types = Vec::new(); + let mut min_datums = Vec::new(); + let mut max_datums = Vec::new(); + let mut null_counts = Vec::new(); + + for (field_idx, field) in write_fields.iter().enumerate() { + let mode = stats_modes + .get(field_idx) + .copied() + .unwrap_or(MetadataStatsMode::None); + let column_stats = if mode == MetadataStatsMode::None { + None + } else { + column_indices + .get(field_idx) + .copied() + .flatten() + .and_then(|column_idx| { + extract_column_value_stats(row_groups, column_idx, field.data_type(), mode) + }) + }; + + // Non-dense stats stay aligned with write_fields, including unavailable stats. + match column_stats { + Some((min_datum, max_datum, null_count)) => { + if stats_dense_store { + column_names.push(field.name().to_string()); + } + column_types.push(field.data_type().clone()); + min_datums.push(min_datum); + max_datums.push(max_datum); + null_counts.push(null_count); + } + None if !stats_dense_store => { + column_types.push(field.data_type().clone()); + min_datums.push(None); + max_datums.push(None); + null_counts.push(None); + } + None => {} + } + } + + let stats = if column_types.is_empty() { + BinaryTableStats::empty() + } else { + binary_table_stats_from_datums(&column_types, &min_datums, &max_datums, null_counts) + }; + // Java omits the dense mapping when stats already cover every write field. + let value_stats_cols = if stats_dense_store && column_types.len() != write_fields.len() { + Some(column_names) + } else { + None + }; + (stats, value_stats_cols) +} + +fn extract_column_value_stats( + row_groups: &[RowGroupMetaData], + column_idx: usize, + data_type: &DataType, + mode: MetadataStatsMode, +) -> Option<(Option, Option, Option)> { + let collect_min_max = matches!( + mode, + MetadataStatsMode::Full | MetadataStatsMode::Truncate(_) + ) && supports_manifest_min_max(data_type); + let mut min_datum: Option = None; + let mut max_datum: Option = None; + let mut min_complete = true; + let mut max_complete = true; + let mut null_count = Some(0_i64); + let mut has_stats = false; + + for row_group in row_groups { + let Some(stats) = row_group.column(column_idx).statistics() else { + min_complete = false; + max_complete = false; + null_count = None; + continue; + }; + has_stats = true; + + match stats + .null_count_opt() + .and_then(|count| i64::try_from(count).ok()) + { + Some(count) => { + if let Some(total) = null_count.as_mut() { + *total += count; + } + } + None => null_count = None, + } + + let row_group_all_null = stats.null_count_opt() == Some(row_group.num_rows().max(0) as u64); + if !collect_min_max || row_group_all_null { + continue; + } + + match parquet_stats_to_datum(stats, data_type, true) { + Some(candidate) => { + if let Some(current) = &min_datum { + match candidate.partial_cmp(current) { + Some(Ordering::Less) => min_datum = Some(candidate), + Some(Ordering::Equal | Ordering::Greater) => {} + None => min_complete = false, + } + } else { + min_datum = Some(candidate); + } + } + None => min_complete = false, + } + + match parquet_stats_to_datum(stats, data_type, false) { + Some(candidate) => { + if let Some(current) = &max_datum { + match candidate.partial_cmp(current) { + Some(Ordering::Greater) => max_datum = Some(candidate), + Some(Ordering::Less | Ordering::Equal) => {} + None => max_complete = false, + } + } else { + max_datum = Some(candidate); + } + } + None => max_complete = false, + } + } + + if !has_stats { + return None; + } + if !min_complete { + min_datum = None; + } + if !max_complete { + max_datum = None; + } + let (min_datum, max_datum) = apply_stats_mode(data_type, mode, min_datum, max_datum); + if min_datum.is_none() && max_datum.is_none() && null_count.is_none() { + None + } else { + Some((min_datum, max_datum, null_count)) + } +} + +fn supports_manifest_min_max(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Char(_) + | DataType::VarChar(_) + | DataType::Decimal(_) + | DataType::Double(_) + | DataType::Float(_) + | DataType::Date(_) + | DataType::Time(_) + | DataType::LocalZonedTimestamp(_) + | DataType::Timestamp(_) + ) +} + +fn apply_stats_mode( + data_type: &DataType, + mode: MetadataStatsMode, + min_datum: Option, + max_datum: Option, +) -> (Option, Option) { + let MetadataStatsMode::Truncate(length) = mode else { + return (min_datum, max_datum); + }; + match data_type { + DataType::Char(_) | DataType::VarChar(_) => { + let min = min_datum.map(|datum| truncate_string_min_datum(datum, length)); + let max = match max_datum { + Some(datum) => match truncate_string_max_datum(datum, length) { + Some(max) => Some(max), + None => return (None, None), + }, + None => None, + }; + (min, max) + } + _ => (min_datum, max_datum), + } +} + +fn truncate_string_min_datum(datum: Datum, length: usize) -> Datum { + match datum { + Datum::String(value) => Datum::String(truncate_string_min(&value, length)), + other => other, + } +} + +fn truncate_string_max_datum(datum: Datum, length: usize) -> Option { + match datum { + Datum::String(value) => truncate_string_max(&value, length).map(Datum::String), + other => Some(other), + } +} + +fn truncate_string_min(value: &str, length: usize) -> String { + value.chars().take(length).collect() +} + +fn truncate_string_max(value: &str, length: usize) -> Option { + let char_count = value.chars().count(); + if char_count <= length { + return Some(value.to_string()); + } + + let mut chars: Vec = value.chars().take(length).collect(); + for idx in (0..chars.len()).rev() { + if let Some(next) = char::from_u32(chars[idx] as u32 + 1) { + chars.truncate(idx); + chars.push(next); + return Some(chars.into_iter().collect()); + } + } + None +} + +fn binary_table_stats_from_datums( + column_types: &[DataType], + min_datums: &[Option], + max_datums: &[Option], + null_counts: Vec>, +) -> BinaryTableStats { + let mut min_builder = BinaryRowBuilder::new(column_types.len() as i32); + let mut max_builder = BinaryRowBuilder::new(column_types.len() as i32); + for (pos, data_type) in column_types.iter().enumerate() { + match &min_datums[pos] { + Some(datum) => min_builder.write_datum(pos, datum, data_type), + None => min_builder.set_null_at(pos), + } + match &max_datums[pos] { + Some(datum) => max_builder.write_datum(pos, datum, data_type), + None => max_builder.set_null_at(pos), + } + } + BinaryTableStats::new( + min_builder.build_serialized(), + max_builder.build_serialized(), + null_counts, + ) +} + // --------------------------------------------------------------------------- // Page-index (ColumnIndex / OffsetIndex) pruning // --------------------------------------------------------------------------- @@ -984,17 +1295,33 @@ fn parquet_stats_to_datum( .copied() .map(Datum::Long) } - (ParquetStatistics::Int64(stats), DataType::Timestamp(ts)) if ts.precision() <= 3 => { + (ParquetStatistics::Int32(stats), DataType::Decimal(d)) => { exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) .copied() - .map(|millis| Datum::Timestamp { millis, nanos: 0 }) + .map(|unscaled| Datum::Decimal { + unscaled: unscaled as i128, + precision: d.precision(), + scale: d.scale(), + }) } - (ParquetStatistics::Int64(stats), DataType::LocalZonedTimestamp(ts)) - if ts.precision() <= 3 => - { + (ParquetStatistics::Int64(stats), DataType::Decimal(d)) => { exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) .copied() - .map(|millis| Datum::LocalZonedTimestamp { millis, nanos: 0 }) + .map(|unscaled| Datum::Decimal { + unscaled: unscaled as i128, + precision: d.precision(), + scale: d.scale(), + }) + } + (ParquetStatistics::Int64(stats), DataType::Timestamp(ts)) => { + exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) + .copied() + .and_then(|value| timestamp_datum_from_parquet_i64(value, ts.precision(), false)) + } + (ParquetStatistics::Int64(stats), DataType::LocalZonedTimestamp(ts)) => { + exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) + .copied() + .and_then(|value| timestamp_datum_from_parquet_i64(value, ts.precision(), true)) } (ParquetStatistics::Float(stats), DataType::Float(_)) => { exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) @@ -1022,10 +1349,58 @@ fn parquet_stats_to_datum( exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) .map(|value| Datum::Bytes(value.data().to_vec())) } + (ParquetStatistics::ByteArray(stats), DataType::Decimal(d)) => { + exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) + .and_then(|value| signed_be_bytes_to_i128(value.data())) + .map(|unscaled| Datum::Decimal { + unscaled, + precision: d.precision(), + scale: d.scale(), + }) + } + (ParquetStatistics::FixedLenByteArray(stats), DataType::Decimal(d)) => { + exact_parquet_value(is_min, stats.min_opt(), stats.max_opt()) + .and_then(|value| signed_be_bytes_to_i128(value.data())) + .map(|unscaled| Datum::Decimal { + unscaled, + precision: d.precision(), + scale: d.scale(), + }) + } _ => None, } } +fn timestamp_datum_from_parquet_i64( + value: i64, + precision: u32, + local_zoned: bool, +) -> Option { + let (millis, nanos) = match precision { + 0..=3 => (value, 0), + 4..=6 => ( + value.div_euclid(1_000), + (value.rem_euclid(1_000) * 1_000) as i32, + ), + _ => return None, + }; + if local_zoned { + Some(Datum::LocalZonedTimestamp { millis, nanos }) + } else { + Some(Datum::Timestamp { millis, nanos }) + } +} + +fn signed_be_bytes_to_i128(bytes: &[u8]) -> Option { + if bytes.is_empty() || bytes.len() > 16 { + return None; + } + let sign_extend = if bytes[0] & 0x80 == 0 { 0 } else { 0xff }; + let mut padded = [sign_extend; 16]; + padded[16 - bytes.len()..].copy_from_slice(bytes); + Some(i128::from_be_bytes(padded)) +} + fn exact_parquet_value<'a, T>( is_min: bool, min: Option<&'a T>, @@ -1766,14 +2141,14 @@ mod tests { let schema = writer_arrow_schema(); let mut writer: Box = Box::new( - ParquetFormatWriter::new(&output, schema.clone(), "zstd", 1) + ParquetFormatWriter::new(&output, schema.clone(), "zstd", 1, None, &HashMap::new()) .await .unwrap(), ); let batch = writer_test_batch(&schema, vec![1, 2, 3], vec![10, 20, 30]); writer.write(&batch).await.unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); // Verify valid parquet by reading back let bytes = file_io.new_input(path).unwrap().read().await.unwrap(); @@ -1791,7 +2166,7 @@ mod tests { let schema = writer_arrow_schema(); let mut writer: Box = Box::new( - ParquetFormatWriter::new(&output, schema.clone(), "zstd", 1) + ParquetFormatWriter::new(&output, schema.clone(), "zstd", 1, None, &HashMap::new()) .await .unwrap(), ); @@ -1804,7 +2179,7 @@ mod tests { .write(&writer_test_batch(&schema, vec![3, 4, 5], vec![30, 40, 50])) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let bytes = file_io.new_input(path).unwrap().read().await.unwrap(); let reader = @@ -1844,12 +2219,12 @@ mod tests { let path = "memory:/test_parquet_inline_vector.parquet"; let output = file_io.new_output(path).unwrap(); let mut writer: Box = Box::new( - ParquetFormatWriter::new(&output, schema.clone(), "zstd", 1) + ParquetFormatWriter::new(&output, schema.clone(), "zstd", 1, None, &HashMap::new()) .await .unwrap(), ); writer.write(&batch).await.unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let bytes = file_io.new_input(path).unwrap().read().await.unwrap(); let reader = @@ -1936,7 +2311,7 @@ mod tests { .await .unwrap(); writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; let raw_bytes = file_io.new_input(&path).unwrap().read().await.unwrap(); let raw_batches = @@ -2069,7 +2444,7 @@ mod tests { .write(&make_batch(vec![3], &late_variants)) .await .unwrap(); - let file_size = writer.close().await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; let raw_bytes = file_io.new_input(&path).unwrap().read().await.unwrap(); let raw_batches = @@ -2153,7 +2528,7 @@ mod tests { .write(&writer_test_batch(&schema, ids, values)) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); } buf } @@ -2296,7 +2671,7 @@ mod tests { .write(&writer_test_batch(&schema, ids, values)) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); } buf } @@ -2361,7 +2736,7 @@ mod tests { .write(&writer_test_batch(&schema, ids, values)) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); } let metadata = load_metadata_with_page_index(&buf, true); let fields = vec![int_field("id"), int_field("value")]; @@ -2417,7 +2792,7 @@ mod tests { let mut writer = AsyncArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); writer.write(&batch).await.unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); } buf } diff --git a/crates/paimon/src/arrow/format/row.rs b/crates/paimon/src/arrow/format/row.rs index 0f869938..94b740f4 100644 --- a/crates/paimon/src/arrow/format/row.rs +++ b/crates/paimon/src/arrow/format/row.rs @@ -20,7 +20,7 @@ //! Layout reference: //! `org.apache.paimon.format.row.RowFormatWriter` in paimon-java. -use super::{FilePredicates, FormatFileReader, FormatFileWriter}; +use super::{FilePredicates, FormatFileReader, FormatFileWriter, FormatWriteResult}; use crate::arrow::{ arrow_to_paimon_type, build_target_arrow_schema, is_variant_arrow_fields, paimon_type_to_arrow, variant_arrow_type, @@ -172,7 +172,7 @@ impl FormatFileWriter for RowFormatWriter { self.flush_block().await } - async fn close(mut self: Box) -> crate::Result { + async fn close(mut self: Box) -> crate::Result { self.flush_block().await?; let index_offset = self.bytes_written; @@ -206,7 +206,7 @@ impl FormatFileWriter for RowFormatWriter { self.writer.write(Bytes::from(footer_bytes)).await?; self.bytes_written += FOOTER_SIZE; self.writer.close().await?; - Ok(self.bytes_written) + Ok(FormatWriteResult::new(self.bytes_written)) } } diff --git a/crates/paimon/src/arrow/format/shredding.rs b/crates/paimon/src/arrow/format/shredding.rs index d33e98a5..df283a98 100644 --- a/crates/paimon/src/arrow/format/shredding.rs +++ b/crates/paimon/src/arrow/format/shredding.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{FilePredicates, FormatFileReader, FormatFileWriter}; +use super::{FilePredicates, FormatFileReader, FormatFileWriter, FormatWriteResult}; use crate::arrow::build_target_arrow_schema; use crate::arrow::shredding::{ assemble_shredded_variant_batch, batch_to_shredded_physical, @@ -289,12 +289,12 @@ impl FormatFileWriter for ShreddingFormatWriter { } } - async fn close(mut self: Box) -> crate::Result { + async fn close(mut self: Box) -> crate::Result { self.finalize_inferred_writer().await?; match std::mem::replace(&mut self.state, ShreddingWriterState::Closed) { ShreddingWriterState::Ready { inner, .. } => inner.close().await, ShreddingWriterState::Infer { .. } => unreachable!("infer writer finalized above"), - ShreddingWriterState::Closed => Ok(0), + ShreddingWriterState::Closed => Ok(FormatWriteResult::new(0)), } } } diff --git a/crates/paimon/src/arrow/format/vortex.rs b/crates/paimon/src/arrow/format/vortex.rs index 12250a76..631f8add 100644 --- a/crates/paimon/src/arrow/format/vortex.rs +++ b/crates/paimon/src/arrow/format/vortex.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use super::{FilePredicates, FormatFileReader, FormatFileWriter}; +use super::{FilePredicates, FormatFileReader, FormatFileWriter, FormatWriteResult}; use crate::arrow::residual::{ filter_record_batch_by_predicates, same_data_field, widen_scan_fields, }; @@ -480,7 +480,7 @@ impl FormatFileWriter for VortexFormatWriter { Ok(()) } - async fn close(self: Box) -> crate::Result { + async fn close(self: Box) -> crate::Result { let this = *self; let VortexFormatWriter { dtype, @@ -502,7 +502,7 @@ impl FormatFileWriter for VortexFormatWriter { output.write(bytes::Bytes::from(buffer)).await?; bytes_written.store(size, Ordering::Relaxed); - Ok(size) + Ok(FormatWriteResult::new(size)) } } @@ -628,10 +628,10 @@ mod tests { .enable_all() .build() .unwrap(); - let bytes = verifier_runtime + let result = verifier_runtime .block_on(async { Box::new(writer).close().await }) .unwrap(); - assert!(bytes > 0); + assert!(result.file_size > 0); } #[test] @@ -701,8 +701,8 @@ mod tests { let batch = test_batch(&schema, vec![1, 2, 3], vec![10, 20, 30]); writer.write(&batch).await.unwrap(); - let bytes = writer.close().await.unwrap(); - assert!(bytes > 0); + let result = writer.close().await.unwrap(); + assert!(result.file_size > 0); // Read back using VortexFormatReader. let input = file_io.new_input(path).unwrap(); @@ -767,7 +767,7 @@ mod tests { ) .unwrap(); writer.write(&batch).await.unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_reader = input.reader().await.unwrap(); @@ -836,7 +836,7 @@ mod tests { .write(&test_batch(&schema, vec![3, 4, 5], vec![30, 40, 50])) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_reader = input.reader().await.unwrap(); @@ -897,7 +897,7 @@ mod tests { )) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_reader = input.reader().await.unwrap(); @@ -965,7 +965,7 @@ mod tests { )) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_reader = input.reader().await.unwrap(); @@ -1022,7 +1022,7 @@ mod tests { ) .unwrap(); writer.write(&batch).await.unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_bytes = input.read().await.unwrap(); @@ -1098,7 +1098,7 @@ mod tests { )) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_reader = input.reader().await.unwrap(); @@ -1229,7 +1229,7 @@ mod tests { )) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let fields = test_file_fields(); let builder = PredicateBuilder::new(&fields); @@ -1349,7 +1349,7 @@ mod tests { )) .await .unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); let input = file_io.new_input(path).unwrap(); let file_reader = input.reader().await.unwrap(); diff --git a/crates/paimon/src/spec/binary_row.rs b/crates/paimon/src/spec/binary_row.rs index d96fa951..4a4017f3 100644 --- a/crates/paimon/src/spec/binary_row.rs +++ b/crates/paimon/src/spec/binary_row.rs @@ -859,48 +859,24 @@ pub fn extract_datum_from_arrow( } DataType::Variant(_) => extract_variant_datum_from_arrow(col, row_idx, col_idx)?, DataType::Timestamp(ts) => { - if ts.precision() <= 3 { - let arr = col - .as_any() - .downcast_ref::() - .ok_or_else(|| type_mismatch_err("Timestamp(ms)", col_idx))?; - Datum::Timestamp { - millis: arr.value(row_idx), - nanos: 0, - } - } else { - let arr = col - .as_any() - .downcast_ref::() - .ok_or_else(|| type_mismatch_err("Timestamp(us)", col_idx))?; - let micros = arr.value(row_idx); - Datum::Timestamp { - millis: micros.div_euclid(1_000), - nanos: (micros.rem_euclid(1_000) * 1_000) as i32, - } - } + let (millis, nanos) = extract_timestamp_parts_from_arrow( + col, + row_idx, + col_idx, + ts.precision(), + "Timestamp", + )?; + Datum::Timestamp { millis, nanos } } DataType::LocalZonedTimestamp(ts) => { - if ts.precision() <= 3 { - let arr = col - .as_any() - .downcast_ref::() - .ok_or_else(|| type_mismatch_err("LocalZonedTimestamp(ms)", col_idx))?; - Datum::LocalZonedTimestamp { - millis: arr.value(row_idx), - nanos: 0, - } - } else { - let arr = col - .as_any() - .downcast_ref::() - .ok_or_else(|| type_mismatch_err("LocalZonedTimestamp(us)", col_idx))?; - let micros = arr.value(row_idx); - Datum::LocalZonedTimestamp { - millis: micros.div_euclid(1_000), - nanos: (micros.rem_euclid(1_000) * 1_000) as i32, - } - } + let (millis, nanos) = extract_timestamp_parts_from_arrow( + col, + row_idx, + col_idx, + ts.precision(), + "LocalZonedTimestamp", + )?; + Datum::LocalZonedTimestamp { millis, nanos } } _ => { return Err(crate::Error::Unsupported { @@ -915,6 +891,55 @@ pub fn extract_datum_from_arrow( Ok(Some(datum)) } +fn extract_timestamp_parts_from_arrow( + col: &std::sync::Arc, + row_idx: usize, + col_idx: usize, + precision: u32, + expected: &str, +) -> crate::Result<(i64, i32)> { + match precision { + 0..=3 => { + let arr = col + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err(&format!("{expected}(ms)"), col_idx))?; + Ok((arr.value(row_idx), 0)) + } + 4..=6 => { + let arr = col + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err(&format!("{expected}(us)"), col_idx))?; + Ok(timestamp_parts_from_micros(arr.value(row_idx))) + } + 7..=9 => { + let arr = col + .as_any() + .downcast_ref::() + .ok_or_else(|| type_mismatch_err(&format!("{expected}(ns)"), col_idx))?; + Ok(timestamp_parts_from_nanos(arr.value(row_idx))) + } + _ => Err(crate::Error::Unsupported { + message: format!("Unsupported {expected} precision {precision}"), + }), + } +} + +fn timestamp_parts_from_micros(micros: i64) -> (i64, i32) { + ( + micros.div_euclid(1_000), + (micros.rem_euclid(1_000) * 1_000) as i32, + ) +} + +fn timestamp_parts_from_nanos(nanos: i64) -> (i64, i32) { + ( + nanos.div_euclid(1_000_000), + nanos.rem_euclid(1_000_000) as i32, + ) +} + fn encode_variant_bytes(value: &[u8], metadata: &[u8]) -> crate::Result> { VariantType::validate_payload(value, metadata)?; let mut bytes = Vec::with_capacity(4 + value.len() + metadata.len()); @@ -1045,6 +1070,7 @@ enum TypedColumn<'a> { Variant(&'a arrow_array::StructArray), TimestampMs(&'a arrow_array::TimestampMillisecondArray), TimestampUs(&'a arrow_array::TimestampMicrosecondArray), + TimestampNs(&'a arrow_array::TimestampNanosecondArray), } /// Downcast Arrow columns once, returning typed references paired with their DataType. @@ -1059,118 +1085,134 @@ fn downcast_columns<'a>( .map(|&col_idx| { let field = &fields[col_idx]; let col = batch.column(col_idx); - let typed = - match field.data_type() { - DataType::Boolean(_) => TypedColumn::Boolean( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Boolean", col_idx))?, - ), - DataType::TinyInt(_) => TypedColumn::Int8( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("TinyInt", col_idx))?, - ), - DataType::SmallInt(_) => TypedColumn::Int16( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("SmallInt", col_idx))?, - ), - DataType::Int(_) => TypedColumn::Int32( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Int", col_idx))?, - ), - DataType::BigInt(_) => TypedColumn::Int64( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("BigInt", col_idx))?, - ), - DataType::Float(_) => TypedColumn::Float32( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Float", col_idx))?, - ), - DataType::Double(_) => TypedColumn::Float64( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Double", col_idx))?, - ), - DataType::Char(_) | DataType::VarChar(_) => { - if let Some(arr) = col.as_any().downcast_ref::() { - TypedColumn::Utf8(arr) - } else if let Some(arr) = - col.as_any().downcast_ref::() - { - TypedColumn::Utf8View(arr) - } else if let Some(arr) = - col.as_any().downcast_ref::() - { - TypedColumn::LargeUtf8(arr) - } else { - return Err(type_mismatch_err("String", col_idx)); - } + let typed = match field.data_type() { + DataType::Boolean(_) => TypedColumn::Boolean( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Boolean", col_idx))?, + ), + DataType::TinyInt(_) => TypedColumn::Int8( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("TinyInt", col_idx))?, + ), + DataType::SmallInt(_) => TypedColumn::Int16( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("SmallInt", col_idx))?, + ), + DataType::Int(_) => TypedColumn::Int32( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Int", col_idx))?, + ), + DataType::BigInt(_) => TypedColumn::Int64( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("BigInt", col_idx))?, + ), + DataType::Float(_) => TypedColumn::Float32( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Float", col_idx))?, + ), + DataType::Double(_) => TypedColumn::Float64( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Double", col_idx))?, + ), + DataType::Char(_) | DataType::VarChar(_) => { + if let Some(arr) = col.as_any().downcast_ref::() { + TypedColumn::Utf8(arr) + } else if let Some(arr) = + col.as_any().downcast_ref::() + { + TypedColumn::Utf8View(arr) + } else if let Some(arr) = + col.as_any().downcast_ref::() + { + TypedColumn::LargeUtf8(arr) + } else { + return Err(type_mismatch_err("String", col_idx)); } - DataType::Date(_) => TypedColumn::Date32( + } + DataType::Date(_) => TypedColumn::Date32( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Date", col_idx))?, + ), + DataType::Decimal(d) => TypedColumn::Decimal128( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Decimal", col_idx))?, + d.precision(), + d.scale(), + ), + DataType::Binary(_) | DataType::VarBinary(_) => TypedColumn::Binary( + col.as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Binary", col_idx))?, + ), + DataType::Variant(_) => { + let arr = col + .as_any() + .downcast_ref() + .ok_or_else(|| type_mismatch_err("Variant", col_idx))?; + validate_variant_struct_array(arr, col_idx)?; + TypedColumn::Variant(arr) + } + DataType::Timestamp(ts) => match ts.precision() { + 0..=3 => TypedColumn::TimestampMs( col.as_any() .downcast_ref() - .ok_or_else(|| type_mismatch_err("Date", col_idx))?, + .ok_or_else(|| type_mismatch_err("Timestamp(ms)", col_idx))?, ), - DataType::Decimal(d) => TypedColumn::Decimal128( + 4..=6 => TypedColumn::TimestampUs( col.as_any() .downcast_ref() - .ok_or_else(|| type_mismatch_err("Decimal", col_idx))?, - d.precision(), - d.scale(), + .ok_or_else(|| type_mismatch_err("Timestamp(us)", col_idx))?, ), - DataType::Binary(_) | DataType::VarBinary(_) => TypedColumn::Binary( + 7..=9 => TypedColumn::TimestampNs( col.as_any() .downcast_ref() - .ok_or_else(|| type_mismatch_err("Binary", col_idx))?, + .ok_or_else(|| type_mismatch_err("Timestamp(ns)", col_idx))?, ), - DataType::Variant(_) => { - let arr = col - .as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Variant", col_idx))?; - validate_variant_struct_array(arr, col_idx)?; - TypedColumn::Variant(arr) - } - DataType::Timestamp(ts) => { - if ts.precision() <= 3 { - TypedColumn::TimestampMs( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Timestamp(ms)", col_idx))?, - ) - } else { - TypedColumn::TimestampUs( - col.as_any() - .downcast_ref() - .ok_or_else(|| type_mismatch_err("Timestamp(us)", col_idx))?, - ) - } - } - DataType::LocalZonedTimestamp(ts) => { - if ts.precision() <= 3 { - TypedColumn::TimestampMs(col.as_any().downcast_ref().ok_or_else( - || type_mismatch_err("LocalZonedTimestamp(ms)", col_idx), - )?) - } else { - TypedColumn::TimestampUs(col.as_any().downcast_ref().ok_or_else( - || type_mismatch_err("LocalZonedTimestamp(us)", col_idx), - )?) - } - } - other => { + _ => { return Err(crate::Error::Unsupported { - message: format!( - "Unsupported data type {:?} for batch column downcast at column {}", - other, col_idx - ), + message: format!("Unsupported Timestamp precision {}", ts.precision()), }); } - }; + }, + DataType::LocalZonedTimestamp(ts) => { + match ts.precision() { + 0..=3 => TypedColumn::TimestampMs(col.as_any().downcast_ref().ok_or_else( + || type_mismatch_err("LocalZonedTimestamp(ms)", col_idx), + )?), + 4..=6 => TypedColumn::TimestampUs(col.as_any().downcast_ref().ok_or_else( + || type_mismatch_err("LocalZonedTimestamp(us)", col_idx), + )?), + 7..=9 => TypedColumn::TimestampNs(col.as_any().downcast_ref().ok_or_else( + || type_mismatch_err("LocalZonedTimestamp(ns)", col_idx), + )?), + _ => { + return Err(crate::Error::Unsupported { + message: format!( + "Unsupported LocalZonedTimestamp precision {}", + ts.precision() + ), + }); + } + } + } + other => { + return Err(crate::Error::Unsupported { + message: format!( + "Unsupported data type {:?} for batch column downcast at column {}", + other, col_idx + ), + }); + } + }; Ok((typed, field)) }) .collect() @@ -1339,9 +1381,15 @@ fn write_typed_value( if arr.is_null(row_idx) { builder.set_null_at(pos); } else { - let micros = arr.value(row_idx); - let millis = micros.div_euclid(1_000); - let nanos = (micros.rem_euclid(1_000) * 1_000) as i32; + let (millis, nanos) = timestamp_parts_from_micros(arr.value(row_idx)); + builder.write_timestamp_non_compact(pos, millis, nanos); + } + } + TypedColumn::TimestampNs(arr) => { + if arr.is_null(row_idx) { + builder.set_null_at(pos); + } else { + let (millis, nanos) = timestamp_parts_from_nanos(arr.value(row_idx)); builder.write_timestamp_non_compact(pos, millis, nanos); } } diff --git a/crates/paimon/src/spec/core_options.rs b/crates/paimon/src/spec/core_options.rs index b77ab95d..e2278cae 100644 --- a/crates/paimon/src/spec/core_options.rs +++ b/crates/paimon/src/spec/core_options.rs @@ -55,6 +55,14 @@ const CHANGELOG_FILE_PREFIX_OPTION: &str = "changelog-file.prefix"; const CHANGELOG_FILE_FORMAT_OPTION: &str = "changelog-file.format"; const CHANGELOG_FILE_COMPRESSION_OPTION: &str = "changelog-file.compression"; const CHANGELOG_FILE_STATS_MODE_OPTION: &str = "changelog-file.stats-mode"; +const METADATA_STATS_MODE_OPTION: &str = "metadata.stats-mode"; +const METADATA_STATS_DENSE_STORE_OPTION: &str = "metadata.stats-dense-store"; +const METADATA_STATS_KEEP_FIRST_N_COLUMNS_OPTION: &str = "metadata.stats-keep-first-n-columns"; +const DEFAULT_METADATA_STATS_MODE: &str = "truncate(16)"; +const DEFAULT_METADATA_STATS_DENSE_STORE: bool = true; +const DEFAULT_METADATA_STATS_KEEP_FIRST_N_COLUMNS: i32 = -1; +const FIELDS_PREFIX: &str = "fields"; +const STATS_MODE_SUFFIX: &str = "stats-mode"; const ROW_TRACKING_ENABLED_OPTION: &str = "row-tracking.enabled"; pub(crate) const TABLE_TYPE_OPTION: &str = "type"; pub(crate) const FORMAT_TABLE_TYPE: &str = "format-table"; @@ -153,6 +161,56 @@ pub enum GlobalIndexSearchMode { Detail, } +/// Metadata stats collection mode. +/// +/// Reference: Java `SimpleColStatsCollector.from`. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum MetadataStatsMode { + None, + Counts, + Full, + Truncate(usize), +} + +impl MetadataStatsMode { + pub(crate) fn parse(option_name: &str, value: &str) -> crate::Result { + let value = value.trim(); + let upper = value.to_ascii_uppercase(); + match upper.as_str() { + "NONE" => Ok(Self::None), + "COUNTS" => Ok(Self::Counts), + "FULL" => Ok(Self::Full), + _ => { + let Some(length) = upper + .strip_prefix("TRUNCATE(") + .and_then(|value| value.strip_suffix(')')) + else { + return Err(crate::Error::Unsupported { + message: format!("Unsupported {option_name}: '{value}'"), + }); + }; + let length = length + .parse::() + .map_err(|e| crate::Error::DataInvalid { + message: format!( + "Option '{option_name}' must use truncate(N) with a positive integer, got: {value}" + ), + source: Some(Box::new(e)), + })?; + if length == 0 { + return Err(crate::Error::DataInvalid { + message: format!( + "Option '{option_name}' must use truncate(N) with N > 0, got: {value}" + ), + source: None, + }); + } + Ok(Self::Truncate(length)) + } + } + } +} + /// Bucket function used to map bucket keys to fixed bucket ids. /// /// Reference: Java `CoreOptions.BucketFunctionType`. @@ -820,6 +878,79 @@ impl<'a> CoreOptions<'a> { .map(String::as_str) } + /// Whether metadata stats should omit columns without collected stats. + pub(crate) fn metadata_stats_dense_store(&self) -> bool { + self.options + .get(METADATA_STATS_DENSE_STORE_OPTION) + .map(|value| value.eq_ignore_ascii_case("true")) + .unwrap_or(DEFAULT_METADATA_STATS_DENSE_STORE) + } + + /// Table-wide metadata stats mode. + pub(crate) fn metadata_stats_mode(&self) -> crate::Result { + let value = self + .options + .get(METADATA_STATS_MODE_OPTION) + .map(String::as_str) + .unwrap_or(DEFAULT_METADATA_STATS_MODE); + MetadataStatsMode::parse(METADATA_STATS_MODE_OPTION, value) + } + + /// Number of leading columns whose stats should be kept. + /// + /// A negative value means the option is ignored, matching Java Paimon. + pub(crate) fn metadata_stats_keep_first_n_columns(&self) -> crate::Result { + self.options + .get(METADATA_STATS_KEEP_FIRST_N_COLUMNS_OPTION) + .map(|value| { + value.parse::().map_err(|e| crate::Error::DataInvalid { + message: format!( + "Invalid value for {METADATA_STATS_KEEP_FIRST_N_COLUMNS_OPTION}: '{value}'" + ), + source: Some(Box::new(e)), + }) + }) + .transpose() + .map(|value| value.unwrap_or(DEFAULT_METADATA_STATS_KEEP_FIRST_N_COLUMNS)) + } + + /// Per-field metadata stats mode. + pub(crate) fn field_metadata_stats_mode( + &self, + field_name: &str, + ) -> crate::Result> { + let option_name = format!("{FIELDS_PREFIX}.{field_name}.{STATS_MODE_SUFFIX}"); + self.options + .get(&option_name) + .map(|value| MetadataStatsMode::parse(&option_name, value)) + .transpose() + } + + /// Resolve metadata stats modes for fields using Java's priority: + /// field override > keep-first-n > table-wide mode. + pub(crate) fn metadata_stats_modes<'b, I>( + &self, + field_names: I, + ) -> crate::Result> + where + I: IntoIterator, + { + let table_mode = self.metadata_stats_mode()?; + let keep_first_n = self.metadata_stats_keep_first_n_columns()?; + let mut modes = Vec::new(); + for (column_count, field_name) in field_names.into_iter().enumerate() { + let mode = if let Some(field_mode) = self.field_metadata_stats_mode(field_name)? { + field_mode + } else if keep_first_n >= 0 && column_count >= keep_first_n as usize { + MetadataStatsMode::None + } else { + table_mode + }; + modes.push(mode); + } + Ok(modes) + } + /// Avro compression codec for manifest, manifest-list and index-manifest files. /// Default is `"zstd"`, matching Java Paimon `CoreOptions.MANIFEST_COMPRESSION`. pub fn manifest_compression(&self) -> &str { @@ -1325,6 +1456,49 @@ mod tests { assert_eq!(custom_core.changelog_file_stats_mode(), Some("counts")); } + #[test] + fn test_metadata_stats_modes_follow_java_priority() { + let options = HashMap::from([ + (METADATA_STATS_MODE_OPTION.to_string(), "counts".to_string()), + ( + METADATA_STATS_KEEP_FIRST_N_COLUMNS_OPTION.to_string(), + "2".to_string(), + ), + ("fields.name.stats-mode".to_string(), "full".to_string()), + ( + "fields.payload.stats-mode".to_string(), + "truncate(8)".to_string(), + ), + ]); + let core = CoreOptions::new(&options); + + assert_eq!( + core.metadata_stats_modes(["id", "name", "payload", "extra"]) + .unwrap(), + vec![ + MetadataStatsMode::Counts, + MetadataStatsMode::Full, + MetadataStatsMode::Truncate(8), + MetadataStatsMode::None, + ] + ); + } + + #[test] + fn test_metadata_stats_mode_rejects_invalid_values() { + let options = HashMap::from([( + METADATA_STATS_MODE_OPTION.to_string(), + "truncate(0)".to_string(), + )]); + let core = CoreOptions::new(&options); + + let err = core + .metadata_stats_mode() + .expect_err("zero truncate length should fail"); + assert!(matches!(err, crate::Error::DataInvalid { message, .. } + if message.contains(METADATA_STATS_MODE_OPTION))); + } + #[test] fn test_vector_file_options_defaults_and_overrides() { let default_options = diff --git a/crates/paimon/src/table/data_evolution_reader.rs b/crates/paimon/src/table/data_evolution_reader.rs index 1c5e9a83..513d4dbc 100644 --- a/crates/paimon/src/table/data_evolution_reader.rs +++ b/crates/paimon/src/table/data_evolution_reader.rs @@ -3281,7 +3281,7 @@ mod tests { .await .unwrap(); writer.write(&batch).await.unwrap(); - writer.close().await.unwrap(); + let _ = writer.close().await.unwrap(); } let table_schema = TableSchema::new( diff --git a/crates/paimon/src/table/data_file_reader.rs b/crates/paimon/src/table/data_file_reader.rs index 84109929..05df8b3f 100644 --- a/crates/paimon/src/table/data_file_reader.rs +++ b/crates/paimon/src/table/data_file_reader.rs @@ -752,7 +752,7 @@ mod row_tests { .await .unwrap(); writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap() as i64; + let file_size = writer.close().await.unwrap().file_size as i64; let schema_id = 1; let split = DataSplitBuilder::new() @@ -823,7 +823,7 @@ mod row_tests { .await .unwrap(); writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap() as i64; + let file_size = writer.close().await.unwrap().file_size as i64; let schema_id = 1; let split = DataSplitBuilder::new() @@ -904,7 +904,7 @@ mod row_tests { .await .unwrap(); writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap() as i64; + let file_size = writer.close().await.unwrap().file_size as i64; let split = DataSplitBuilder::new() .with_snapshot(1) @@ -1545,12 +1545,19 @@ mod vector_parquet_tests { let file_path = format!("{bucket_path}/{file_name}"); let output = file_io.new_output(&file_path).unwrap(); let mut writer: Box = Box::new( - ParquetFormatWriter::new(&output, arrow_schema.clone(), "zstd", 1) - .await - .unwrap(), + ParquetFormatWriter::new( + &output, + arrow_schema.clone(), + "zstd", + 1, + None, + &std::collections::HashMap::new(), + ) + .await + .unwrap(), ); writer.write(&batch).await.unwrap(); - let file_size = writer.close().await.unwrap(); + let file_size = writer.close().await.unwrap().file_size; // Build a split whose data file's schema_id matches the table schema_id, so the // read path uses `read_type` directly (no SchemaManager lookup needed). diff --git a/crates/paimon/src/table/data_file_writer.rs b/crates/paimon/src/table/data_file_writer.rs index 58824ed5..75f17c6f 100644 --- a/crates/paimon/src/table/data_file_writer.rs +++ b/crates/paimon/src/table/data_file_writer.rs @@ -22,7 +22,7 @@ //! handles file rolling when `target_file_size` is reached, and collects //! [`DataFileMeta`] for the commit path. -use crate::arrow::format::{create_format_writer, FormatFileWriter}; +use crate::arrow::format::{create_format_writer, FormatFileWriter, FormatValueStats}; use crate::io::FileIO; use crate::spec::stats::BinaryTableStats; use crate::spec::{bucket_dir_name, DataField, DataFileMeta, EMPTY_SERIALIZED_ROW}; @@ -117,8 +117,8 @@ impl DataFileWriter { self.open_new_file(batch.schema()).await?; } - self.current_row_count += batch.num_rows() as i64; self.current_writer.as_mut().unwrap().write(batch).await?; + self.current_row_count += batch.num_rows() as i64; // Roll to a new file if target size is reached — close in background if self.current_writer.as_ref().unwrap().num_bytes() as i64 >= self.target_file_size { @@ -142,7 +142,6 @@ impl DataFileWriter { 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 { @@ -183,16 +182,17 @@ impl DataFileWriter { let row_count = self.current_row_count; self.current_row_count = 0; - let file_size = writer.close().await? as i64; + let write_result = writer.close().await?; let meta = Self::build_meta( file_name, - file_size, + 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); Ok(()) @@ -213,15 +213,16 @@ impl DataFileWriter { let write_cols = self.write_cols.clone(); self.in_flight_closes.spawn(async move { - let file_size = writer.close().await? as i64; + let write_result = writer.close().await?; Ok(Self::build_meta( file_name, - file_size, + write_result.file_size as i64, row_count, schema_id, file_source, first_row_id, write_cols, + write_result.value_stats, )) }); } @@ -239,6 +240,7 @@ impl DataFileWriter { Ok(std::mem::take(&mut self.written_files)) } + #[allow(clippy::too_many_arguments)] fn build_meta( file_name: String, file_size: i64, @@ -247,23 +249,20 @@ impl DataFileWriter { file_source: Option, first_row_id: Option, write_cols: Option>, + format_value_stats: Option, ) -> DataFileMeta { + let (value_stats, value_stats_cols) = match format_value_stats { + Some(stats) => (stats.stats, stats.columns), + None => (BinaryTableStats::empty(), Some(Vec::new())), + }; DataFileMeta { file_name, file_size, row_count, min_key: EMPTY_SERIALIZED_ROW.clone(), max_key: EMPTY_SERIALIZED_ROW.clone(), - key_stats: BinaryTableStats::new( - EMPTY_SERIALIZED_ROW.clone(), - EMPTY_SERIALIZED_ROW.clone(), - vec![], - ), - value_stats: BinaryTableStats::new( - EMPTY_SERIALIZED_ROW.clone(), - EMPTY_SERIALIZED_ROW.clone(), - vec![], - ), + key_stats: BinaryTableStats::empty(), + value_stats, min_sequence_number: 0, max_sequence_number: 0, schema_id, @@ -273,7 +272,7 @@ impl DataFileWriter { delete_row_count: Some(0), embedded_index: None, file_source, - value_stats_cols: Some(vec![]), + value_stats_cols, external_path: None, first_row_id, write_cols, diff --git a/crates/paimon/src/table/dedicated_format_file_writer.rs b/crates/paimon/src/table/dedicated_format_file_writer.rs index 15f8803d..6417b9cf 100644 --- a/crates/paimon/src/table/dedicated_format_file_writer.rs +++ b/crates/paimon/src/table/dedicated_format_file_writer.rs @@ -131,6 +131,10 @@ impl AppendDedicatedFormatFileWriter { } let normal_schema = Arc::new(arrow_schema::Schema::new(normal_arrow_fields)); + let normal_field_names = normal_table_fields + .iter() + .map(|field| field.name().to_string()) + .collect(); let vector_writer = if let Some(vector_file_format) = vector_file_format { if vector_table_fields.is_empty() { None @@ -178,7 +182,7 @@ impl AppendDedicatedFormatFileWriter { format_options.clone(), Some(0), None, - None, + Some(normal_field_names), ); Self { diff --git a/crates/paimon/src/table/kv_file_writer.rs b/crates/paimon/src/table/kv_file_writer.rs index dd5ac7d4..54ffa463 100644 --- a/crates/paimon/src/table/kv_file_writer.rs +++ b/crates/paimon/src/table/kv_file_writer.rs @@ -423,7 +423,7 @@ impl KeyValueFileWriter { writer.write(&chunk_batch).await?; } - let file_size = writer.close().await? as i64; + let file_size = writer.close().await?.file_size as i64; let key_columns: Vec> = self .config diff --git a/crates/paimon/src/table/postpone_file_writer.rs b/crates/paimon/src/table/postpone_file_writer.rs index afe318ed..8ee83cba 100644 --- a/crates/paimon/src/table/postpone_file_writer.rs +++ b/crates/paimon/src/table/postpone_file_writer.rs @@ -184,7 +184,7 @@ impl PostponeFileWriter { let creation_time = self.current_file_creation_time; self.in_flight_closes.spawn(async move { - let file_size = writer.close().await? as i64; + let file_size = writer.close().await?.file_size as i64; Ok(build_meta( file_name, file_size, @@ -249,7 +249,7 @@ impl PostponeFileWriter { let file_name = self.current_file_name.take().unwrap(); let row_count = self.current_row_count; self.current_row_count = 0; - let file_size = writer.close().await? as i64; + let file_size = writer.close().await?.file_size as i64; let min_seq = self.current_file_start_seq; let max_seq = self.next_sequence_number - 1; diff --git a/crates/paimon/src/table/table_write.rs b/crates/paimon/src/table/table_write.rs index 034e6348..a8e36e73 100644 --- a/crates/paimon/src/table/table_write.rs +++ b/crates/paimon/src/table/table_write.rs @@ -786,15 +786,16 @@ mod tests { use crate::catalog::Identifier; use crate::io::{FileIO, FileIOBuilder}; use crate::spec::{ - bucket_dir_name, BigIntType, BinaryRowBuilder, BlobType, DataField, DataType, DecimalType, - FileKind, FloatType, IndexManifest, IntType, LocalZonedTimestampType, Manifest, - ManifestList, Schema, TableSchema, TimestampType, TinyIntType, VarCharType, VectorType, + bucket_dir_name, BigIntType, BinaryRow, BinaryRowBuilder, BinaryType, BlobType, DataField, + DataType, Datum, DecimalType, FileKind, FloatType, IndexManifest, IntType, + LocalZonedTimestampType, Manifest, ManifestList, PredicateBuilder, Schema, TableSchema, + TimeType, TimestampType, TinyIntType, VarBinaryType, VarCharType, VectorType, SEQUENCE_NUMBER_FIELD_ID, SEQUENCE_NUMBER_FIELD_NAME, VALUE_KIND_FIELD_ID, VALUE_KIND_FIELD_NAME, }; use crate::table::{SnapshotManager, TableCommit}; use arrow_array::RecordBatchReader as _; - use arrow_array::{Int32Array, Int64Array, Int8Array, StringArray}; + use arrow_array::{Int32Array, Int64Array, Int8Array, StringArray, Time32MillisecondArray}; use arrow_schema::{ DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, TimeUnit, }; @@ -894,6 +895,21 @@ mod tests { .unwrap() } + fn make_nullable_batch(ids: Vec>, values: Vec>) -> RecordBatch { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("value", ArrowDataType::Int32, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(Int32Array::from(values)), + ], + ) + .unwrap() + } + fn make_vector_batch(ids: Vec, vectors: Vec>) -> RecordBatch { use arrow_array::builder::{FixedSizeListBuilder, Float32Builder}; @@ -1108,6 +1124,445 @@ mod tests { assert_eq!(snapshot.total_record_count(), Some(3)); } + #[tokio::test] + async fn test_append_write_populates_value_stats() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_value_stats"; + setup_dirs(&file_io, table_path).await; + + let table = test_table(&file_io, table_path); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&make_nullable_batch( + vec![Some(3), None], + vec![Some(30), Some(10)], + )) + .await + .unwrap(); + table_write + .write_arrow_batch(&make_nullable_batch(vec![Some(1)], vec![None])) + .await + .unwrap(); + + let messages = table_write.prepare_commit().await.unwrap(); + assert_eq!(messages.len(), 1); + assert_eq!(messages[0].new_files.len(), 1); + let file = &messages[0].new_files[0]; + assert_eq!(file.row_count, 3); + assert_eq!(file.value_stats_cols, None); + assert_eq!(file.value_stats.null_counts(), &vec![Some(1), Some(1)]); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert_eq!(min_values.get_int(0).unwrap(), 1); + assert_eq!(max_values.get_int(0).unwrap(), 3); + assert_eq!(min_values.get_int(1).unwrap(), 10); + assert_eq!(max_values.get_int(1).unwrap(), 30); + } + + #[tokio::test] + async fn test_append_write_populates_microsecond_timestamp_value_stats() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_timestamp6_value_stats"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("ts", DataType::Timestamp(TimestampType::new(6).unwrap())) + .column( + "lzts", + DataType::LocalZonedTimestamp(LocalZonedTimestampType::new(6).unwrap()), + ) + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_timestamp9_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Microsecond, None), + true, + ), + ArrowField::new( + "lzts", + ArrowDataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into())), + true, + ), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(arrow_array::TimestampMicrosecondArray::from(vec![ + Some(1_704_067_200_987_654_i64), + Some(1_704_067_200_123_456_i64), + ])), + Arc::new( + arrow_array::TimestampMicrosecondArray::from(vec![ + Some(1_704_067_201_999_999_i64), + Some(1_704_067_201_000_001_i64), + ]) + .with_timezone("UTC"), + ), + ], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!(file.value_stats_cols, None); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert_eq!( + min_values.get_timestamp_raw(0, 6).unwrap(), + (1_704_067_200_123, 456_000) + ); + assert_eq!( + max_values.get_timestamp_raw(0, 6).unwrap(), + (1_704_067_200_987, 654_000) + ); + assert_eq!( + min_values.get_timestamp_raw(1, 6).unwrap(), + (1_704_067_201_000, 1_000) + ); + assert_eq!( + max_values.get_timestamp_raw(1, 6).unwrap(), + (1_704_067_201_999, 999_000) + ); + } + + #[tokio::test] + async fn test_append_write_nanosecond_timestamp_value_stats_keeps_counts_only() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_timestamp9_value_stats"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("ts", DataType::Timestamp(TimestampType::new(9).unwrap())) + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_timestamp9_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![Arc::new(arrow_array::TimestampNanosecondArray::from(vec![ + Some(1_704_067_200_987_654_321_i64), + None, + ]))], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!(file.value_stats_cols, None); + assert_eq!(file.value_stats.null_counts(), &vec![Some(1)]); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert!(min_values.is_null_at(0)); + assert!(max_values.is_null_at(0)); + } + + #[tokio::test] + async fn test_append_write_truncates_string_value_stats_and_keeps_binary_counts() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_skip_variable_length_stats"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .column( + "payload", + DataType::VarBinary(VarBinaryType::new(1024).unwrap()), + ) + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_variable_length_stats_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, false), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("payload", ArrowDataType::Binary, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(vec![2, 1])), + Arc::new(StringArray::from(vec![ + Some("a long string value"), + Some("another long string value"), + ])), + Arc::new(arrow_array::BinaryArray::from(vec![ + Some(b"large-binary-value" as &[u8]), + Some(b"another-large-binary-value" as &[u8]), + ])), + ], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!(file.value_stats_cols, None); + assert_eq!( + file.value_stats.null_counts(), + &vec![Some(0), Some(0), Some(0)] + ); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert_eq!(min_values.get_int(0).unwrap(), 1); + assert_eq!(max_values.get_int(0).unwrap(), 2); + assert_eq!(min_values.get_string(1).unwrap(), "a long string va"); + assert_eq!(max_values.get_string(1).unwrap(), "another long sts"); + assert!(min_values.is_null_at(2)); + assert!(max_values.is_null_at(2)); + } + + #[tokio::test] + async fn test_append_write_applies_configured_metadata_stats_modes() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_configured_metadata_stats_modes"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("name", DataType::VarChar(VarCharType::string_type())) + .column("payload", DataType::Binary(BinaryType::new(1024).unwrap())) + .column("value", DataType::Int(IntType::new())) + .option("metadata.stats-mode", "counts") + .option("metadata.stats-keep-first-n-columns", "2") + .option("fields.name.stats-mode", "full") + .option("fields.payload.stats-mode", "full") + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_configured_stats_modes_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("name", ArrowDataType::Utf8, true), + ArrowField::new("payload", ArrowDataType::Binary, true), + ArrowField::new("value", ArrowDataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(vec![Some(2), None])), + Arc::new(StringArray::from(vec![ + Some("alpha-long-value-12345"), + Some("zeta-long-value-99999"), + ])), + Arc::new(arrow_array::BinaryArray::from(vec![ + Some(b"first-binary-value" as &[u8]), + None, + ])), + Arc::new(Int32Array::from(vec![Some(10), Some(20)])), + ], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!( + file.value_stats_cols, + Some(vec![ + "id".to_string(), + "name".to_string(), + "payload".to_string() + ]) + ); + assert_eq!( + file.value_stats.null_counts(), + &vec![Some(1), Some(0), Some(1)] + ); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert!(min_values.is_null_at(0)); + assert!(max_values.is_null_at(0)); + assert_eq!(min_values.get_string(1).unwrap(), "alpha-long-value-12345"); + assert_eq!(max_values.get_string(1).unwrap(), "zeta-long-value-99999"); + assert!(min_values.is_null_at(2)); + assert!(max_values.is_null_at(2)); + } + + #[tokio::test] + async fn test_append_write_respects_non_dense_metadata_stats_storage() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_non_dense_metadata_stats"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("id", DataType::Int(IntType::new())) + .column("ignored", DataType::Int(IntType::new())) + .option("metadata.stats-mode", "none") + .option("metadata.stats-dense-store", "false") + .option("fields.id.stats-mode", "counts") + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_non_dense_stats_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", ArrowDataType::Int32, true), + ArrowField::new("ignored", ArrowDataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![ + Arc::new(Int32Array::from(vec![Some(1), None])), + Arc::new(Int32Array::from(vec![Some(10), Some(20)])), + ], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!(file.value_stats_cols, None); + assert_eq!(file.value_stats.null_counts(), &vec![Some(1), None]); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert_eq!(min_values.arity(), 2); + assert_eq!(max_values.arity(), 2); + assert!(min_values.is_null_at(0)); + assert!(max_values.is_null_at(0)); + assert!(min_values.is_null_at(1)); + assert!(max_values.is_null_at(1)); + } + + #[tokio::test] + async fn test_append_write_populates_time_value_stats() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_time_value_stats"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("tm", DataType::Time(TimeType::new(3).unwrap())) + .build() + .unwrap(); + let table = Table::new( + file_io.clone(), + Identifier::new("default", "test_time_value_stats_table"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + + let arrow_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "tm", + ArrowDataType::Time32(TimeUnit::Millisecond), + true, + )])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![Arc::new(Time32MillisecondArray::from(vec![ + Some(3_000), + None, + Some(1_000), + ]))], + ) + .unwrap(); + + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let file = &messages[0].new_files[0]; + assert_eq!(file.value_stats_cols, None); + assert_eq!(file.value_stats.null_counts(), &vec![Some(1)]); + + let min_values = BinaryRow::from_serialized_bytes(file.value_stats.min_values()).unwrap(); + let max_values = BinaryRow::from_serialized_bytes(file.value_stats.max_values()).unwrap(); + assert_eq!(min_values.get_int(0).unwrap(), 1_000); + assert_eq!(max_values.get_int(0).unwrap(), 3_000); + } + + #[tokio::test] + async fn test_scan_prunes_real_append_files_using_written_value_stats() { + let file_io = test_file_io(); + let table_path = "memory:/test_table_write_value_stats_scan"; + setup_dirs(&file_io, table_path).await; + + let table = test_table(&file_io, table_path); + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&make_batch(vec![1, 2], vec![10, 20])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write + .write_arrow_batch(&make_batch(vec![100, 101], vec![1000, 1010])) + .await + .unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .greater_than("id", Datum::Int(10)) + .unwrap(); + let mut reader = table.new_read_builder(); + reader.with_filter(predicate); + let (_plan, trace) = reader.new_scan().plan_with_trace().await.unwrap(); + + assert_eq!(trace.final_files, 1, "scan trace: {trace:?}"); + assert!( + trace.manifest_entries_pruned_by_data_stats >= 1, + "scan trace: {trace:?}" + ); + } + #[tokio::test] async fn test_table_write_row_format_roundtrip() { let file_io = test_file_io(); @@ -1233,6 +1688,7 @@ mod tests { assert_eq!(parquet_files[0].row_count, 3); assert_eq!(blob_files[0].row_count, 3); + assert_eq!(parquet_files[0].write_cols, Some(vec!["id".to_string()])); assert_eq!(blob_files[0].write_cols, Some(vec!["payload".to_string()])); // Commit and verify snapshot @@ -1244,6 +1700,71 @@ mod tests { assert_eq!(snapshot.id(), 1); } + #[tokio::test] + async fn test_dedicated_blob_stats_use_normal_write_columns() { + let file_io = test_file_io(); + let table_path = "memory:/test_dedicated_blob_stats_mapping"; + setup_dirs(&file_io, table_path).await; + + let schema = Schema::builder() + .column("payload", DataType::Blob(BlobType::new())) + .column("a", DataType::Int(IntType::new())) + .column("b", DataType::Int(IntType::new())) + .option("data-evolution.enabled", "true") + .build() + .unwrap(); + let table = Table::new( + file_io, + Identifier::new("default", "test_dedicated_blob_stats_mapping"), + table_path.to_string(), + TableSchema::new(0, &schema), + None, + ); + let batch = RecordBatch::try_new( + Arc::new(ArrowSchema::new(vec![ + ArrowField::new("payload", ArrowDataType::Binary, true), + ArrowField::new("a", ArrowDataType::Int32, false), + ArrowField::new("b", ArrowDataType::Int32, false), + ])), + vec![ + Arc::new(arrow_array::BinaryArray::from(vec![Some( + b"payload" as &[u8], + )])), + Arc::new(Int32Array::from(vec![100])), + Arc::new(Int32Array::from(vec![0])), + ], + ) + .unwrap(); + + let mut table_write = TableWrite::new(&table, "test-user".to_string()).unwrap(); + table_write.write_arrow_batch(&batch).await.unwrap(); + let messages = table_write.prepare_commit().await.unwrap(); + let normal_write_cols = messages[0] + .new_files + .iter() + .find(|file| file.file_name.ends_with(".parquet")) + .unwrap() + .write_cols + .clone(); + TableCommit::new(table.clone(), "test-user".to_string()) + .commit(messages) + .await + .unwrap(); + + let predicate = PredicateBuilder::new(table.schema().fields()) + .greater_than("a", Datum::Int(50)) + .unwrap(); + let mut reader = table.new_read_builder(); + reader.with_filter(predicate); + let (_plan, trace) = reader.new_scan().plan_with_trace().await.unwrap(); + + assert!(trace.final_files > 0, "scan trace: {trace:?}"); + assert_eq!( + normal_write_cols, + Some(vec!["a".to_string(), "b".to_string()]) + ); + } + async fn assert_vector_write_uses_dedicated_file(table_path: &str, vector_file_format: &str) { let file_io = test_file_io(); setup_dirs(&file_io, table_path).await; @@ -1284,6 +1805,7 @@ mod tests { assert_eq!(vector_files.len(), 1); assert_eq!(normal_files[0].row_count, 3); assert_eq!(vector_files[0].row_count, 3); + assert_eq!(normal_files[0].write_cols, Some(vec!["id".to_string()])); assert_eq!( vector_files[0].write_cols, Some(vec!["embedding".to_string()])