crates/integrations/datafusion/tests/update_tests.rs

// 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.

//! UPDATE integration tests for append-only tables (CoW path).
//!
//! Test cases adapted from Java Paimon's `UpdateTableTestBase.scala`.

mod common;

use paimon_datafusion::SQLContext;

use common::{
    assert_sql_error, create_sql_context, create_test_env, dml_count, exec, query_int_str_int,
};

// ======================= Helpers =======================

async fn setup() -> (tempfile::TempDir, SQLContext) {
    let (tmp, catalog) = create_test_env();
    let sql_context = create_sql_context(catalog).await;
    sql_context
        .sql("CREATE SCHEMA paimon.test_db")
        .await
        .unwrap();
    sql_context
        .sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
        .await
        .unwrap();
    exec(
        &sql_context,
        "INSERT INTO paimon.test_db.t VALUES (1, 'a', 10), (2, 'b', 20), (3, 'c', 30), (4, 'd', 40)",
    )
    .await;
    (tmp, sql_context)
}

async fn setup_partitioned() -> (tempfile::TempDir, SQLContext) {
    let (tmp, catalog) = create_test_env();
    let sql_context = create_sql_context(catalog).await;
    sql_context
        .sql("CREATE SCHEMA paimon.test_db")
        .await
        .unwrap();
    sql_context
        .sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, pt INT) PARTITIONED BY (pt)")
        .await
        .unwrap();
    exec(
        &sql_context,
        "INSERT INTO paimon.test_db.t VALUES (1, 'a', 1), (2, 'b', 1), (3, 'c', 2), (4, 'd', 2)",
    )
    .await;
    (tmp, sql_context)
}

async fn query(sql_context: &SQLContext) -> Vec<(i32, String, i32)> {
    query_int_str_int(
        sql_context,
        "SELECT id, name, age FROM paimon.test_db.t ORDER BY id",
    )
    .await
}

async fn query_pt(sql_context: &SQLContext) -> Vec<(i32, String, i32)> {
    query_int_str_int(
        sql_context,
        "SELECT id, name, pt FROM paimon.test_db.t ORDER BY id",
    )
    .await
}

// ======================= Basic UPDATE =======================

#[tokio::test]
async fn test_update_with_where() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'a_new' WHERE id = 1",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a_new".into(), 10),
            (2, "b".into(), 20),
            (3, "c".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

#[tokio::test]
async fn test_update_with_expression() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET age = age + 1 WHERE id = 1",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 11),
            (2, "b".into(), 20),
            (3, "c".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

#[tokio::test]
async fn test_update_no_match_empty_commit() {
    let (_tmp, sql_context) = setup().await;

    let cnt = dml_count(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE id = 99",
    )
    .await;
    assert_eq!(cnt, 0);

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 10),
            (2, "b".into(), 20),
            (3, "c".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

#[tokio::test]
async fn test_update_without_where() {
    let (_tmp, sql_context) = setup().await;

    exec(&sql_context, "UPDATE paimon.test_db.t SET age = 0").await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 0),
            (2, "b".into(), 0),
            (3, "c".into(), 0),
            (4, "d".into(), 0),
        ]
    );
}

// ======================= Partitioned table UPDATE =======================

#[tokio::test]
async fn test_update_partitioned_by_non_partition_column() {
    let (_tmp, sql_context) = setup_partitioned().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE id = 1",
    )
    .await;

    assert_eq!(
        query_pt(&sql_context).await,
        vec![
            (1, "x".into(), 1),
            (2, "b".into(), 1),
            (3, "c".into(), 2),
            (4, "d".into(), 2),
        ]
    );
}

#[tokio::test]
async fn test_update_partitioned_by_partition_column() {
    let (_tmp, sql_context) = setup_partitioned().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE pt = 1",
    )
    .await;

    assert_eq!(
        query_pt(&sql_context).await,
        vec![
            (1, "x".into(), 1),
            (2, "x".into(), 1),
            (3, "c".into(), 2),
            (4, "d".into(), 2),
        ]
    );
}

#[tokio::test]
async fn test_update_partitioned_by_both_columns() {
    let (_tmp, sql_context) = setup_partitioned().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE pt = 2 AND id = 3",
    )
    .await;

    assert_eq!(
        query_pt(&sql_context).await,
        vec![
            (1, "a".into(), 1),
            (2, "b".into(), 1),
            (3, "x".into(), 2),
            (4, "d".into(), 2),
        ]
    );
}

// ======================= IN / NOT IN =======================

#[tokio::test]
async fn test_update_in_condition() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE id IN (1, 3)",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "x".into(), 10),
            (2, "b".into(), 20),
            (3, "x".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

#[tokio::test]
async fn test_update_not_in_condition() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE id NOT IN (1, 3)",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 10),
            (2, "x".into(), 20),
            (3, "c".into(), 30),
            (4, "x".into(), 40),
        ]
    );
}

// ======================= Cross-column references =======================

#[tokio::test]
async fn test_update_cross_column_reference() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET age = id * 100 WHERE id <= 2",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 100),
            (2, "b".into(), 200),
            (3, "c".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

#[tokio::test]
async fn test_update_multiple_columns_with_expressions() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'updated', age = age + id WHERE id >= 3",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 10),
            (2, "b".into(), 20),
            (3, "updated".into(), 33),
            (4, "updated".into(), 44),
        ]
    );
}

