Add COPY .. TO .. syntax support

Author: alamb

datafusion/core/tests/sqllogictests/test_files/copy.slt

@@ -0,0 +1,80 @@
+# 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.
+
+# tests for copy command
+
+statement ok
+create table source_table(col1 integer, col2 varchar) as values (1, 'Foo'), (2, 'Bar');
+
+# create an external parquet file
+# TODO figure out how to use a real tempfile name
+#statement ok
+#COPY source_table to '/tmp/table.parquet';
+
+
+# create an external parquet file from a query
+# TODO figure out how to use a real tempfile name
+statement ok
+COPY (select col2, sum(col1) from source_table group by col2 order by col2)  to '/tmp/table.parquet';
+
+# Read the data back in
+statement ok
+CREATE external table new_table STORED as PARQUET LOCATION '/tmp/table.parquet';
+
+# Read the data back in
+query IT
+select * from new_table;
+----
+1 Foo
+2 Bar
+
+statement ok
+drop table new_table;
+
+
+# TODO error cases:
+# incomplete statement
+# pass non literal to COPY options
+
+
+
+
+
+# TODO add support for reading directly for files without needing to create an external table
+# (TODO find ticket)
+
+
+
+# TODO test copying to partitioned target (directory)
+
+
+# TODO test copying with options
+
+
+
+# TODO test query source
+
+
+
+# TODO: support CSV, JSON and AVRO
+
+# TODO: document different options
+
+
+
+statement ok
+drop table source_table;

datafusion/sql/src/parser.rs

@@ -18,7 +18,7 @@
 //! DataFusion SQL Parser based on [`sqlparser`]
 
 use datafusion_common::parsers::CompressionTypeVariant;
