1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

// 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::{
    fmt::{self, Display},
    sync::Arc,
};

use datafusion_common::{
    file_options::StatementOptions, DFSchemaRef, FileType, FileTypeWriterOptions,
    OwnedTableReference,
};

use crate::LogicalPlan;

/// Operator that copies the contents of a database to file(s)
#[derive(Clone, PartialEq, Eq, Hash)]
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,
    /// The file format to output (explicitly defined or inferred from file extension)
    pub file_format: FileType,
    /// If false, it is assumed output_url is a file to which all data should be written
    /// regardless of input partitioning. Otherwise, output_url is assumed to be a directory
    /// to which each output partition is written to its own output file
    pub single_file_output: bool,
    /// Arbitrary options as tuples
    pub copy_options: CopyOptions,
}

/// When the logical plan is constructed from SQL, CopyOptions
/// will contain arbitrary string tuples which must be parsed into
/// FileTypeWriterOptions. When the logical plan is constructed directly
/// from rust code (such as via the DataFrame API), FileTypeWriterOptions
/// can be provided directly, avoiding the run time cost and fallibility of
/// parsing string based options.
#[derive(Clone)]
pub enum CopyOptions {
    /// Holds StatementOptions parsed from a SQL statement
    SQLOptions(StatementOptions),
    /// Holds FileTypeWriterOptions directly provided
    WriterOptions(Box<FileTypeWriterOptions>),
}

impl PartialEq for CopyOptions {
    fn eq(&self, other: &CopyOptions) -> bool {
        match self {
            Self::SQLOptions(statement1) => match other {
                Self::SQLOptions(statement2) => statement1.eq(statement2),
                Self::WriterOptions(_) => false,
            },
            Self::WriterOptions(_) => false,
        }
    }
}

impl Eq for CopyOptions {}

impl std::hash::Hash for CopyOptions {
    fn hash<H>(&self, hasher: &mut H)
    where
        H: std::hash::Hasher,
    {
        match self {
            Self::SQLOptions(statement) => statement.hash(hasher),
            Self::WriterOptions(_) => (),
        }
    }
}

/// The operator that modifies the content of a database (adapted from
/// substrait WriteRel)
#[derive(Clone, PartialEq, Eq, Hash)]
pub struct DmlStatement {
    /// The table name
    pub table_name: OwnedTableReference,
    /// The schema of the table (must align with Rel input)
    pub table_schema: DFSchemaRef,
    /// 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>,
}

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

#[derive(Clone, PartialEq, Eq, Hash)]
pub enum WriteOp {
    InsertOverwrite,
    InsertInto,
    Delete,
    Update,
    Ctas,
}

impl WriteOp {
    /// Return a descriptive name of this [`WriteOp`]
    pub fn name(&self) -> &str {
        match self {
            WriteOp::InsertOverwrite => "Insert Overwrite",
            WriteOp::InsertInto => "Insert Into",
            WriteOp::Delete => "Delete",
            WriteOp::Update => "Update",
            WriteOp::Ctas => "Ctas",
        }
    }
}

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