datafusion_expr/logical_plan/

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

use std::cmp::Ordering;
use std::collections::HashMap;
use std::fmt::{self, Debug, Display, Formatter};
use std::hash::{Hash, Hasher};
use std::sync::Arc;

use arrow::datatypes::{DataType, Field, Schema};
use datafusion_common::file_options::file_type::FileType;
use datafusion_common::{DFSchemaRef, Result, TableReference, internal_err};

use crate::{Expr, LogicalPlan, TableSource};

/// Operator that copies the contents of a database to file(s)
#[derive(Clone)]
pub struct CopyTo {
    /// The relation that determines the tuples to write to the output file(s)
    pub input: Arc<LogicalPlan>,
    /// The location to write the file(s)
    pub output_url: String,
    /// Determines which, if any, columns should be used for hive-style partitioned writes
    pub partition_by: Vec<String>,
    /// File type trait
    pub file_type: Arc<dyn FileType>,
    /// SQL Options that can affect the formats
    pub options: HashMap<String, String>,
    /// The schema of the output (a single column "count")
    pub output_schema: DFSchemaRef,
}

impl Debug for CopyTo {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("CopyTo")
            .field("input", &self.input)
            .field("output_url", &self.output_url)
            .field("partition_by", &self.partition_by)
            .field("file_type", &"...")
            .field("options", &self.options)
            .field("output_schema", &self.output_schema)
            .finish_non_exhaustive()
    }
}

// Implement PartialEq manually
impl PartialEq for CopyTo {
    fn eq(&self, other: &Self) -> bool {
        self.input == other.input && self.output_url == other.output_url
    }
}

// Implement Eq (no need for additional logic over PartialEq)
impl Eq for CopyTo {}

// Manual implementation needed because of `file_type` and `options` fields.
// Comparison excludes these field.
impl PartialOrd for CopyTo {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        match self.input.partial_cmp(&other.input) {
            Some(Ordering::Equal) => match self.output_url.partial_cmp(&other.output_url)
            {
                Some(Ordering::Equal) => {
                    self.partition_by.partial_cmp(&other.partition_by)
                }
                cmp => cmp,
            },
            cmp => cmp,
        }
        // TODO (https://github.com/apache/datafusion/issues/17477) avoid recomparing all fields
        .filter(|cmp| *cmp != Ordering::Equal || self == other)
    }
}

// Implement Hash manually
impl Hash for CopyTo {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.input.hash(state);
        self.output_url.hash(state);
    }
}

impl CopyTo {
    pub fn new(
        input: Arc<LogicalPlan>,
        output_url: String,
        partition_by: Vec<String>,
        file_type: Arc<dyn FileType>,
        options: HashMap<String, String>,
    ) -> Self {
        Self {
            input,
            output_url,
            partition_by,
            file_type,
            options,
            // The output schema is always a single column "count" with the number of rows copied
            output_schema: make_count_schema(),
        }
    }
}

/// Modifies the content of a database
///
/// This operator is used to perform DML operations such as INSERT, DELETE,
/// UPDATE, and CTAS (CREATE TABLE AS SELECT).
///
/// * `INSERT` - Appends new rows to the existing table. Calls
///   [`TableProvider::insert_into`]
///
/// * `DELETE` - Removes rows from the table. Calls [`TableProvider::delete_from`]
///
/// * `UPDATE` - Modifies existing rows in the table. Calls [`TableProvider::update`]
///
/// * `CREATE TABLE AS SELECT` - Creates a new table and populates it with data
///   from a query. This is similar to the `INSERT` operation, but it creates a new
///   table instead of modifying an existing one.
///
/// Note that the structure is adapted from substrait WriteRel)
///
/// [`TableProvider`]: https://docs.rs/datafusion/latest/datafusion/datasource/trait.TableProvider.html
/// [`TableProvider::insert_into`]: https://docs.rs/datafusion/latest/datafusion/datasource/trait.TableProvider.html#method.insert_into
/// [`TableProvider::delete_from`]: https://docs.rs/datafusion/latest/datafusion/datasource/trait.TableProvider.html#method.delete_from
/// [`TableProvider::update`]: https://docs.rs/datafusion/latest/datafusion/datasource/trait.TableProvider.html#method.update
#[derive(Clone)]
pub struct DmlStatement {
    /// The table name
    pub table_name: TableReference,
    /// this is target table to insert into
    pub target: Arc<dyn TableSource>,
    /// The type of operation to perform
    pub op: WriteOp,
    /// The relation that determines the tuples to add/remove/modify the schema must match with table_schema
    pub input: Arc<LogicalPlan>,
    /// The schema of the output relation
    pub output_schema: DFSchemaRef,
}
impl Eq for DmlStatement {}
impl Hash for DmlStatement {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.table_name.hash(state);
        self.target.schema().hash(state);
        self.op.hash(state);
        self.input.hash(state);
        self.output_schema.hash(state);
    }
}

