Skip to content

Commit 8df32eb

Browse files
slfan1989slfan1989
authored
fix(datafusion): preserve string literals when rewriting MERGE references (#1266)
1 parent 4a7a896 commit 8df32eb

1 file changed

Lines changed: 143 additions & 74 deletions

File tree

‎crates/integrations/datafusion/src/merge_into.rs‎

Lines changed: 143 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
//! via [`paimon::table::CopyOnWriteMergeWriter`].
2424
2525
use std::collections::{HashMap, HashSet};
26+
use std::ops::ControlFlow;
2627
use std::sync::atomic::{AtomicU64, Ordering};
2728
use std::sync::Arc;
2829

@@ -34,8 +35,8 @@ use datafusion::arrow::datatypes::{DataType as ArrowDataType, Field, Schema};
3435
use datafusion::error::{DataFusionError, Result as DFResult};
3536
use datafusion::prelude::{DataFrame, SessionContext};
3637
use datafusion::sql::sqlparser::ast::{
37-
AssignmentTarget, BinaryOperator, Expr as SqlExpr, Ident, Merge, MergeAction, MergeClauseKind,
38-
MergeInsertExpr, MergeInsertKind, TableFactor,
38+
visit_expressions_mut, AssignmentTarget, BinaryOperator, Expr as SqlExpr, Ident, Merge,
39+
MergeAction, MergeClauseKind, MergeInsertExpr, MergeInsertKind, TableFactor, VisitMut,
3940
};
4041
use futures::TryStreamExt;
4142

@@ -314,7 +315,28 @@ async fn execute_cow_merge_once(
314315
table: &Table,
315316
enable_ident_normalization: bool,
316317
) -> DFResult<DataFrame> {
317-
let mut clauses = extract_cow_merge_clauses(merge, enable_ident_normalization)?;
318+
let (source_ref, source_alias) = extract_source_ref(&merge.source)?;
319+
let (target_ref, target_alias) = extract_table_ref(&merge.table)?;
320+
let s_alias = source_alias.as_deref().unwrap_or(&source_ref);
321+
let t_alias = target_alias.as_deref().unwrap_or("__cow_t");
322+
323+
// Rewrite column qualifiers before rendering SQL, preserving literal contents.
324+
let mut aliased_merge = merge.clone();
325+
rewrite_table_references(
326+
&mut aliased_merge.on,
327+
&target_ref,
328+
t_alias,
329+
&source_ref,
330+
s_alias,
331+
);
332+
rewrite_table_references(
333+
&mut aliased_merge.clauses,
334+
&target_ref,
335+
t_alias,
336+
&source_ref,
337+
s_alias,
338+
);
339+
let clauses = extract_cow_merge_clauses(&aliased_merge, enable_ident_normalization)?;
318340
validate_merge_insert_columns(&clauses.inserts, table.schema().fields())?;
319341
for matched in &clauses.matched {
320342
if let CowMatchedAction::Update(update) = &matched.action {
@@ -334,12 +356,8 @@ async fn execute_cow_merge_once(
334356
}
335357
}
336358

337-
let (source_ref, source_alias) = extract_source_ref(&merge.source)?;
338-
let (target_ref, target_alias) = extract_table_ref(&merge.table)?;
339359
let on_expr = &merge.on;
340-
let on_condition = on_expr.to_string();
341-
let s_alias = source_alias.as_deref().unwrap_or(&source_ref);
342-
let t_alias = target_alias.as_deref().unwrap_or("__cow_t");
360+
let on_condition = aliased_merge.on.to_string();
343361

344362
// Build partition filter from source data to avoid scanning all partitions
345363
let partition_set =
@@ -349,27 +367,6 @@ async fn execute_cow_merge_once(
349367
.await
350368
.map_err(to_datafusion_error)?;
351369

352-
// Rewrite ON condition and all clause expressions: replace original table references with aliases
353-
let on_condition = rewrite_condition(&on_condition, &target_ref, t_alias, &source_ref, s_alias);
354-
for mc in &mut clauses.matched {
355-
if let Some(ref mut pred) = mc.predicate {
356-
*pred = rewrite_condition(pred, &target_ref, t_alias, &source_ref, s_alias);
357-
}
358-
if let CowMatchedAction::Update(ref mut upd) = mc.action {
359-
for expr in &mut upd.exprs {
360-
*expr = rewrite_condition(expr, &target_ref, t_alias, &source_ref, s_alias);
361-
}
362-
}
363-
}
364-
for ins in &mut clauses.inserts {
365-
for expr in &mut ins.value_exprs {
366-
*expr = rewrite_condition(expr, &target_ref, t_alias, &source_ref, s_alias);
367-
}
368-
if let Some(ref mut pred) = ins.predicate {
369-
*pred = rewrite_condition(pred, &target_ref, t_alias, &source_ref, s_alias);
370-
}
371-
}
372-
373370
// Read each target file individually, attach __paimon_file_idx and __paimon_row_offset
374371
let mut temp_tracker = TempTableTracker::new(ctx);
375372
let (has_target_data, cow_table_name) =
@@ -1633,60 +1630,43 @@ fn is_qualified_column(parts: &[Ident], alias: &str, column: &str) -> bool {
16331630
&& parts[1].value.eq_ignore_ascii_case(column)
16341631
}
16351632

1636-
/// Rewrite SQL expressions by replacing original table references with aliases.
1633+
/// Replace table qualifiers in column references, leaving literals and column names intact.
16371634
///
16381635
/// For example, `paimon.test_db.target.a = source.a` becomes `t.a = s.a`
16391636
/// when target_ref="paimon.test_db.target", t_alias="t", source_ref="source", s_alias="s".
16401637
///
1641-
/// Uses word-boundary-aware replacement to avoid mangling identifiers that
1642-
/// contain the table name as a substring (e.g. `"my_source.x"` won't match `"source."`).
1643-
///
1644-
/// Limitation: this is best-effort text replacement, not AST-level rewriting.
1645-
/// It may produce incorrect results for quoted identifiers containing dots
1646-
/// (e.g. `"my.table".col`) or other unusual naming patterns.
1647-
///
1648-
/// TODO: migrate to AST-level rewriting for correctness with edge-case identifiers.
1649-
fn rewrite_condition(
1650-
condition: &str,
1638+
/// Match the longest qualifier first, and rewrite it only once so an alias that
1639+
/// equals another table name cannot be rewritten again.
1640+
fn rewrite_table_references<T: VisitMut>(
1641+
node: &mut T,
16511642
target_ref: &str,
16521643
t_alias: &str,
16531644
source_ref: &str,
16541645
s_alias: &str,
1655-
) -> String {
1656-
let mut result = condition.to_string();
1657-
// Replace longer (more qualified) names first to avoid partial matches
1658-
if target_ref.len() >= source_ref.len() {
1659-
result = replace_table_ref(&result, target_ref, t_alias);
1660-
result = replace_table_ref(&result, source_ref, s_alias);
1661-
} else {
1662-
result = replace_table_ref(&result, source_ref, s_alias);
1663-
result = replace_table_ref(&result, target_ref, t_alias);
1664-
}
1665-
result
1666-
}
1667-
1668-
/// Replace `"ref."` with `"alias."` only when `ref` is not preceded by a word character.
1669-
fn replace_table_ref(input: &str, table_ref: &str, alias: &str) -> String {
1670-
let needle = format!("{table_ref}.");
1671-
let replacement = format!("{alias}.");
1672-
let mut result = String::with_capacity(input.len());
1673-
let mut remaining = input;
1674-
1675-
while let Some(pos) = remaining.find(&needle) {
1676-
let preceding_is_word = pos > 0 && {
1677-
let prev = remaining.as_bytes()[pos - 1];
1678-
prev.is_ascii_alphanumeric() || prev == b'_'
1679-
};
1680-
if preceding_is_word {
1681-
result.push_str(&remaining[..pos + needle.len()]);
1682-
} else {
1683-
result.push_str(&remaining[..pos]);
1684-
result.push_str(&replacement);
1646+
) {
1647+
let _: ControlFlow<()> = visit_expressions_mut(node, |expr| {
1648+
if let SqlExpr::CompoundIdentifier(parts) = expr {
1649+
for prefix_len in (1..parts.len()).rev() {
1650+
let qualifier = parts[..prefix_len]
1651+
.iter()
1652+
.map(ToString::to_string)
1653+
.collect::<Vec<_>>()
1654+
.join(".");
1655+
let alias = if qualifier == target_ref {
1656+
Some(t_alias)
1657+
} else if qualifier == source_ref {
1658+
Some(s_alias)
1659+
} else {
1660+
None
1661+
};
1662+
if let Some(alias) = alias {
1663+
parts.splice(..prefix_len, [Ident::new(alias)]);
1664+
break;
1665+
}
1666+
}
16851667
}
1686-
remaining = &remaining[pos + needle.len()..];
1687-
}
1688-
result.push_str(remaining);
1689-
result
1668+
ControlFlow::Continue(())
1669+
});
16901670
}
16911671

16921672
/// Return a DataFrame with a single "count" column.
@@ -1777,6 +1757,37 @@ mod tests {
17771757
}
17781758
}
17791759

1760+
#[test]
1761+
fn test_merge_reference_rewrite_preserves_literal_contents() {
1762+
for (input, expected) in [
1763+
("target.id = source.id", "t.id = s.id"),
1764+
("'source.payload'", "'source.payload'"),
1765+
("'it''s target.payload'", "'it''s target.payload'"),
1766+
(
1767+
"'{\"path\":\"source.payload\"}'",
1768+
"'{\"path\":\"source.payload\"}'",
1769+
),
1770+
("source.name LIKE 'target.%'", "s.name LIKE 'target.%'"),
1771+
("source.\"source.payload\"", "s.\"source.payload\""),
1772+
("my_source.id = source.id", "my_source.id = s.id"),
1773+
(
1774+
"target.nested.value = source.value",
1775+
"t.nested.value = s.value",
1776+
),
1777+
] {
1778+
let parse_expr = |sql| {
1779+
Parser::new(&GenericDialect {})
1780+
.try_with_sql(sql)
1781+
.unwrap()
1782+
.parse_expr()
1783+
.unwrap()
1784+
};
1785+
let mut actual = parse_expr(input);
1786+
rewrite_table_references(&mut actual, "target", "t", "source", "s");
1787+
assert_eq!(actual, parse_expr(expected), "{input}");
1788+
}
1789+
}
1790+
17801791
#[test]
17811792
fn test_merge_update_assignments_normalize_identifiers() {
17821793
for (target, enabled, expected) in [
@@ -2612,6 +2623,64 @@ mod tests {
26122623
}
26132624
}
26142625

2626+
#[tokio::test]
2627+
async fn test_cow_merge_preserves_string_literals() {
2628+
for (on_predicate, matched_predicate, insert_predicate) in [
2629+
("", "", ""),
2630+
(" AND s.name = 'paimon.test_db.source.payload'", "", ""),
2631+
("", " AND s.name = 'paimon.test_db.source.payload'", ""),
2632+
("", "", " AND s.name = 'paimon.test_db.source.payload'"),
2633+
] {
2634+
let (_tmp, sql_context, table) = setup_append_only_table("target").await;
2635+
for sql in [
2636+
"CREATE TABLE paimon.test_db.source (id INT, name VARCHAR)",
2637+
"INSERT INTO paimon.test_db.source VALUES (1, 'paimon.test_db.source.payload'), (4, 'paimon.test_db.source.payload')",
2638+
] {
2639+
sql_context.sql(sql).await.unwrap().collect().await.unwrap();
2640+
}
2641+
2642+
let merge = parse_merge(&format!(
2643+
"MERGE INTO paimon.test_db.target t USING paimon.test_db.source s \
2644+
ON paimon.test_db.target.id = paimon.test_db.source.id{on_predicate} \
2645+
WHEN MATCHED{matched_predicate} THEN UPDATE SET name = 'paimon.test_db.target.payload' \
2646+
WHEN NOT MATCHED{insert_predicate} THEN INSERT (id, name, value) VALUES (s.id, 'paimon.test_db.source.payload', 40)"
2647+
));
2648+
let result = execute_merge_into(&sql_context, &merge, table, true)
2649+
.await
2650+
.unwrap()
2651+
.collect()
2652+
.await
2653+
.unwrap();
2654+
assert_eq!(
2655+
result[0]
2656+
.column(0)
2657+
.as_any()
2658+
.downcast_ref::<UInt64Array>()
2659+
.unwrap()
2660+
.value(0),
2661+
2,
2662+
);
2663+
2664+
let batches = sql_context
2665+
.sql("SELECT id, name, value FROM paimon.test_db.target ORDER BY id")
2666+
.await
2667+
.unwrap()
2668+
.collect()
2669+
.await
2670+
.unwrap();
2671+
assert_eq!(
2672+
collect_rows(&batches),
2673+
vec![
2674+
(1, "paimon.test_db.target.payload".to_string(), 10),
2675+
(2, "bob".to_string(), 20),
2676+
(3, "charlie".to_string(), 30),
2677+
(4, "paimon.test_db.source.payload".to_string(), 40),
2678+
],
2679+
"ON: {on_predicate}, MATCHED: {matched_predicate}, INSERT: {insert_predicate}",
2680+
);
2681+
}
2682+
}
2683+
26152684
#[tokio::test]
26162685
async fn test_cow_merge_no_matches() {
26172686
let (_tmp, sql_context, table) = setup_append_only_table("t_cow_nomatch").await;

0 commit comments

Comments
 (0)