feat: export the data from a table to parquet files (#1000)

* feat: copy table parser

* feat: coopy table

* chore: minor fix

* chore: give stmt a more clearer name

* chore: unified naming

* chore: minor change

* chore: add a todo

* chore: end up with an empty file when occur an empty table

* feat: format with copy table

* feat: with options

* chore: by cr

* chore: default 5M rows per segment

* Update src/datanode/src/sql/copy_table.rs

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>

* Update src/datanode/src/sql/copy_table.rs

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>

* Update src/datanode/src/error.rs

Co-authored-by: fys <40801205+Fengys123@users.noreply.github.com>

---------

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>
Co-authored-by: fys <40801205+Fengys123@users.noreply.github.com>

src/sql/src/error.rs

@@ -135,6 +135,9 @@ pub enum Error {
         timestamp: Timestamp,
         target_unit: TimeUnit,
     },
+
+    #[snafu(display("Unsupported format option: {}", name))]
+    UnsupportedCopyFormatOption { name: String },
 }
 
 impl ErrorExt for Error {
@@ -156,12 +159,14 @@ impl ErrorExt for Error {
             InvalidColumnOption { .. }
             | InvalidDatabaseName { .. }
             | ColumnTypeMismatch { .. }
-            | InvalidTableName { .. } => StatusCode::InvalidArguments,
+            | InvalidTableName { .. }
+            | InvalidSqlValue { .. }
+            | TimestampOverflow { .. }
+            | UnsupportedCopyFormatOption { .. } => StatusCode::InvalidArguments,
+
             UnsupportedAlterTableStatement { .. } => StatusCode::InvalidSyntax,
             SerializeColumnDefaultConstraint { source, .. } => source.status_code(),
             ConvertToGrpcDataType { source, .. } => source.status_code(),
-            InvalidSqlValue { .. } => StatusCode::InvalidArguments,
-            TimestampOverflow { .. } => StatusCode::InvalidArguments,
         }
     }
 

src/sql/src/parser.rs

@@ -117,6 +117,8 @@ impl<'a> ParserContext<'a> {
                         Ok(Statement::Use(database_name.value))
                     }
 
+                    Keyword::COPY => self.parse_copy(),
+
                     // todo(hl) support more statements.
                     _ => self.unsupported(self.peek_token_as_string()),
                 }

src/sql/src/parsers.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod alter_parser;
+pub(crate) mod copy_parser;
 pub(crate) mod create_parser;
 pub(crate) mod delete_parser;
 pub(crate) mod insert_parser;

src/sql/src/parsers/copy_parser.rs