impl PartialEq for DmlStatement {
    fn eq(&self, other: &Self) -> bool {
        self.table_name == other.table_name
            && self.target.schema() == other.target.schema()
            && self.op == other.op
            && self.input == other.input
            && self.output_schema == other.output_schema
    }
}

impl Debug for DmlStatement {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_struct("DmlStatement")
            .field("table_name", &self.table_name)
            .field("target", &"...")
            .field("target_schema", &self.target.schema())
            .field("op", &self.op)
            .field("input", &self.input)
            .field("output_schema", &self.output_schema)
            .finish()
    }
}

impl DmlStatement {
    /// Creates a new DML statement with the output schema set to a single `count` column.
    pub fn new(
        table_name: TableReference,
        target: Arc<dyn TableSource>,
        op: WriteOp,
        input: Arc<LogicalPlan>,
    ) -> Self {
        Self {
            table_name,
            target,
            op,
            input,

            // The output schema is always a single column with the number of rows affected
            output_schema: make_count_schema(),
        }
    }

    /// Return a descriptive name of this [`DmlStatement`]
    pub fn name(&self) -> &str {
        self.op.name()
    }
}

// Manual implementation needed because of `table_schema` and `output_schema` fields.
// Comparison excludes these fields.
impl PartialOrd for DmlStatement {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        match self.table_name.partial_cmp(&other.table_name) {
            Some(Ordering::Equal) => match self.op.partial_cmp(&other.op) {
                Some(Ordering::Equal) => self.input.partial_cmp(&other.input),
                cmp => cmp,
            },
            cmp => cmp,
        }
        // TODO (https://github.com/apache/datafusion/issues/17477) avoid recomparing all fields
        .filter(|cmp| *cmp != Ordering::Equal || self == other)
    }
}

/// The type of DML operation to perform.
///
/// See [`DmlStatement`] for more details.
///
/// Marked `#[non_exhaustive]` so adding new variants in future releases is
/// not a SemVer break for downstream matchers.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
#[non_exhaustive]
pub enum WriteOp {
    /// `INSERT INTO` operation
    Insert(InsertOp),
    /// `DELETE` operation
    Delete,
    /// `UPDATE` operation
    Update,
    /// `CREATE TABLE AS SELECT` operation
    Ctas,
    /// `TRUNCATE` operation
    Truncate,
    /// `MERGE INTO` operation
    MergeInto(Box<MergeIntoOp>),
}

impl WriteOp {
    /// Return a descriptive name of this [`WriteOp`]
    pub fn name(&self) -> &str {
        match self {
            WriteOp::Insert(insert) => insert.name(),
            WriteOp::Delete => "Delete",
            WriteOp::Update => "Update",
            WriteOp::Ctas => "Ctas",
            WriteOp::Truncate => "Truncate",
            WriteOp::MergeInto(_) => "MergeInto",
        }
    }
}

impl Display for WriteOp {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)]
pub enum InsertOp {
    /// Appends new rows to the existing table without modifying any
    /// existing rows. This corresponds to the SQL `INSERT INTO` query.
    Append,
    /// Overwrites all existing rows in the table with the new rows.
    /// This corresponds to the SQL `INSERT OVERWRITE` query.
    Overwrite,
    /// If any existing rows collides with the inserted rows (typically based
    /// on a unique key or primary key), those existing rows are replaced.
    /// This corresponds to the SQL `REPLACE INTO` query and its equivalents.
    Replace,
}

impl InsertOp {
    /// Return a descriptive name of this [`InsertOp`]
    pub fn name(&self) -> &str {
        match self {
            InsertOp::Append => "Insert Into",
            InsertOp::Overwrite => "Insert Overwrite",
            InsertOp::Replace => "Replace Into",
        }
    }
}

impl Display for InsertOp {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

/// Describes a MERGE INTO operation's parameters.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub struct MergeIntoOp {
    /// The join condition from `ON <expr>`.
    pub on: Expr,
    /// The WHEN clauses, in the order they appeared in the SQL.
    pub clauses: Vec<MergeIntoClause>,
}

impl MergeIntoOp {
    /// Count of top-level [`Expr`]s owned by this operation (no allocation).
    ///
    /// Matches the length of [`Self::exprs`] and the `exprs` vec consumed by
    /// [`Self::with_new_exprs`].
    fn expr_count(&self) -> usize {
        1 + self
            .clauses
            .iter()
            .map(|c| {
                c.predicate.is_some() as usize
                    + match &c.action {
                        MergeIntoAction::Update(a) => a.len(),
                        MergeIntoAction::Insert { values, .. } => values.len(),
                        MergeIntoAction::Delete => 0,
                    }
            })
            .sum::<usize>()
    }