-use sqlparser::ast::OrderByExpr;
+use sqlparser::ast::{OrderByExpr, Query, Value};
 use sqlparser::{
     ast::{
         ColumnDef, ColumnOptionDef, ObjectName, Statement as SQLStatement,
@@ -42,6 +42,46 @@ fn parse_file_type(s: &str) -> Result<String, ParserError> {
     Ok(s.to_uppercase())
 }
 
+/// DataFusion extension DDL for `COPY`
+///
+/// Syntax:
+///
+/// ```text
+/// COPY <table_name | (<query>)>
+/// TO
+/// <destination_url>
+/// (key_value_list)
+///
+/// ```
+/// Examples
+/// ``sql
+/// COPY lineitem  TO 'lineitem'
+///  (format parquet,
+///   partitions 16,
+///   row_group_limit_rows 100000,
+//    row_group_limit_bytes
+///  )
+///
+/// COPY (SELECT l_orderkey from lineitem) to 'lineitem.parquet';
+/// ```
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct CopyToStatement {
+    /// From where the data comes from
+    pub source: CopyToSource,
+    /// The URL to where the data is heading
+    pub target: String,
+    /// Target specific options
+    pub options: HashMap<String, Value>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum CopyToSource {
+    /// `COPY <table> TO ...`
+    Relation(ObjectName),
+    /// COPY (query...) TO ...
+    Query(Query),
+}
+
 /// DataFusion extension DDL for `CREATE EXTERNAL TABLE`
 ///
 /// Syntax:
@@ -117,12 +157,14 @@ pub struct DescribeTableStmt {
 /// Tokens parsed by [`DFParser`] are converted into these values.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Statement {
-    /// ANSI SQL AST node
+    /// ANSI SQL AST node (from sqlparser-rs)
     Statement(Box<SQLStatement>),
     /// Extension: `CREATE EXTERNAL TABLE`
     CreateExternalTable(CreateExternalTable),
     /// Extension: `DESCRIBE TABLE`
     DescribeTableStmt(DescribeTableStmt),
+    /// Extension: `COPY TO`
+    CopyTo(CopyToStatement),
 }
 
 /// DataFusion SQL Parser based on [`sqlparser`]
@@ -211,6 +253,11 @@ impl<'a> DFParser<'a> {
                         // use custom parsing
                         self.parse_create()
                     }
+                    Keyword::COPY => {
+                        // move one token forward
+                        self.parser.next_token();
+                        self.parse_copy()
+                    }
                     Keyword::DESCRIBE => {
                         // move one token forward
                         self.parser.next_token();
@@ -242,6 +289,37 @@ impl<'a> DFParser<'a> {
         }))
     }
 
+    /// Parse a SQL `COPY TO` statement
+    pub fn parse_copy(&mut self) -> Result<Statement, ParserError> {
+        // parse as a query
+        let source = if self.parser.consume_token(&Token::LParen) {
+            let query = self.parser.parse_query()?;
+            self.parser.expect_token(&Token::RParen)?;
+            CopyToSource::Query(query)
+        } else {
+            // parse as table reference
+            let table_name = self.parser.parse_object_name()?;
+            CopyToSource::Relation(table_name)
+        };
+
+        self.parser.expect_keyword(Keyword::TO)?;
+
+        let target = self.parser.parse_literal_string()?;
+
+        // check for options in parens
+        let options = if self.parser.peek_token().token == Token::LParen {
+            self.parse_value_options()?
+        } else {
+            HashMap::new()
+        };
+
+        Ok(Statement::CopyTo(CopyToStatement {
+            source,
+            target,
+            options,
+        }))
+    }
+
     /// Parse a SQL `CREATE` statement handling `CREATE EXTERNAL TABLE`
     pub fn parse_create(&mut self) -> Result<Statement, ParserError> {
         if self.parser.parse_keyword(Keyword::EXTERNAL) {
@@ -457,7 +535,7 @@ impl<'a> DFParser<'a> {
                 builder.table_partition_cols = Some(self.parse_partitions()?)
             } else if self.parser.parse_keyword(Keyword::OPTIONS) {
                 ensure_not_set(&builder.options, "OPTIONS")?;
-                builder.options = Some(self.parse_options()?);
+                builder.options = Some(self.parse_string_options()?);
             } else {
                 break;
             }
@@ -513,14 +591,40 @@ impl<'a> DFParser<'a> {
         }
     }
 
-    fn parse_options(&mut self) -> Result<HashMap<String, String>, ParserError> {
-        let mut options: HashMap<String, String> = HashMap::new();
+    /// Parses (key value) style options, but values can only be literal strings
+    /// TODO maybe change this to be real expressions rather than just strings
+    /// the reason
+    fn parse_string_options(&mut self) -> Result<HashMap<String, String>, ParserError> {
+        let mut options = HashMap::new();
         self.parser.expect_token(&Token::LParen)?;
 
         loop {
             let key = self.parser.parse_literal_string()?;
             let value = self.parser.parse_literal_string()?;
-            options.insert(key.to_string(), value.to_string());
+            options.insert(key, value);
+            let comma = self.parser.consume_token(&Token::Comma);
+            if self.parser.consume_token(&Token::RParen) {
+                // allow a trailing comma, even though it's not in standard
+                break;
+            } else if !comma {
+                return self.expected(
+                    "',' or ')' after option definition",
+                    self.parser.peek_token(),
+                );
+            }
+        }
+        Ok(options)
+    }
+
+    /// parses (foo bar) style options into a map of String --> [`Value`]
+    fn parse_value_options(&mut self) -> Result<HashMap<String, Value>, ParserError> {
+        let mut options = HashMap::new();
+        self.parser.expect_token(&Token::LParen)?;
+
+        loop {
+            let key = self.parser.parse_literal_string()?;
+            let value = self.parser.parse_value()?;
+            options.insert(key, value);
             let comma = self.parser.consume_token(&Token::Comma);
             if self.parser.consume_token(&Token::RParen) {
                 // allow a trailing comma, even though it's not in standard
@@ -560,7 +664,7 @@ mod tests {
             1,
             "Expected to parse exactly one statement"
         );
-        assert_eq!(statements[0], expected);
+        assert_eq!(statements[0], expected, "actual:\n{:#?}", statements[0]);
         Ok(())
     }
 
@@ -980,4 +1084,60 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn copy_to_table_to_table() -> Result<(), ParserError> {
+        // positive case
+        let sql = "COPY foo TO bar";
+        let expected = Statement::CopyTo(CopyToStatement {
+            source: object_name("foo"),
+            target: "bar".to_string(),
+            options: HashMap::new(),
+        });
+
+        expect_parse_ok(sql, expected)?;
+        Ok(())
+    }
+
+    #[test]
+    fn copy_to_query_to_table() -> Result<(), ParserError> {
+        let mut statements = Parser::parse_sql(&GenericDialect {}, "select 1")?;
+        assert_eq!(statements.len(), 1);
+        let statement = statements.pop().unwrap();
+        let query = if let SQLStatement::Query(query) = statement {
+            *query
+        } else {
+            panic!("Expected query, got {statement:?}");
+        };
+
+        let sql = "COPY (select 1) TO bar";
+        let expected = Statement::CopyTo(CopyToStatement {
+            source: CopyToSource::Query(query),
+            target: "bar".to_string(),
+            options: HashMap::new(),
+        });
+        expect_parse_ok(sql, expected)?;
+        Ok(())
+    }
+
+    #[test]
+    fn copy_to_options() -> Result<(), ParserError> {
+        let sql = "COPY foo TO bar (row_group_size 55)";
+        let expected = Statement::CopyTo(CopyToStatement {
+            source: object_name("foo"),
+            target: "bar".to_string(),
+            options: HashMap::from([(
+                "row_group_size".to_string(),
+                Value::Number("55".to_string(), false),
+            )]),
+        });
+        expect_parse_ok(sql, expected)?;
+        Ok(())
+    }
+
+    // For error cases, see: `copy.slt`
+
+    fn object_name(name: &str) -> CopyToSource {
+        CopyToSource::Relation(ObjectName(vec![Ident::new(name)]))
+    }
 }

datafusion/sql/src/relation/mod.rs

@@ -23,6 +23,7 @@ use sqlparser::ast::TableFactor;
 mod join;
 
 impl<'a, S: ContextProvider> SqlToRel<'a, S> {
+    /// Create a `LogicalPlan` that scans the named relation
     fn create_relation(
         &self,
         relation: TableFactor,