// ======================= Successive updates =======================

#[tokio::test]
async fn test_successive_updates() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET age = 0 WHERE id = 1",
    )
    .await;
    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET age = 99 WHERE id = 1",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 99),
            (2, "b".into(), 20),
            (3, "c".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

// ======================= OR condition =======================

#[tokio::test]
async fn test_update_or_condition() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'x' WHERE id = 1 OR id = 4",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "x".into(), 10),
            (2, "b".into(), 20),
            (3, "c".into(), 30),
            (4, "x".into(), 40),
        ]
    );
}

// ======================= Range condition =======================

#[tokio::test]
async fn test_update_range_condition() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET age = 0 WHERE id >= 2 AND id <= 3",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "a".into(), 10),
            (2, "b".into(), 0),
            (3, "c".into(), 0),
            (4, "d".into(), 40),
        ]
    );
}

// ======================= Subquery condition =======================

#[tokio::test]
async fn test_update_in_subquery() {
    let (_tmp, sql_context) = setup().await;

    exec(
        &sql_context,
        "CREATE TEMPORARY TABLE paimon.test_db.src AS SELECT * FROM (VALUES (1), (3)) AS t(id)",
    )
    .await;

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET name = 'sub' WHERE id IN (SELECT id FROM paimon.test_db.src)",
    )
    .await;

    assert_eq!(
        query(&sql_context).await,
        vec![
            (1, "sub".into(), 10),
            (2, "b".into(), 20),
            (3, "sub".into(), 30),
            (4, "d".into(), 40),
        ]
    );
}

// ======================= NULL handling =======================

#[tokio::test]
async fn test_update_with_null_values() {
    let (tmp, catalog) = create_test_env();
    let sql_context = create_sql_context(catalog).await;
    sql_context
        .sql("CREATE SCHEMA paimon.test_db")
        .await
        .unwrap();
    sql_context
        .sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
        .await
        .unwrap();
    sql_context
        .sql("INSERT INTO paimon.test_db.t VALUES (1, 'a', 10), (2, NULL, 20), (3, 'c', NULL)")
        .await
        .unwrap()
        .collect()
        .await
        .unwrap();

    exec(
        &sql_context,
        "UPDATE paimon.test_db.t SET age = 0 WHERE name = 'a'",
    )
    .await;

    let rows = query(&sql_context).await;
    assert_eq!(rows[0], (1, "a".into(), 0));
    assert_eq!(rows[1].0, 2);
    assert_eq!(rows[1].2, 20);
    assert_eq!(rows[2].0, 3);
    drop(tmp);
}

// ======================= Empty table =======================

#[tokio::test]
async fn test_update_empty_table() {
    let (tmp, catalog) = create_test_env();
    let sql_context = create_sql_context(catalog).await;
    sql_context
        .sql("CREATE SCHEMA paimon.test_db")
        .await
        .unwrap();
    sql_context
        .sql("CREATE TABLE paimon.test_db.t (id INT, name VARCHAR, age INT)")
        .await
        .unwrap();

    let cnt = dml_count(&sql_context, "UPDATE paimon.test_db.t SET name = 'x'").await;
    assert_eq!(cnt, 0);
    drop(tmp);
}

// ======================= Unsupported cases =======================

#[tokio::test]
async fn test_update_rejects_primary_key_table_without_data_evolution() {
    let (tmp, catalog) = create_test_env();
    let sql_context = create_sql_context(catalog).await;
    sql_context
        .sql("CREATE SCHEMA paimon.test_db")
        .await
        .unwrap();
    sql_context
        .sql(
            "CREATE TABLE paimon.test_db.pk_t (\
                id INT NOT NULL, name VARCHAR, PRIMARY KEY (id)\
            ) WITH ('bucket' = '1')",
        )
        .await
        .unwrap();

    assert_sql_error(
        &sql_context,
        "UPDATE paimon.test_db.pk_t SET name = 'x' WHERE id = 1",
        "primary-key tables without data-evolution",
    )
    .await;
    drop(tmp);
}

#[tokio::test]
async fn test_update_rejects_primary_key_table_with_data_evolution() {
    let (tmp, catalog) = create_test_env();
    let sql_context = create_sql_context(catalog).await;
    sql_context
        .sql("CREATE SCHEMA paimon.test_db")
        .await
        .unwrap();
    sql_context
        .sql(
            "CREATE TABLE paimon.test_db.pk_de_t (\
                id INT NOT NULL, name VARCHAR, PRIMARY KEY (id)\
            ) WITH (\
                'bucket' = '1',\
                'row-tracking.enabled' = 'true',\
                'data-evolution.enabled' = 'true'\
            )",
        )
        .await
        .unwrap();

    assert_sql_error(
        &sql_context,
        "UPDATE paimon.test_db.pk_de_t SET name = 'x' WHERE id = 1",
        "does not support primary keys",
    )
    .await;
    drop(tmp);
}

#[tokio::test]
async fn test_update_rejects_table_alias() {
    let (_tmp, sql_context) = setup().await;

    assert_sql_error(
        &sql_context,
        "UPDATE paimon.test_db.t AS target SET name = 'x' WHERE id = 1",
        "Table alias 'target' in UPDATE is not yet supported",
    )
    .await;
}