polars_sql/
table_functions.rs

1use std::str::FromStr;
2
3#[cfg(any(
4    feature = "csv",
5    feature = "parquet",
6    feature = "ipc",
7    feature = "json"
8))]
9use polars_core::prelude::polars_ensure;
10use polars_core::prelude::{PolarsError, PolarsResult, polars_bail};
11#[cfg(feature = "csv")]
12use polars_lazy::prelude::LazyCsvReader;
13use polars_lazy::prelude::LazyFrame;
14use polars_utils::plpath::PlPath;
15use sqlparser::ast::{
16    Expr as SQLExpr, FunctionArg as SQLFunctionArg, FunctionArgExpr as SQLFunctionArgExpr,
17    Value as SQLValue, ValueWithSpan as SQLValueWithSpan,
18};
19
20/// Table functions that are supported by Polars
21#[allow(clippy::enum_variant_names)]
22pub(crate) enum PolarsTableFunctions {
23    /// SQL 'read_csv' function.
24    /// ```sql
25    /// SELECT * FROM read_csv('path/to/file.csv')
26    /// ```
27    #[cfg(feature = "csv")]
28    ReadCsv,
29    /// SQL 'read_parquet' function.
30    /// ```sql
31    /// SELECT * FROM read_parquet('path/to/file.parquet')
32    /// ```
33    #[cfg(feature = "parquet")]
34    ReadParquet,
35    /// SQL 'read_ipc' function.
36    /// ```sql
37    /// SELECT * FROM read_ipc('path/to/file.ipc')
38    /// ```
39    #[cfg(feature = "ipc")]
40    ReadIpc,
41    /// SQL 'read_json' function (*only ndjson is currently supported*).
42    /// ```sql
43    /// SELECT * FROM read_json('path/to/file.json')
44    /// ```
45    #[cfg(feature = "json")]
46    ReadJson,
47}
48
49impl FromStr for PolarsTableFunctions {
50    type Err = PolarsError;
51
52    #[allow(unreachable_code)]
53    fn from_str(s: &str) -> Result<Self, Self::Err> {
54        Ok(match s.to_lowercase().as_str() {
55            #[cfg(feature = "csv")]
56            "read_csv" => PolarsTableFunctions::ReadCsv,
57            #[cfg(feature = "parquet")]
58            "read_parquet" => PolarsTableFunctions::ReadParquet,
59            #[cfg(feature = "ipc")]
60            "read_ipc" => PolarsTableFunctions::ReadIpc,
61            #[cfg(feature = "json")]
62            "read_json" => PolarsTableFunctions::ReadJson,
63            _ => polars_bail!(SQLInterface: "'{}' is not a supported table function", s),
64        })
65    }
66}
67
68impl PolarsTableFunctions {
69    #[allow(unused_variables, unreachable_patterns)]
70    pub(crate) fn execute(&self, args: &[SQLFunctionArg]) -> PolarsResult<(PlPath, LazyFrame)> {
71        match self {
72            #[cfg(feature = "csv")]
73            PolarsTableFunctions::ReadCsv => self.read_csv(args),
74            #[cfg(feature = "parquet")]
75            PolarsTableFunctions::ReadParquet => self.read_parquet(args),
76            #[cfg(feature = "ipc")]
77            PolarsTableFunctions::ReadIpc => self.read_ipc(args),
78            #[cfg(feature = "json")]
79            PolarsTableFunctions::ReadJson => self.read_ndjson(args),
80            _ => unreachable!(),
81        }
82    }
83
84    #[cfg(feature = "csv")]
85    fn read_csv(&self, args: &[SQLFunctionArg]) -> PolarsResult<(PlPath, LazyFrame)> {
86        polars_ensure!(args.len() == 1, SQLSyntax: "`read_csv` expects a single file path; found {:?} arguments", args.len());
87
88        use polars_lazy::frame::LazyFileListReader;
89        let path = self.get_file_path_from_arg(&args[0])?;
90        let lf = LazyCsvReader::new(path.clone())
91            .with_try_parse_dates(true)
92            .with_missing_is_null(true)
93            .finish()?;
94        Ok((path, lf))
95    }
96
97    #[cfg(feature = "parquet")]
98    fn read_parquet(&self, args: &[SQLFunctionArg]) -> PolarsResult<(PlPath, LazyFrame)> {
99        polars_ensure!(args.len() == 1, SQLSyntax: "`read_parquet` expects a single file path; found {:?} arguments", args.len());
100
101        let path = self.get_file_path_from_arg(&args[0])?;
102        let lf = LazyFrame::scan_parquet(path.clone(), Default::default())?;
103        Ok((path, lf))
104    }
105
106    #[cfg(feature = "ipc")]
107    fn read_ipc(&self, args: &[SQLFunctionArg]) -> PolarsResult<(PlPath, LazyFrame)> {
108        polars_ensure!(args.len() == 1, SQLSyntax: "`read_ipc` expects a single file path; found {:?} arguments", args.len());
109
110        let path = self.get_file_path_from_arg(&args[0])?;
111        let lf = LazyFrame::scan_ipc(path.clone(), Default::default(), Default::default())?;
112        Ok((path, lf))
113    }
114    #[cfg(feature = "json")]
115    fn read_ndjson(&self, args: &[SQLFunctionArg]) -> PolarsResult<(PlPath, LazyFrame)> {
116        polars_ensure!(args.len() == 1, SQLSyntax: "`read_ndjson` expects a single file path; found {:?} arguments", args.len());
117
118        use polars_lazy::frame::LazyFileListReader;
119        use polars_lazy::prelude::LazyJsonLineReader;
120
121        let path = self.get_file_path_from_arg(&args[0])?;
122        let lf = LazyJsonLineReader::new(path.clone()).finish()?;
123        Ok((path, lf))
124    }
125
126    #[allow(dead_code)]
127    fn get_file_path_from_arg(&self, arg: &SQLFunctionArg) -> PolarsResult<PlPath> {
128        match arg {
129            SQLFunctionArg::Unnamed(SQLFunctionArgExpr::Expr(SQLExpr::Value(
130                SQLValueWithSpan {
131                    value: SQLValue::SingleQuotedString(s),
132                    ..
133                },
134            ))) => Ok(PlPath::from_str(s)),
135            _ => polars_bail!(
136                SQLSyntax:
137                "expected a valid file path as a single-quoted string; found: {}", arg,
138            ),
139        }
140    }
141}
142
143impl PolarsTableFunctions {
144    // list sql names of all table functions
145    pub(crate) fn keywords() -> &'static [&'static str] {
146        &[
147            #[cfg(feature = "csv")]
148            "read_csv",
149            #[cfg(feature = "parquet")]
150            "read_parquet",
151            #[cfg(feature = "ipc")]
152            "read_ipc",
153            #[cfg(feature = "json")]
154            "read_json",
155        ]
156    }
157}
polars_sql/table_functions.rs

polars_sql/
table_functions.rs
