You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@paimon.apache.org by "JingsongLi (via GitHub)" <gi...@apache.org> on 2026/04/13 08:59:51 UTC
[PR] [WIP] feat(write): add MERGE INTO support with RowIdUpdateWriter [paimon-rust]
JingsongLi opened a new pull request, #241:
URL: https://github.com/apache/paimon-rust/pull/241
<!--
Thank you very much for contributing to Paimon Rust - we are happy that you want to help us improve it. To help the community review your contribution in the best possible way, please go through the checklist below, which will get the contribution into a shape in which it can be best reviewed.
## Contribution Checklist
- Make sure that the pull request corresponds to a [GitHub issue](https://github.com/apache/paimon-rust/issues). Exceptions are made for typos in documentation or comments, which need no issue.
- Fill out the template below to describe the changes contributed by the pull request. That will give reviewers the context they need to do the review.
- Make sure that the change passes the automated tests, i.e., `cargo test` passes.
- Each pull request should address only one issue, not mix up code from multiple issues.
**(The sections below can be removed for hotfixes or typos)**
-->
### Purpose
<!-- Linking this pull request to the issue -->
- Add RowIdUpdateWriter: engine-agnostic Table-layer API for partial-column updates via _ROW_ID, returning Vec<CommitMessage> for caller-controlled commit
- Add shared DataFileWriter extracted from TableWrite, parameterized via FileMetaTemplate for both append and data-evolution write modes
- Add DataEvolutionWriter for partial-column parquet files with write_cols and first_row_id metadata
- Add MERGE INTO execution in DataFusion with single LEFT JOIN split by _ROW_ID null/not-null: matched rows go through RowIdUpdateWriter, not-matched rows go through TableWrite, committed atomically
- Support conditional WHEN NOT MATCHED AND <predicate> THEN INSERT
- Add 18 integration tests aligned with Java Paimon RowTrackingTestBase
<!-- What is the purpose of the change -->
### Brief change log
<!-- Please describe the changes made in this pull request and explain how they address the issue -->
### Tests
<!-- List unit tests or integration cases to verify this change -->
### API and Format
<!-- Does this change affect API or storage format -->
### Documentation
<!-- Does this change introduce a new feature or require documentation updates -->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(write): add MERGE & UPDATE support with RowIdUpdateWriter [paimon-rust]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi commented on PR #241:
URL: https://github.com/apache/paimon-rust/pull/241#issuecomment-4248453028
@littlecoder04 Thanks for the review, fixed comments and add e2e tests.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(write): add MERGE & UPDATE with DataEvolutionWriter [paimon-rust]
Posted by "littlecoder04 (via GitHub)" <gi...@apache.org>.
littlecoder04 commented on PR #241:
URL: https://github.com/apache/paimon-rust/pull/241#issuecomment-4248744688
+1
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(write): add MERGE & UPDATE with DataEvolutionWriter [paimon-rust]
Posted by "JingsongLi (via GitHub)" <gi...@apache.org>.
JingsongLi merged PR #241:
URL: https://github.com/apache/paimon-rust/pull/241
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(write): add MERGE & UPDATE support with RowIdUpdateWriter [paimon-rust]
Posted by "littlecoder04 (via GitHub)" <gi...@apache.org>.
littlecoder04 commented on code in PR #241:
URL: https://github.com/apache/paimon-rust/pull/241#discussion_r3080644921
##########
crates/paimon/src/table/row_id_update_write.rs:
##########
@@ -0,0 +1,693 @@
+// 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.
+
+//! Row-ID-based update writer for data evolution tables.
+//!
+//! [`RowIdUpdateWriter`] accepts rows to update (identified by `_ROW_ID`) along with
+//! new column values, then handles file metadata lookup, row grouping,
+//! reading original columns, applying updates, and writing partial-column files.
+//!
+//! The writer does NOT commit — it returns `Vec<CommitMessage>` that the caller
+//! passes to [`TableCommit`](super::TableCommit).
+//! This separation allows callers to compose multiple operations into a single commit,
+//! similar to Iceberg's Transaction/Action pattern.
+
+use crate::io::FileIO;
+use crate::spec::{BinaryRow, CoreOptions, DataFileMeta, PartitionComputer};
+use crate::table::commit_message::CommitMessage;
+use crate::table::data_file_writer::DataFileWriter;
+use crate::table::DataSplitBuilder;
+use crate::table::Table;
+use crate::Result;
+use arrow_array::{Array, ArrayRef, Int64Array, RecordBatch};
+use arrow_select::concat::concat_batches;
+use arrow_select::interleave::interleave;
+use futures::TryStreamExt;
+use std::collections::HashMap;
+
+/// Engine-agnostic writer for partial-column updates via `_ROW_ID`.
+///
+/// Usage:
+/// 1. Create via [`RowIdUpdateWriter::new`] (validates preconditions).
+/// 2. Feed matched rows via [`add_matched_batch`](Self::add_matched_batch).
+/// Each batch must contain a `_ROW_ID` (Int64) column plus the update columns.
+/// 3. Call [`prepare_commit`](Self::prepare_commit) to produce `CommitMessage`s.
+/// 4. Commit via [`TableCommit`](super::TableCommit) (caller's responsibility).
+///
+/// The query engine (DataFusion, custom, etc.) is responsible for:
+/// - Parsing the MERGE INTO SQL
+/// - Executing the JOIN to find matched rows
+/// - Computing new column values
+/// - Passing the results as `RecordBatch`es with `_ROW_ID` + update columns
+#[must_use = "writer must be used to call prepare_commit()"]
+pub struct RowIdUpdateWriter {
+ table: Table,
+ update_columns: Vec<String>,
+ matched_batches: Vec<RecordBatch>,
+}
+
+impl RowIdUpdateWriter {
+ /// Create a new writer for the given table and update columns.
+ ///
+ /// Validates:
+ /// - `data-evolution.enabled = true`
+ /// - `row-tracking.enabled = true`
+ /// - No primary keys
+ /// - Update columns don't include partition keys
+ pub fn new(table: &Table, update_columns: Vec<String>) -> Result<Self> {
+ let schema = table.schema();
+ let core_options = CoreOptions::new(schema.options());
+
+ if !core_options.data_evolution_enabled() {
+ return Err(crate::Error::Unsupported {
+ message:
+ "MERGE INTO is only supported for tables with 'data-evolution.enabled' = 'true'"
+ .to_string(),
+ });
+ }
+ if !core_options.row_tracking_enabled() {
+ return Err(crate::Error::Unsupported {
+ message: "MERGE INTO requires 'row-tracking.enabled' = 'true'".to_string(),
+ });
+ }
+ if !schema.primary_keys().is_empty() {
+ return Err(crate::Error::Unsupported {
+ message: "MERGE INTO on data evolution tables does not support primary keys"
+ .to_string(),
+ });
+ }
+
+ let partition_keys = schema.partition_keys();
+ for col in &update_columns {
+ if partition_keys.contains(col) {
+ return Err(crate::Error::Unsupported {
+ message: format!("Cannot update partition column '{col}' in MERGE INTO"),
+ });
+ }
+ }
+
+ Ok(Self {
+ table: table.clone(),
+ update_columns,
+ matched_batches: Vec::new(),
+ })
+ }
+
+ /// Add a batch of matched rows.
+ ///
+ /// The batch must contain:
+ /// - A `_ROW_ID` column (Int64) identifying which rows to update
+ /// - One column for each entry in `update_columns` with the new values
+ pub fn add_matched_batch(&mut self, batch: RecordBatch) -> Result<()> {
+ if batch.num_rows() == 0 {
+ return Ok(());
+ }
+
+ // Validate _ROW_ID column exists
+ if batch.column_by_name("_ROW_ID").is_none() {
+ return Err(crate::Error::DataInvalid {
+ message: "Matched batch must contain a '_ROW_ID' column".to_string(),
+ source: None,
+ });
+ }
+
+ self.matched_batches.push(batch);
+ Ok(())
+ }
+
+ /// Scan file metadata, group matched rows by file, read originals,
+ /// apply updates, and write partial-column files.
+ ///
+ /// Returns `CommitMessage`s for the caller to commit via [`TableCommit`](super::TableCommit).
+ #[must_use = "commit messages must be passed to TableCommit"]
+ pub async fn prepare_commit(self) -> Result<Vec<CommitMessage>> {
+ let total_matched: usize = self.matched_batches.iter().map(|b| b.num_rows()).sum();
+ if total_matched == 0 {
+ return Ok(Vec::new());
+ }
+
+ // 1. Scan file metadata and build row_id -> file index
+ let scan = self.table.new_read_builder().new_scan();
+ let plan = scan.plan().await?;
+
+ let mut file_index: Vec<FileRowRange> = Vec::new();
+ for split in plan.splits() {
+ let partition_bytes = split.partition().to_serialized_bytes();
+ let bucket = split.bucket();
+ let bucket_path = split.bucket_path().to_string();
+ let snapshot_id = split.snapshot_id();
+ let total_buckets = split.total_buckets();
+ for file in split.data_files() {
+ if let Some(first_row_id) = file.first_row_id {
+ file_index.push(FileRowRange {
+ first_row_id,
+ last_row_id: first_row_id + file.row_count - 1,
+ partition: partition_bytes.clone(),
+ bucket,
+ bucket_path: bucket_path.clone(),
+ snapshot_id,
+ total_buckets,
+ file: file.clone(),
+ });
+ }
+ }
+ }
+ file_index.sort_by_key(|f| f.first_row_id);
+
+ if file_index.is_empty() {
+ return Err(crate::Error::DataInvalid {
+ message: "No files with row tracking found in target table".to_string(),
+ source: None,
+ });
+ }
+
+ // 2. Group matched rows by their owning file
+ let mut file_matches: HashMap<usize, Vec<MatchedRow>> = HashMap::new();
+
+ for (batch_idx, batch) in self.matched_batches.iter().enumerate() {
+ let row_id_col = batch
+ .column_by_name("_ROW_ID")
+ .unwrap()
+ .as_any()
+ .downcast_ref::<Int64Array>()
+ .ok_or_else(|| crate::Error::DataInvalid {
+ message: "_ROW_ID column must be Int64".to_string(),
+ source: None,
+ })?;
+
+ for row_idx in 0..batch.num_rows() {
+ let row_id = row_id_col.value(row_idx);
+ let (file_pos, file_range) =
+ find_owning_file(&file_index, row_id).ok_or_else(|| {
+ crate::Error::DataInvalid {
+ message: format!("No file found for _ROW_ID {row_id}"),
+ source: None,
+ }
+ })?;
+
+ let offset = (row_id - file_range.first_row_id) as usize;
+ file_matches.entry(file_pos).or_default().push(MatchedRow {
+ offset,
+ batch_idx,
+ row_idx,
+ });
+ }
+ }
+
+ // 3. For each affected file: read original columns, apply updates, write partial files
+ let mut writer = PartialColumnsWriter::new(&self.table, self.update_columns.clone())?;
+
+ for (&file_pos, matched_rows) in &file_matches {
+ let file_range = &file_index[file_pos];
+ let file = &file_range.file;
+ let first_row_id = file.first_row_id.unwrap();
+ let row_count = file.row_count as usize;
+
+ // Read original columns from this file
+ let col_refs: Vec<&str> = self.update_columns.iter().map(|s| s.as_str()).collect();
+ let mut rb = self.table.new_read_builder();
+ rb.with_projection(&col_refs);
+ let read = rb.new_read()?;
+
+ let split = DataSplitBuilder::new()
+ .with_snapshot(file_range.snapshot_id)
+ .with_partition(BinaryRow::from_serialized_bytes(&file_range.partition)?)
+ .with_bucket(file_range.bucket)
+ .with_bucket_path(file_range.bucket_path.clone())
+ .with_total_buckets(file_range.total_buckets)
+ .with_data_files(vec![file.clone()])
Review Comment:
At this point, `_ROW_ID` has already been narrowed to a single `DataFileMeta`, and this read only loads that file back. However, the current logical row range may span multiple files with the same `first_row_id` (base file + partial-column files). Upstream Java/Python paths seem to read the whole `first_row_id` group.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(write): add MERGE & UPDATE support with RowIdUpdateWriter [paimon-rust]
Posted by "littlecoder04 (via GitHub)" <gi...@apache.org>.
littlecoder04 commented on code in PR #241:
URL: https://github.com/apache/paimon-rust/pull/241#discussion_r3080782149
##########
crates/integrations/datafusion/src/merge_into.rs:
##########
@@ -0,0 +1,766 @@
+// 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.
+
+//! MERGE INTO execution for data evolution tables.
+//!
+//! This module provides the DataFusion-specific SQL parsing and JOIN execution layer.
+//! The engine-agnostic merge logic (file metadata lookup, row grouping, reading originals,
+//! applying updates, writing partial files, committing) lives in
+//! [`paimon::table::RowIdUpdateWriter`].
+
+use std::sync::Arc;
+
+use datafusion::arrow::array::{RecordBatch, UInt64Array};
+use datafusion::arrow::compute;
+use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
+use datafusion::datasource::MemTable;
+use datafusion::error::{DataFusionError, Result as DFResult};
+use datafusion::prelude::{DataFrame, SessionContext};
+use datafusion::sql::sqlparser::ast::{
+ AssignmentTarget, Merge, MergeAction, MergeClauseKind, MergeInsertKind, TableFactor,
+};
+
+use paimon::table::{RowIdUpdateWriter, Table};
+
+use crate::error::to_datafusion_error;
+
+/// Maximum number of retries when MERGE INTO conflicts with concurrent compaction.
+const MERGE_INTO_MAX_RETRIES: u32 = 5;
+
+/// Execute a MERGE INTO statement on a data evolution Paimon table.
+///
+/// The `ctx` must already have the target and source tables registered.
+/// The `merge` AST node comes from sqlparser.
+///
+/// When a concurrent compaction rewrites files that this MERGE INTO references,
+/// the commit will fail with a row ID conflict. In that case, the entire operation
+/// (scan + write + commit) is retried from scratch against the new file layout.
+pub(crate) async fn execute_merge_into(
+ ctx: &SessionContext,
+ merge: &Merge,
+ table: Table,
+) -> DFResult<DataFrame> {
+ for retry in 0..MERGE_INTO_MAX_RETRIES {
+ match execute_merge_into_once(ctx, merge, &table).await {
+ Ok(df) => return Ok(df),
+ Err(e) if is_row_id_conflict(&e) => {
+ if retry + 1 >= MERGE_INTO_MAX_RETRIES {
+ return Err(DataFusionError::External(Box::new(std::io::Error::other(
+ format!(
+ "MERGE INTO failed after {} retries due to concurrent compaction: {}",
+ MERGE_INTO_MAX_RETRIES, e
+ ),
+ ))));
+ }
+ // Retry: re-execute the entire MERGE INTO against the new file layout
+ continue;
+ }
+ Err(e) => return Err(e),
+ }
+ }
+ unreachable!()
+}
+
+/// Check if a DataFusion error is caused by a row ID conflict during commit.
+pub(crate) fn is_row_id_conflict(err: &DataFusionError) -> bool {
+ match err {
+ DataFusionError::External(e) => e.to_string().contains("Row ID conflict"),
+ _ => false,
+ }
+}
+
+/// Single attempt of MERGE INTO execution.
+async fn execute_merge_into_once(
+ ctx: &SessionContext,
+ merge: &Merge,
+ table: &Table,
+) -> DFResult<DataFrame> {
+ // 1. Parse all MERGE clauses
+ let parsed = extract_merge_clauses(merge)?;
+
+ // Validate preconditions early and create writer (before executing any SQL)
+ let update_writer = if let Some(ref upd) = parsed.update {
+ Some(RowIdUpdateWriter::new(table, upd.columns.clone()).map_err(to_datafusion_error)?)
+ } else {
+ None
+ };
+
+ let (target_ref, target_alias) = extract_table_ref(&merge.table)?;
+ let (source_ref, source_alias) = extract_source_ref(&merge.source)?;
+ let on_condition = merge.on.to_string();
+ let t_alias = target_alias.as_deref().unwrap_or(&target_ref);
+ let s_alias = source_alias.as_deref().unwrap_or(&source_ref);
+
+ // 2. Build a single LEFT JOIN: source LEFT JOIN target
+ // _ROW_ID IS NOT NULL → matched (UPDATE path)
+ // _ROW_ID IS NULL → not matched (INSERT path)
+ let mut select_parts = vec![format!("{t_alias}.\"_ROW_ID\"")];
+
+ // Add update expressions (prefixed to avoid collisions)
+ if let Some(ref upd) = parsed.update {
+ for (col, expr) in upd.columns.iter().zip(upd.exprs.iter()) {
+ select_parts.push(format!("{expr} AS \"__upd_{col}\""));
+ }
+ }
+
+ // Add source columns for INSERT path (all source columns via s.*)
+ // We also need insert expressions if they differ from source columns
+ if !parsed.inserts.is_empty() {
+ select_parts.push(format!("{s_alias}.*"));
+ }
+
+ let select_clause = select_parts.join(", ");
+ // Safety: all interpolated values (select_clause, source_ref, s_alias, t_alias, on_condition)
+ // originate from sqlparser AST's Display impl, so they are well-formed SQL fragments.
+ let join_sql = format!(
+ "SELECT {select_clause} FROM {source_ref} AS {s_alias} \
+ LEFT JOIN {target_ref} AS {t_alias} ON {on_condition}"
+ );
+
+ let join_result = ctx.sql(&join_sql).await?.collect().await?;
+
+ // 3. Split by _ROW_ID null/not-null
+ let mut all_messages = Vec::new();
+ let mut total_count: u64 = 0;
+
+ // Separate matched and not-matched rows
+ let (matched_batches, not_matched_batches) = split_by_row_id(&join_result)?;
+
+ // 4. Handle matched rows (UPDATE)
+ if let Some(mut writer) = update_writer {
+ let upd = parsed.update.as_ref().unwrap();
+ let matched_count: usize = matched_batches.iter().map(|b| b.num_rows()).sum();
+ if matched_count > 0 {
+ // Extract _ROW_ID + update columns (rename __upd_X → X)
+ let update_batches = project_update_columns(&matched_batches, &upd.columns)?;
+ for batch in update_batches {
+ writer
+ .add_matched_batch(batch)
+ .map_err(to_datafusion_error)?;
+ }
+ let update_messages = writer.prepare_commit().await.map_err(to_datafusion_error)?;
+ all_messages.extend(update_messages);
+ total_count += matched_count as u64;
+ }
+ }
+
+ // 5. Handle not-matched rows (INSERT)
+ if !parsed.inserts.is_empty() {
+ // Collect the exact set of injected column names to strip from JOIN result
+ let mut injected_columns: Vec<String> = vec!["_ROW_ID".to_string()];
+ if let Some(ref upd) = parsed.update {
+ for col in &upd.columns {
+ injected_columns.push(format!("__upd_{col}"));
+ }
+ }
+ let insert_batches = build_insert_batches(
+ ctx,
+ ¬_matched_batches,
+ &parsed.inserts,
+ s_alias,
+ &injected_columns,
+ )
+ .await?;
+ let insert_count: usize = insert_batches.iter().map(|b| b.num_rows()).sum();
+ if insert_count > 0 {
+ let mut table_write = table
+ .new_write_builder()
+ .new_write()
+ .map_err(to_datafusion_error)?;
+ for batch in &insert_batches {
+ table_write
+ .write_arrow_batch(batch)
+ .await
+ .map_err(to_datafusion_error)?;
+ }
+ let insert_messages = table_write
+ .prepare_commit()
+ .await
+ .map_err(to_datafusion_error)?;
+ all_messages.extend(insert_messages);
+ total_count += insert_count as u64;
+ }
+ }
+
+ // 6. Commit all messages atomically
+ if !all_messages.is_empty() {
+ let commit = table.new_write_builder().new_commit();
+ commit
+ .commit(all_messages)
+ .await
+ .map_err(to_datafusion_error)?;
+ }
+
+ ok_result(ctx, total_count)
+}
+
+/// Split join result into matched (_ROW_ID not null) and not-matched (_ROW_ID null) batches.
+fn split_by_row_id(batches: &[RecordBatch]) -> DFResult<(Vec<RecordBatch>, Vec<RecordBatch>)> {
+ let mut matched = Vec::new();
+ let mut not_matched = Vec::new();
+
+ for batch in batches {
+ if batch.num_rows() == 0 {
+ continue;
+ }
+ let row_id_col = batch.column_by_name("_ROW_ID").ok_or_else(|| {
+ DataFusionError::Internal("_ROW_ID column not found in join result".to_string())
+ })?;
+
+ let is_not_null = compute::is_not_null(row_id_col)?;
+ let is_null = compute::is_null(row_id_col)?;
+
+ let matched_batch = compute::filter_record_batch(batch, &is_not_null)?;
+ if matched_batch.num_rows() > 0 {
+ matched.push(matched_batch);
+ }
+
+ let not_matched_batch = compute::filter_record_batch(batch, &is_null)?;
+ if not_matched_batch.num_rows() > 0 {
+ not_matched.push(not_matched_batch);
+ }
+ }
+
+ Ok((matched, not_matched))
+}
+
+/// Extract _ROW_ID + __upd_X columns from matched batches, renaming __upd_X → X.
+pub(crate) fn project_update_columns(
+ batches: &[RecordBatch],
+ update_columns: &[String],
+) -> DFResult<Vec<RecordBatch>> {
+ let mut result = Vec::new();
+ for batch in batches {
+ let row_id_idx = batch
+ .schema()
+ .index_of("_ROW_ID")
+ .map_err(|e| DataFusionError::Internal(format!("_ROW_ID not found: {e}")))?;
+
+ let mut columns = vec![batch.column(row_id_idx).clone()];
+ let mut fields = vec![batch.schema().field(row_id_idx).clone()];
+
+ for col in update_columns {
+ let prefixed = format!("__upd_{col}");
+ let idx = batch.schema().index_of(&prefixed).map_err(|e| {
+ DataFusionError::Internal(format!("Column {prefixed} not found: {e}"))
+ })?;
+ columns.push(batch.column(idx).clone());
+ fields.push(Field::new(
+ col,
+ batch.schema().field(idx).data_type().clone(),
+ true,
+ ));
+ }
+
+ let schema = Arc::new(Schema::new(fields));
+ let projected = RecordBatch::try_new(schema, columns)?;
+ result.push(projected);
+ }
+ Ok(result)
+}
+
+/// Build insert batches from not-matched rows, applying INSERT clause projections and predicates.
+async fn build_insert_batches(
+ ctx: &SessionContext,
+ not_matched_batches: &[RecordBatch],
+ inserts: &[MergeInsertClause],
+ s_alias: &str,
+ injected_columns: &[String],
+) -> DFResult<Vec<RecordBatch>> {
+ if not_matched_batches.is_empty() || not_matched_batches.iter().all(|b| b.num_rows() == 0) {
+ return Ok(Vec::new());
+ }
+
+ // Strip injected columns (_ROW_ID, __upd_*) — keep only source columns
+ let source_batches = strip_non_source_columns(not_matched_batches, injected_columns)?;
+
+ // Register as temp table for SQL-based projection/filtering
+ let first_schema = source_batches[0].schema();
+ let mem_table = MemTable::try_new(first_schema, vec![source_batches])?;
+ let tmp_name = format!("__merge_not_matched_{}", std::process::id());
+ ctx.register_table(&tmp_name, Arc::new(mem_table))?;
+
+ let result = build_insert_batches_inner(ctx, inserts, s_alias, &tmp_name).await;
+
+ // Always clean up temp table, even on error
+ let _ = ctx.deregister_table(&tmp_name);
+
+ result
+}
+
+/// Execute INSERT clause queries against the registered temp table.
+async fn build_insert_batches_inner(
+ ctx: &SessionContext,
+ inserts: &[MergeInsertClause],
+ s_alias: &str,
+ tmp_name: &str,
+) -> DFResult<Vec<RecordBatch>> {
+ let mut all_batches = Vec::new();
+ let mut consumed_predicates: Vec<String> = Vec::new();
+
+ for ins in inserts {
+ let mut conditions = Vec::new();
+ for prev in &consumed_predicates {
+ conditions.push(format!("NOT ({prev})"));
+ }
+ if let Some(ref pred) = ins.predicate {
+ conditions.push(pred.clone());
+ consumed_predicates.push(pred.clone());
+ }
+
+ let where_clause = if conditions.is_empty() {
+ String::new()
+ } else {
+ format!(" WHERE {}", conditions.join(" AND "))
+ };
+
+ let select_clause = insert_select_clause(ins);
+ let sql = format!("SELECT {select_clause} FROM {tmp_name} AS {s_alias}{where_clause}");
+
+ let batches = ctx.sql(&sql).await?.collect().await?;
+ all_batches.extend(batches);
+ }
+
+ Ok(all_batches)
+}
+
+/// Remove injected columns from batches, keeping only source columns.
+fn strip_non_source_columns(
+ batches: &[RecordBatch],
+ injected_columns: &[String],
+) -> DFResult<Vec<RecordBatch>> {
+ let mut result = Vec::new();
+ for batch in batches {
+ let schema = batch.schema();
+ let mut indices = Vec::new();
+ let mut fields = Vec::new();
+ for (i, field) in schema.fields().iter().enumerate() {
+ if injected_columns.contains(field.name()) {
+ continue;
+ }
+ indices.push(i);
+ fields.push(field.as_ref().clone());
+ }
+ let new_schema = Arc::new(Schema::new(fields));
+ let columns: Vec<_> = indices.iter().map(|&i| batch.column(i).clone()).collect();
+ let projected = RecordBatch::try_new(new_schema, columns)?;
+ result.push(projected);
+ }
+ Ok(result)
+}
+
+/// Build the SELECT clause for an INSERT clause: `*` or `expr AS col, ...`.
+fn insert_select_clause(ins: &MergeInsertClause) -> String {
+ if ins.columns.is_empty() && ins.value_exprs.is_empty() {
+ "*".to_string()
+ } else {
+ ins.columns
+ .iter()
+ .zip(ins.value_exprs.iter())
+ .map(|(col, expr)| format!("{expr} AS {col}"))
Review Comment:
This keeps the generated batch in `INSERT (...)` order. But the write path later reads partition/bucket fields by target schema index, not by column name, so a reordered insert list can mis-map columns on partitioned / fixed-bucket tables.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@paimon.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org