Add support for Hive's LOAD DATA expr

Author: wugeer

src/ast/mod.rs

@@ -3317,6 +3317,22 @@ pub enum Statement {
         channel: Ident,
         payload: Option<String>,
     },
+    /// ```sql
+    /// LOAD DATA [LOCAL] INPATH 'filepath' [OVERWRITE] INTO TABLE tablename
+    /// [PARTITION (partcol1=val1, partcol2=val2 ...)]
+    /// [INPUTFORMAT 'inputformat' SERDE 'serde']
+    /// ```
+    /// Loading files into tables
+    ///
+    /// See Hive <https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=27362036#LanguageManualDML-Loadingfilesintotables>
+    LoadData {
+        local: bool,
+        inpath: String,
+        overwrite: bool,
+        table_name: ObjectName,
+        partitioned: Option<Vec<Expr>>,
+        table_format: Option<HiveLoadDataOption>,
+    },
 }
 
 impl fmt::Display for Statement {
@@ -3919,6 +3935,36 @@ impl fmt::Display for Statement {
                 Ok(())
             }
             Statement::CreateTable(create_table) => create_table.fmt(f),
+            Statement::LoadData {
+                local,
+                inpath,
+                overwrite,
+                table_name,
+                partitioned,
+                table_format,
+            } => {
+                write!(
+                    f,
+                    "LOAD DATA {local}INPATH '{inpath}' {overwrite}INTO TABLE {table_name}",
+                    local = if *local { "LOCAL " } else { "" },
+                    inpath = inpath,
+                    overwrite = if *overwrite { "OVERWRITE " } else { "" },
+                    table_name = table_name,
+                )?;
+                if let Some(ref parts) = &partitioned {
+                    if !parts.is_empty() {
+                        write!(f, " PARTITION ({})", display_comma_separated(parts))?;
+                    }
+                }
+                if let Some(HiveLoadDataOption {
+                    serde,
+                    input_format,
+                }) = &table_format
+                {
+                    write!(f, " INPUTFORMAT {input_format} SERDE {serde}")?;
+                }
+                Ok(())
+            }
             Statement::CreateVirtualTable {
                 name,
                 if_not_exists,
@@ -5811,6 +5857,14 @@ pub enum HiveRowFormat {
     DELIMITED { delimiters: Vec<HiveRowDelimiter> },
 }
 
+#[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+#[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]
+pub struct HiveLoadDataOption {
+    pub serde: Expr,
+    pub input_format: Expr,
+}
+
 #[derive(Debug, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 #[cfg_attr(feature = "visitor", derive(Visit, VisitMut))]

src/dialect/duckdb.rs

@@ -66,4 +66,9 @@ impl Dialect for DuckDbDialect {
     fn supports_explain_with_utility_options(&self) -> bool {
         true
     }
+
+    /// See DuckDB <https://duckdb.org/docs/sql/statements/load_and_install.html#load>
+    fn supports_load_extension(&self) -> bool {
+        true
+    }
 }

src/dialect/hive.rs

@@ -51,4 +51,9 @@ impl Dialect for HiveDialect {
     fn require_interval_qualifier(&self) -> bool {
         true
     }
+
+    /// See Hive <https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=27362036#LanguageManualDML-Loadingfilesintotables>
+    fn supports_load_data(&self) -> bool {
+        true
+    }
 }

src/dialect/mod.rs

@@ -601,6 +601,16 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if the dialect supports the `LOAD DATA` statement
+    fn supports_load_data(&self) -> bool {
+        false
+    }
+
+    /// Returns true if the dialect supports the `LOAD extension` statement
+    fn supports_load_extension(&self) -> bool {
+        false
+    }
+
     /// Returns true if this dialect expects the the `TOP` option
     /// before the `ALL`/`DISTINCT` options in a `SELECT` statement.
     fn supports_top_before_distinct(&self) -> bool {

src/keywords.rs

@@ -387,6 +387,7 @@ define_keywords!(
     INITIALLY,
     INNER,
     INOUT,
+    INPATH,
     INPUT,
     INPUTFORMAT,
     INSENSITIVE,

src/parser/mod.rs

@@ -27,16 +27,15 @@ use core::{
 
 use log::debug;
 
-use recursion::RecursionCounter;
-use IsLateral::*;
-use IsOptional::*;
-
 use crate::ast::helpers::stmt_create_table::{CreateTableBuilder, CreateTableConfiguration};
 use crate::ast::Statement::CreatePolicy;
 use crate::ast::*;
 use crate::dialect::*;
 use crate::keywords::{Keyword, ALL_KEYWORDS};
 use crate::tokenizer::*;
+use recursion::RecursionCounter;
+use IsLateral::*;
+use IsOptional::*;
 
 mod alter;
 
@@ -543,10 +542,7 @@ impl<'a> Parser<'a> {
                 Keyword::INSTALL if dialect_of!(self is DuckDbDialect | GenericDialect) => {
                     self.parse_install()
                 }
-                // `LOAD` is duckdb specific https://duckdb.org/docs/extensions/overview
-                Keyword::LOAD if dialect_of!(self is DuckDbDialect | GenericDialect) => {
-                    self.parse_load()
-                }
+                Keyword::LOAD => self.parse_load(),
                 // `OPTIMIZE` is clickhouse specific https://clickhouse.tech/docs/en/sql-reference/statements/optimize/
                 Keyword::OPTIMIZE if dialect_of!(self is ClickHouseDialect | GenericDialect) => {
                     self.parse_optimize_table()
@@ -11040,6 +11036,22 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn parse_load_data_table_format(
+        &mut self,
+    ) -> Result<Option<HiveLoadDataOption>, ParserError> {
+        if self.parse_keyword(Keyword::INPUTFORMAT) {
+            let input_format = self.parse_expr()?;
+            self.expect_keyword(Keyword::SERDE)?;
+            let serde = self.parse_expr()?;
+            Ok(Some(HiveLoadDataOption {
+                input_format,
+                serde,
+            }))
+        } else {
+            Ok(None)
+        }
+    }
+
     /// Parse an UPDATE statement, returning a `Box`ed SetExpr
     ///
     /// This is used to reduce the size of the stack frames in debug builds
@@ -12042,10 +12054,35 @@ impl<'a> Parser<'a> {
         Ok(Statement::Install { extension_name })
     }
 
-    /// `LOAD [extension_name]`
+    /// Parse a SQL LOAD statement
     pub fn parse_load(&mut self) -> Result<Statement, ParserError> {
-        let extension_name = self.parse_identifier(false)?;
-        Ok(Statement::Load { extension_name })
+        if self.dialect.supports_load_extension() {
+            let extension_name = self.parse_identifier(false)?;
+            Ok(Statement::Load { extension_name })
+        } else if self.parse_keyword(Keyword::DATA) && self.dialect.supports_load_data() {
+            let local = self.parse_one_of_keywords(&[Keyword::LOCAL]).is_some();
+            self.expect_keyword(Keyword::INPATH)?;
+            let inpath = self.parse_literal_string()?;
+            let overwrite = self.parse_one_of_keywords(&[Keyword::OVERWRITE]).is_some();
+            self.expect_keyword(Keyword::INTO)?;
+            self.expect_keyword(Keyword::TABLE)?;
+            let table_name = self.parse_object_name(false)?;
+            let partitioned = self.parse_insert_partition()?;
+            let table_format = self.parse_load_data_table_format()?;
+            Ok(Statement::LoadData {
+                local,
+                inpath,
+                overwrite,
+                table_name,
+                partitioned,
+                table_format,
+            })
+        } else {
+            self.expected(
+                "Expected: dialect supports `LOAD DATA` or `LOAD extension` to parse `LOAD` statements",
+                self.peek_token(),
+            )
+        }
     }
 
     /// ```sql

tests/sqlparser_common.rs

@@ -11523,6 +11523,160 @@ fn parse_notify_channel() {
     }
 }
 
+#[test]
+fn parse_load_data() {
+    let dialects = all_dialects_where(|d| d.supports_load_data());
+
+    match dialects
+        .verified_stmt("LOAD DATA INPATH '/local/path/to/data.txt' INTO TABLE test.my_table")
+    {
+        Statement::LoadData {
+            local,
+            inpath,
+            overwrite,
+            table_name,
+            partitioned,
+            table_format,
+        } => {
+            assert_eq!(false, local);
+            assert_eq!("/local/path/to/data.txt", inpath);
+            assert_eq!(false, overwrite);
+            assert_eq!(
+                ObjectName(vec![Ident::new("test"), Ident::new("my_table")]),
+                table_name
+            );
+            assert_eq!(None, partitioned);
+            assert_eq!(None, table_format);
+        }
+        _ => unreachable!(),
+    };
+
+    // with OVERWRITE keyword
+    match dialects
+        .verified_stmt("LOAD DATA INPATH '/local/path/to/data.txt' OVERWRITE INTO TABLE my_table")
+    {
+        Statement::LoadData {
+            local,
+            inpath,
+            overwrite,
+            table_name,
+            partitioned,
+            table_format,
+        } => {
+            assert_eq!(false, local);
+            assert_eq!("/local/path/to/data.txt", inpath);
+            assert_eq!(true, overwrite);
+            assert_eq!(ObjectName(vec![Ident::new("my_table")]), table_name);
+            assert_eq!(None, partitioned);
+            assert_eq!(None, table_format);
+        }
+        _ => unreachable!(),
+    };
+
+    // with LOCAL keyword
+    match dialects
+        .verified_stmt("LOAD DATA LOCAL INPATH '/local/path/to/data.txt' INTO TABLE test.my_table")
+    {
+        Statement::LoadData {
+            local,
+            inpath,
+            overwrite,
+            table_name,
+            partitioned,
+            table_format,
+        } => {
+            assert_eq!(true, local);
+            assert_eq!("/local/path/to/data.txt", inpath);
+            assert_eq!(false, overwrite);
+            assert_eq!(
+                ObjectName(vec![Ident::new("test"), Ident::new("my_table")]),
+                table_name
+            );
+            assert_eq!(None, partitioned);
+            assert_eq!(None, table_format);
+        }
+        _ => unreachable!(),
+    };
+
+    // with PARTITION  clause
+    match dialects.verified_stmt("LOAD DATA LOCAL INPATH '/local/path/to/data.txt' INTO TABLE my_table PARTITION (year = 2024, month = 11)") {
+        Statement::LoadData {local, inpath, overwrite, table_name, partitioned, table_format} => {
+            assert_eq!(true, local);
+            assert_eq!("/local/path/to/data.txt", inpath);
+            assert_eq!(false, overwrite);
+            assert_eq!(ObjectName(vec![Ident::new("my_table")]), table_name);
+            assert_eq!(Some(vec![
+                Expr::BinaryOp{
+                    left: Box::new(Expr::Identifier(Ident::new("year"))),
+                    op: BinaryOperator::Eq,
+                    right:  Box::new(Expr::Value(Value::Number("2024".parse().unwrap(), false))),
+                },
+                Expr::BinaryOp{
+                    left: Box::new(Expr::Identifier(Ident::new("month"))),
+                    op: BinaryOperator::Eq,
+                    right:  Box::new(Expr::Value(Value::Number("11".parse().unwrap(), false))),
+                }]), partitioned);
+            assert_eq!(None, table_format);
+        }
+        _ => unreachable!(),
+    };
+
+    // with PARTITION  clause
+    match dialects.verified_stmt("LOAD DATA LOCAL INPATH '/local/path/to/data.txt' OVERWRITE INTO TABLE good.my_table PARTITION (year = 2024, month = 11) INPUTFORMAT 'org.apache.hadoop.mapred.TextInputFormat' SERDE 'org.apache.hadoop.hive.serde2.OpenCSVSerde'") {
+            Statement::LoadData {local, inpath, overwrite, table_name, partitioned, table_format} => {
+                assert_eq!(true, local);
+                assert_eq!("/local/path/to/data.txt", inpath);
+                assert_eq!(true, overwrite);
+                assert_eq!(ObjectName(vec![Ident::new("good"), Ident::new("my_table")]), table_name);
+                assert_eq!(Some(vec![
+                    Expr::BinaryOp{
+                        left: Box::new(Expr::Identifier(Ident::new("year"))),
+                        op: BinaryOperator::Eq,
+                        right:  Box::new(Expr::Value(Value::Number("2024".parse().unwrap(), false))),
+                    },
+                    Expr::BinaryOp{
+                        left: Box::new(Expr::Identifier(Ident::new("month"))),
+                        op: BinaryOperator::Eq,
+                        right:  Box::new(Expr::Value(Value::Number("11".parse().unwrap(), false))),
+                    }]), partitioned);
+                assert_eq!(Some(HiveLoadDataOption {serde: Expr::Value(Value::SingleQuotedString("org.apache.hadoop.hive.serde2.OpenCSVSerde".to_string())), input_format: Expr::Value(Value::SingleQuotedString("org.apache.hadoop.mapred.TextInputFormat".to_string()))}), table_format);
+            }
+            _ => unreachable!(),
+        };
+
+    let dialects = all_dialects_where(|d| !d.supports_load_data() && d.supports_load_extension());
+
+    assert_eq!(
+        dialects
+            .parse_sql_statements(
+                "LOAD DATA LOCAL INPATH '/local/path/to/data.txt' INTO TABLE test.my_table"
+            )
+            .unwrap_err(),
+        ParserError::ParserError("Expected: end of statement, found: LOCAL".to_string())
+    );
+
+    assert_eq!(
+        dialects
+            .parse_sql_statements(
+                "LOAD DATA INPATH '/local/path/to/data.txt' INTO TABLE test.my_table"
+            )
+            .unwrap_err(),
+        ParserError::ParserError("Expected: end of statement, found: INPATH".to_string())
+    );
+
+    let dialects = all_dialects_where(|d| !d.supports_load_data() && !d.supports_load_extension());
+
+    assert_eq!(
+        dialects.parse_sql_statements("LOAD DATA LOCAL INPATH '/local/path/to/data.txt' INTO TABLE test.my_table").unwrap_err(),
+        ParserError::ParserError("Expected: Expected: dialect supports `LOAD DATA` or `LOAD extension` to parse `LOAD` statements, found: LOCAL".to_string())
+    );
+
+    assert_eq!(
+        dialects.parse_sql_statements("LOAD DATA INPATH '/local/path/to/data.txt' INTO TABLE test.my_table").unwrap_err(),
+        ParserError::ParserError("Expected: Expected: dialect supports `LOAD DATA` or `LOAD extension` to parse `LOAD` statements, found: INPATH".to_string())
+    );
+}
+
 #[test]
 fn test_select_top() {
     let dialects = all_dialects_where(|d| d.supports_top_before_distinct());