@@ -0,0 +1,132 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed 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 snafu::ResultExt;
+use sqlparser::ast::Value;
+use sqlparser::keywords::Keyword;
+
+use crate::error::{self, Result};
+use crate::parser::ParserContext;
+use crate::statements::copy::{CopyTable, Format};
+use crate::statements::statement::Statement;
+
+// COPY tbl TO 'output.parquet';
+impl<'a> ParserContext<'a> {
+    pub(crate) fn parse_copy(&mut self) -> Result<Statement> {
+        self.parser.next_token();
+        let copy_table = self.parse_copy_table()?;
+        Ok(Statement::Copy(copy_table))
+    }
+
+    fn parse_copy_table(&mut self) -> Result<CopyTable> {
+        let table_name =
+            self.parser
+                .parse_object_name()
+                .with_context(|_| error::UnexpectedSnafu {
+                    sql: self.sql,
+                    expected: "a table name",
+                    actual: self.peek_token_as_string(),
+                })?;
+
+        self.parser
+            .expect_keyword(Keyword::TO)
+            .context(error::SyntaxSnafu { sql: self.sql })?;
+
+        let file_name =
+            self.parser
+                .parse_literal_string()
+                .with_context(|_| error::UnexpectedSnafu {
+                    sql: self.sql,
+                    expected: "a file name",
+                    actual: self.peek_token_as_string(),
+                })?;
+
+        let options = self
+            .parser
+            .parse_options(Keyword::WITH)
+            .context(error::SyntaxSnafu { sql: self.sql })?;
+
+        // default format is parquet
+        let mut format = Format::Parquet;
+        for option in options {
+            if option.name.value.to_uppercase() == "FORMAT" {
+                if let Value::SingleQuotedString(fmt_str) = option.value {
+                    format = Format::try_from(fmt_str)?;
+                }
+            }
+        }
+
+        Ok(CopyTable::new(table_name, file_name, format))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::assert_matches::assert_matches;
+
+    use sqlparser::dialect::GenericDialect;
+
+    use super::*;
+
+    #[test]
+    fn test_parse_copy_table() {
+        let sql0 = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet'";
+        let sql1 = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'parquet')";
+        let result0 = ParserContext::create_with_dialect(sql0, &GenericDialect {}).unwrap();
+        let result1 = ParserContext::create_with_dialect(sql1, &GenericDialect {}).unwrap();
+
+        for mut result in vec![result0, result1] {
+            assert_eq!(1, result.len());
+
+            let statement = result.remove(0);
+            assert_matches!(statement, Statement::Copy { .. });
+            match statement {
+                Statement::Copy(copy_table) => {
+                    let (catalog, schema, table) =
+                        if let [catalog, schema, table] = &copy_table.table_name().0[..] {
+                            (
+                                catalog.value.clone(),
+                                schema.value.clone(),
+                                table.value.clone(),
+                            )
+                        } else {
+                            unreachable!()
+                        };
+
+                    assert_eq!("catalog0", catalog);
+                    assert_eq!("schema0", schema);
+                    assert_eq!("tbl", table);
+
+                    let file_name = copy_table.file_name();
+                    assert_eq!("tbl_file.parquet", file_name);
+
+                    let format = copy_table.format();
+                    assert_eq!(Format::Parquet, *format);
+                }
+                _ => unreachable!(),
+            }
+        }
+    }
+
+    #[test]
+    fn test_parse_copy_table_with_unsupopoted_format() {
+        let sql = "COPY catalog0.schema0.tbl TO 'tbl_file.parquet' WITH (FORMAT = 'unknow_format')";
+        let result = ParserContext::create_with_dialect(sql, &GenericDialect {});
+        assert!(result.is_err());
+        assert_matches!(
+            result.err().unwrap(),
+            error::Error::UnsupportedCopyFormatOption { .. }
+        );
+    }
+}

src/sql/src/statements.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 pub mod alter;
+pub mod copy;
 pub mod create;
 pub mod delete;
 pub mod describe;

src/sql/src/statements/copy.rs

@@ -0,0 +1,62 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed 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 sqlparser::ast::ObjectName;
+
+use crate::error::{self, Result};
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CopyTable {
+    table_name: ObjectName,
+    file_name: String,
+    format: Format,
+}
+
+impl CopyTable {
+    pub(crate) fn new(table_name: ObjectName, file_name: String, format: Format) -> Self {
+        Self {
+            table_name,
+            file_name,
+            format,
+        }
+    }
+
+    pub fn table_name(&self) -> &ObjectName {
+        &self.table_name
+    }
+
+    pub fn file_name(&self) -> &str {
+        &self.file_name
+    }
+
+    pub fn format(&self) -> &Format {
+        &self.format
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum Format {
+    Parquet,
+}
+
+impl TryFrom<String> for Format {
+    type Error = error::Error;
+
+    fn try_from(name: String) -> Result<Self> {
+        match name.to_uppercase().as_str() {
+            "PARQUET" => Ok(Format::Parquet),
+            _ => error::UnsupportedCopyFormatOptionSnafu { name }.fail(),
+        }
+    }
+}

src/sql/src/statements/statement.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use crate::statements::alter::AlterTable;
+use crate::statements::copy::CopyTable;
 use crate::statements::create::{CreateDatabase, CreateTable};
 use crate::statements::delete::Delete;
 use crate::statements::describe::DescribeTable;
@@ -51,6 +52,8 @@ pub enum Statement {
     // EXPLAIN QUERY
     Explain(Explain),
     Use(String),
+    // COPY
+    Copy(CopyTable),
 }
 
 /// Comment hints from SQL.