Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 99 additions & 26 deletions crates/integrations/datafusion/src/merge_into.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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;
Expand Down Expand Up @@ -569,13 +571,6 @@ async fn execute_cow_merge_inner(

// Handle NOT MATCHED → INSERT
if !clauses.inserts.is_empty() {
let table_fields: Vec<String> = 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} \
Expand All @@ -595,7 +590,7 @@ async fn execute_cow_merge_inner(
&clauses.inserts,
s_alias,
&[],
&table_fields,
table.schema().fields(),
temp_tracker,
)
.await?;
Expand Down Expand Up @@ -734,21 +729,14 @@ async fn execute_merge_into_once(
injected_columns.push(format!("__upd_{col}"));
}
}
// Table schema field names for reordering INSERT columns
let table_fields: Vec<String> = 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,
&not_matched_batches,
&parsed.inserts,
s_alias,
&injected_columns,
&table_fields,
table.schema().fields(),
&mut temp_tracker,
)
.await?;
Expand Down Expand Up @@ -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<Vec<RecordBatch>> {
if not_matched_batches.is_empty() || not_matched_batches.iter().all(|b| b.num_rows() == 0) {
Expand Down Expand Up @@ -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<Vec<RecordBatch>> {
let mut all_batches = Vec::new();
let mut consumed_predicates: Vec<String> = Vec::new();
Expand All @@ -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<RecordBatch> {
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<ArrayRef> {
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],
Expand Down Expand Up @@ -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 {
Expand All @@ -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::<Vec<_>>()
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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::<Int32Array>()
.unwrap();
let second = normalized
.column(1)
.as_any()
.downcast_ref::<Int32Array>()
.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(
Expand Down
6 changes: 3 additions & 3 deletions crates/paimon/src/arrow/format/blob.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -725,7 +725,7 @@ impl FormatFileWriter for BlobFormatWriter {
Ok(())
}

async fn close(mut self: Box<Self>) -> crate::Result<u64> {
async fn close(mut self: Box<Self>) -> crate::Result<FormatWriteResult> {
let index_bytes = encode_delta_varints_write(&self.lengths);
let index_length = index_bytes.len() as i32;

Expand All @@ -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))
}
}

Expand Down
40 changes: 37 additions & 3 deletions crates/paimon/src/arrow/format/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")]
Expand All @@ -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;
Expand Down Expand Up @@ -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<Self>) -> crate::Result<u64>;
async fn close(self: Box<Self>) -> crate::Result<FormatWriteResult>;
}

pub(crate) struct FormatWriteResult {
pub(crate) file_size: u64,
pub(crate) value_stats: Option<FormatValueStats>,
}

pub(crate) struct FormatValueStats {
pub(crate) stats: BinaryTableStats,
pub(crate) columns: Option<Vec<String>>,
}

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<Vec<String>>,
) -> Self {
Self {
file_size,
value_stats: Some(FormatValueStats {
stats: value_stats,
columns,
}),
}
}
}

/// Create a format reader based on the file extension.
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading