polars_sql/
table_functions.rs

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

polars_sql/
table_functions.rs
