This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-rust.git
The following commit(s) were added to refs/heads/main by this push:
new 9545797 fix(datafusion): guard cow merge partition pruning (#475)
9545797 is described below
commit 95457972e389c302d544c51afa774c130d34905b
Author: 承昭 <[email protected]>
AuthorDate: Thu Jul 9 10:37:41 2026 +0800
fix(datafusion): guard cow merge partition pruning (#475)
---
crates/integrations/datafusion/src/merge_into.rs | 131 ++++++++++++++++++++-
.../datafusion/tests/append_merge_into.rs | 28 +++++
2 files changed, 155 insertions(+), 4 deletions(-)
diff --git a/crates/integrations/datafusion/src/merge_into.rs
b/crates/integrations/datafusion/src/merge_into.rs
index 203538d..124dd65 100644
--- a/crates/integrations/datafusion/src/merge_into.rs
+++ b/crates/integrations/datafusion/src/merge_into.rs
@@ -32,7 +32,8 @@ use datafusion::arrow::datatypes::{DataType as ArrowDataType,
Field, Schema};
use datafusion::error::{DataFusionError, Result as DFResult};
use datafusion::prelude::{DataFrame, SessionContext};
use datafusion::sql::sqlparser::ast::{
- AssignmentTarget, Merge, MergeAction, MergeClauseKind, MergeInsertKind,
TableFactor,
+ AssignmentTarget, BinaryOperator, Expr as SqlExpr, Ident, Merge,
MergeAction, MergeClauseKind,
+ MergeInsertKind, TableFactor,
};
use futures::TryStreamExt;
@@ -325,12 +326,14 @@ async fn execute_cow_merge_once(
let (source_ref, source_alias) = extract_source_ref(&merge.source)?;
let (target_ref, target_alias) = extract_table_ref(&merge.table)?;
- let on_condition = merge.on.to_string();
+ let on_expr = &merge.on;
+ let on_condition = on_expr.to_string();
let s_alias = source_alias.as_deref().unwrap_or(&source_ref);
let t_alias = target_alias.as_deref().unwrap_or("__cow_t");
// Build partition filter from source data to avoid scanning all partitions
- let partition_set = build_source_partition_set(ctx, table, &source_ref,
s_alias).await?;
+ let partition_set =
+ build_source_partition_set(ctx, table, &source_ref, s_alias, t_alias,
on_expr).await?;
let mut writer = CopyOnWriteMergeWriter::new(table,
update_columns.clone(), partition_set)
.await
@@ -1365,17 +1368,24 @@ pub(crate) async fn build_partition_set_from_where(
/// Query source table for distinct partition values and build a partition set.
///
/// Returns `None` for non-partitioned tables or when the source lacks matching
-/// partition key columns (falls back to full-partition scan).
+/// partition key columns (falls back to full-partition scan). The source
values
+/// are only safe to apply to the target when the MERGE condition equates every
+/// target partition column with the same source column.
async fn build_source_partition_set(
ctx: &SQLContext,
table: &Table,
source_ref: &str,
s_alias: &str,
+ t_alias: &str,
+ on_expr: &SqlExpr,
) -> DFResult<Option<HashSet<Vec<u8>>>> {
let partition_keys = table.schema().partition_keys();
if partition_keys.is_empty() {
return Ok(None);
}
+ if !can_prune_target_by_source_partitions(partition_keys, t_alias,
s_alias, on_expr) {
+ return Ok(None);
+ }
let cols = partition_keys
.iter()
@@ -1392,6 +1402,71 @@ async fn build_source_partition_set(
}
}
+fn can_prune_target_by_source_partitions(
+ partition_keys: &[String],
+ t_alias: &str,
+ s_alias: &str,
+ on_expr: &SqlExpr,
+) -> bool {
+ partition_keys
+ .iter()
+ .all(|key| on_expr_contains_alias_column_eq(on_expr, t_alias, key,
s_alias, key))
+}
+
+fn on_expr_contains_alias_column_eq(
+ expr: &SqlExpr,
+ left_alias: &str,
+ left_column: &str,
+ right_alias: &str,
+ right_column: &str,
+) -> bool {
+ match expr {
+ SqlExpr::BinaryOp {
+ left,
+ op: BinaryOperator::And,
+ right,
+ } => {
+ on_expr_contains_alias_column_eq(
+ left,
+ left_alias,
+ left_column,
+ right_alias,
+ right_column,
+ ) || on_expr_contains_alias_column_eq(
+ right,
+ left_alias,
+ left_column,
+ right_alias,
+ right_column,
+ )
+ }
+ SqlExpr::BinaryOp {
+ left,
+ op: BinaryOperator::Eq,
+ right,
+ } => {
+ is_alias_column(left, left_alias, left_column)
+ && is_alias_column(right, right_alias, right_column)
+ || is_alias_column(left, right_alias, right_column)
+ && is_alias_column(right, left_alias, left_column)
+ }
+ _ => false,
+ }
+}
+
+fn is_alias_column(expr: &SqlExpr, alias: &str, column: &str) -> bool {
+ match expr {
+ SqlExpr::CompoundIdentifier(parts) => is_qualified_column(parts,
alias, column),
+ _ => false,
+ }
+}
+
+fn is_qualified_column(parts: &[Ident], alias: &str, column: &str) -> bool {
+ parts.len() == 2
+ && parts[0].value.eq_ignore_ascii_case(alias)
+ && parts[1].value.eq_ignore_ascii_case(column)
+}
+
/// Rewrite SQL expressions by replacing original table references with
aliases.
///
/// For example, `paimon.test_db.target.a = source.a` becomes `t.a = s.a`
@@ -1536,6 +1611,54 @@ mod tests {
}
}
+ #[test]
+ fn test_source_partition_pruning_requires_partition_equality() {
+ let merge = parse_merge(
+ "MERGE INTO target t USING source s ON t.a = s.a \
+ WHEN MATCHED THEN UPDATE SET b = s.b",
+ );
+ let partition_keys = vec!["pt".to_string()];
+
+ assert!(!can_prune_target_by_source_partitions(
+ &partition_keys,
+ "t",
+ "s",
+ &merge.on,
+ ));
+ }
+
+ #[test]
+ fn test_source_partition_pruning_accepts_partition_equality() {
+ let merge = parse_merge(
+ "MERGE INTO target t USING source s ON t.a = s.a AND t.pt = s.pt \
+ WHEN MATCHED THEN UPDATE SET b = s.b",
+ );
+ let partition_keys = vec!["pt".to_string()];
+
+ assert!(can_prune_target_by_source_partitions(
+ &partition_keys,
+ "t",
+ "s",
+ &merge.on,
+ ));
+ }
+
+ #[test]
+ fn test_source_partition_pruning_accepts_reversed_partition_equality() {
+ let merge = parse_merge(
+ "MERGE INTO target t USING source s ON s.pt = t.pt AND t.a = s.a \
+ WHEN MATCHED THEN UPDATE SET b = s.b",
+ );
+ let partition_keys = vec!["pt".to_string()];
+
+ assert!(can_prune_target_by_source_partitions(
+ &partition_keys,
+ "t",
+ "s",
+ &merge.on,
+ ));
+ }
+
#[tokio::test]
async fn test_merge_into_updates_matched_rows() {
let (_tmp, sql_context, table) =
setup_data_evolution_table("t_merge").await;
diff --git a/crates/integrations/datafusion/tests/append_merge_into.rs
b/crates/integrations/datafusion/tests/append_merge_into.rs
index c14936c..c9d66bb 100644
--- a/crates/integrations/datafusion/tests/append_merge_into.rs
+++ b/crates/integrations/datafusion/tests/append_merge_into.rs
@@ -1132,6 +1132,34 @@ async fn test_partitioned_update_and_insert() {
);
}
+#[tokio::test]
+async fn
test_partitioned_merge_without_partition_equality_does_not_prune_target() {
+ let (_tmp, sql_context) =
+ setup("CREATE TABLE paimon.test_db.target (a INT, b INT, pt INT)
PARTITIONED BY (pt)")
+ .await;
+ exec(
+ &sql_context,
+ "INSERT INTO paimon.test_db.target VALUES (1, 10, 2)",
+ )
+ .await;
+ exec(
+ &sql_context,
+ "CREATE TEMPORARY TABLE paimon.test_db.source (a INT, b INT, pt INT)
AS SELECT * FROM (VALUES (1, 100, 1)) AS t(a, b, pt)",
+ )
+ .await;
+
+ exec(
+ &sql_context,
+ "MERGE INTO paimon.test_db.target t \
+ USING paimon.test_db.source s ON t.a = s.a \
+ WHEN MATCHED THEN UPDATE SET b = s.b \
+ WHEN NOT MATCHED THEN INSERT (a, b, pt) VALUES (s.a, s.b, s.pt)",
+ )
+ .await;
+
+ assert_eq!(query_a_b_pt(&sql_context).await, vec![(1, 100, 2),]);
+}
+
#[tokio::test]
async fn test_partitioned_successive_merges() {
let (_tmp, sql_context) = setup_partitioned().await;