    /// Top-level [`Expr`]s in stable order: `on`, then per-clause predicate
    /// (if any) and action value expressions.
    pub fn exprs(&self) -> Vec<&Expr> {
        let mut out = Vec::with_capacity(self.expr_count());
        out.push(&self.on);
        for clause in &self.clauses {
            if let Some(predicate) = &clause.predicate {
                out.push(predicate);
            }
            match &clause.action {
                MergeIntoAction::Update(assignments) => {
                    out.extend(assignments.iter().map(|(_, value)| value));
                }
                MergeIntoAction::Insert { values, .. } => {
                    out.extend(values.iter());
                }
                MergeIntoAction::Delete => {}
            }
        }
        out
    }

    /// Rebuild this `MergeIntoOp` from a flat vector of new expressions, in
    /// the same order produced by [`Self::exprs`]. The clause kinds, action
    /// kinds, column lists, and presence/absence of each predicate are
    /// preserved from `self`.
    pub fn with_new_exprs(&self, exprs: Vec<Expr>) -> Result<Self> {
        let expected = self.expr_count();
        if exprs.len() != expected {
            return internal_err!(
                "MergeIntoOp::with_new_exprs expected {expected} expressions, got {}",
                exprs.len()
            );
        }
        let mut iter = exprs.into_iter();
        let on = iter.next().expect("non-empty by length check");
        let clauses = self
            .clauses
            .iter()
            .map(|clause| {
                let predicate = clause
                    .predicate
                    .is_some()
                    .then(|| iter.next().expect("non-empty by length check"));
                let action = match &clause.action {
                    MergeIntoAction::Update(assignments) => {
                        let assignments = assignments
                            .iter()
                            .map(|(name, _)| {
                                (
                                    name.clone(),
                                    iter.next().expect("non-empty by length check"),
                                )
                            })
                            .collect();
                        MergeIntoAction::Update(assignments)
                    }
                    MergeIntoAction::Insert { columns, values } => {
                        let values = values
                            .iter()
                            .map(|_| iter.next().expect("non-empty by length check"))
                            .collect();
                        MergeIntoAction::Insert {
                            columns: columns.clone(),
                            values,
                        }
                    }
                    MergeIntoAction::Delete => MergeIntoAction::Delete,
                };
                MergeIntoClause {
                    kind: clause.kind,
                    predicate,
                    action,
                }
            })
            .collect();
        Ok(Self { on, clauses })
    }
}

/// A single WHEN clause within a MERGE INTO statement.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub struct MergeIntoClause {
    /// Whether this fires on matched or unmatched rows.
    pub kind: MergeIntoClauseKind,
    /// Optional additional predicate (`AND <expr>`).
    pub predicate: Option<Expr>,
    /// The action to take.
    pub action: MergeIntoAction,
}

/// Which rows a MERGE WHEN clause applies to.
///
/// Mirrors `sqlparser::ast::MergeClauseKind` so that the SQL spelling is
/// preserved through the logical plan.
///
/// **Note on `NotMatched` vs `NotMatchedByTarget`:** these two variants are
/// semantically identical — both describe a source row that has no matching
/// target row. `NotMatched` is the SQL standard short form (used by
/// Snowflake, Postgres, SQL Server); `NotMatchedByTarget` is BigQuery's
/// explicit form added for symmetry with `NotMatchedBySource`. Downstream
/// consumers (planners, table providers, optimizers) MUST treat the two
/// variants identically.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)]
pub enum MergeIntoClauseKind {
    /// `WHEN MATCHED`
    Matched,
    /// `WHEN NOT MATCHED` — see type-level note for the equivalence with
    /// [`NotMatchedByTarget`](Self::NotMatchedByTarget).
    NotMatched,
    /// `WHEN NOT MATCHED BY TARGET` — see type-level note for the
    /// equivalence with [`NotMatched`](Self::NotMatched).
    NotMatchedByTarget,
    /// `WHEN NOT MATCHED BY SOURCE`
    NotMatchedBySource,
}

impl MergeIntoClauseKind {
    /// True if this clause fires on a source row that has no matching target
    /// row. Returns `true` for both [`NotMatched`](Self::NotMatched) and
    /// [`NotMatchedByTarget`](Self::NotMatchedByTarget) (see the type-level
    /// note explaining why those two variants are semantically identical).
    ///
    /// Prefer this predicate over hand-written `matches!` arms so the
    /// `NotMatched`/`NotMatchedByTarget` equivalence is enforced in one place.
    pub fn is_not_matched_by_target(&self) -> bool {
        matches!(self, Self::NotMatched | Self::NotMatchedByTarget)
    }

    /// Collapse the SQL-spelling variants into the canonical three semantic
    /// categories: [`Matched`](Self::Matched),
    /// [`NotMatchedByTarget`](Self::NotMatchedByTarget) (covering both
    /// "NOT MATCHED" spellings), and
    /// [`NotMatchedBySource`](Self::NotMatchedBySource).
    ///
    /// Use this in downstream `match` expressions when the SQL spelling
    /// distinction does not matter — e.g. in planners, optimizers, or
    /// table-provider dispatch.
    pub fn canonical(self) -> Self {
        match self {
            Self::NotMatched => Self::NotMatchedByTarget,
            other => other,
        }
    }
}

/// The action for a single WHEN clause.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)]
pub enum MergeIntoAction {
    /// `UPDATE SET col1 = expr1, col2 = expr2, ...`, stored as
    /// `(column_name, value_expr)` pairs.
    Update(Vec<(String, Expr)>),
    /// `INSERT (col1, col2, ...) VALUES (expr1, expr2, ...)`. `columns` may
    /// be empty, meaning all columns.
    Insert {
        columns: Vec<String>,
        values: Vec<Expr>,
    },
    Delete,
}

fn make_count_schema() -> DFSchemaRef {
    Arc::new(
        Schema::new(vec![Field::new("count", DataType::UInt64, false)])
            .try_into()
            .unwrap(),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{col, lit};

    #[test]
    fn write_op_merge_into_name_and_display() {
        let op = WriteOp::MergeInto(Box::new(MergeIntoOp {
            on: col("id").eq(col("source_id")),
            clauses: vec![MergeIntoClause {
                kind: MergeIntoClauseKind::Matched,
                predicate: Some(col("qty").gt(lit(0_i64))),
                action: MergeIntoAction::Update(vec![(
                    "qty".to_string(),
                    col("source_qty"),
                )]),
            }],
        }));
        assert_eq!(op.name(), "MergeInto");
        assert_eq!(format!("{op}"), "MergeInto");
    }

    #[test]
    fn merge_into_clause_kind_is_not_matched_by_target() {
        assert!(!MergeIntoClauseKind::Matched.is_not_matched_by_target());
        assert!(MergeIntoClauseKind::NotMatched.is_not_matched_by_target());
        assert!(MergeIntoClauseKind::NotMatchedByTarget.is_not_matched_by_target());
        assert!(!MergeIntoClauseKind::NotMatchedBySource.is_not_matched_by_target());
    }

    #[test]
    fn merge_into_clause_kind_canonical_collapses_not_matched() {
        assert_eq!(
            MergeIntoClauseKind::NotMatched.canonical(),
            MergeIntoClauseKind::NotMatchedByTarget
        );
        assert_eq!(
            MergeIntoClauseKind::NotMatchedByTarget.canonical(),
            MergeIntoClauseKind::NotMatchedByTarget
        );
        assert_eq!(
            MergeIntoClauseKind::Matched.canonical(),
            MergeIntoClauseKind::Matched
        );
        assert_eq!(
            MergeIntoClauseKind::NotMatchedBySource.canonical(),
            MergeIntoClauseKind::NotMatchedBySource
        );
    }

    #[test]
    fn merge_into_op_exprs_round_trip() {
        let op = MergeIntoOp {
            on: col("id").eq(col("source_id")),
            clauses: vec![
                MergeIntoClause {
                    kind: MergeIntoClauseKind::Matched,
                    predicate: Some(col("qty").gt(lit(0_i64))),
                    action: MergeIntoAction::Update(vec![
                        ("qty".to_string(), col("source_qty")),
                        ("price".to_string(), col("source_price")),
                    ]),
                },
                MergeIntoClause {
                    kind: MergeIntoClauseKind::NotMatched,
                    predicate: None,
                    action: MergeIntoAction::Insert {
                        columns: vec!["id".to_string(), "qty".to_string()],
                        values: vec![col("source_id"), col("source_qty")],
                    },
                },
                MergeIntoClause {
                    kind: MergeIntoClauseKind::NotMatchedBySource,
                    predicate: Some(col("active").eq(lit(true))),
                    action: MergeIntoAction::Delete,
                },
            ],
        };
        let exprs = op.exprs();
        assert_eq!(exprs.len(), 7);

        let owned: Vec<Expr> = exprs.into_iter().cloned().collect();
        let rebuilt = op.with_new_exprs(owned).unwrap();
        assert_eq!(op, rebuilt);
    }

    #[test]
    fn merge_into_op_with_new_exprs_length_mismatch() {
        let op = MergeIntoOp {
            on: col("id").eq(col("source_id")),
            clauses: vec![],
        };
        let err = op.with_new_exprs(vec![]).unwrap_err();
        assert!(
            err.to_string().contains("expected 1 expressions, got 0"),
            "unexpected error: {err}"
        );
    }
}