feat(python,rust,cli): add DATE function for SQL (#11541)

pola-rs · Oct 16, 2023 · 8589836 · 8589836
1 parent 7efc54e
commit 8589836
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 2 deletions.
diff --git a/crates/polars-sql/Cargo.toml b/crates/polars-sql/Cargo.toml
@@ -12,7 +12,7 @@ description = "SQL transpiler for Polars. Converts SQL to Polars logical plans"
 arrow = { workspace = true }
 polars-core = { workspace = true }
 polars-error = { workspace = true }
-polars-lazy = { workspace = true, features = ["strings", "cross_join", "trigonometry", "abs", "round_series", "log", "regex", "is_in", "meta", "cum_agg"] }
+polars-lazy = { workspace = true, features = ["strings", "cross_join", "trigonometry", "abs", "round_series", "log", "regex", "is_in", "meta", "cum_agg", "dtype-date"] }
 polars-plan = { workspace = true }
 
 rand = { workspace = true }

diff --git a/crates/polars-sql/src/functions.rs b/crates/polars-sql/src/functions.rs
@@ -2,8 +2,8 @@ use polars_core::prelude::{polars_bail, polars_err, PolarsResult};
 use polars_lazy::dsl::Expr;
 use polars_plan::dsl::{coalesce, count, when};
 use polars_plan::logical_plan::LiteralValue;
-use polars_plan::prelude::lit;
 use polars_plan::prelude::LiteralValue::Null;
+use polars_plan::prelude::{lit, StrptimeOptions};
 use sqlparser::ast::{
     Expr as SqlExpr, Function as SQLFunction, FunctionArg, FunctionArgExpr, Value as SqlValue,
     WindowSpec, WindowType,
@@ -217,6 +217,16 @@ pub(crate) enum PolarsSqlFunctions {
     /// ```
     Radians,
 
+    // ----
+    // Date Functions
+    // ----
+    /// SQL 'date' function
+    /// ```sql
+    /// SELECT DATE('2021-03-15') from df;
+    /// SELECT DATE('2021-03', '%Y-%m') from df;
+    /// ```
+    Date,
+
     // ----
     // String functions
     // ----
@@ -471,6 +481,7 @@ impl PolarsSqlFunctions {
             "cot",
             "cotd",
             "count",
+            "date",
             "degrees",
             "ends_with",
             "exp",
@@ -559,6 +570,11 @@ impl PolarsSqlFunctions {
             "nullif" => Self::NullIf,
             "coalesce" => Self::Coalesce,
 
+            // ----
+            // Date functions
+            // ----
+            "date" => Self::Date,
+
             // ----
             // String functions
             // ----
@@ -718,6 +734,14 @@ impl SqlFunctionVisitor<'_> {
                 }),
                 _ => polars_bail!(InvalidOperation:"Invalid number of arguments for RegexpLike: {}",function.args.len()),
             },
+            Date => match function.args.len() {
+                1 => self.visit_unary(|e| e.str().to_date(StrptimeOptions::default())),
+                2 => self.visit_binary(|e, fmt| e.str().to_date(fmt)),
+                _ => polars_bail!(InvalidOperation:
+                    "Invalid number of arguments for Date: {}",
+                    function.args.len()
+                ),
+            },
             RTrim => match function.args.len() {
                 1 => self.visit_unary(|e| e.str().strip_chars_end(lit(Null))),
                 2 => self.visit_binary(|e, s| e.str().strip_chars_end(s)),
@@ -1076,6 +1100,24 @@ impl FromSqlExpr for String {
     }
 }
 
+impl FromSqlExpr for StrptimeOptions {
+    fn from_sql_expr(expr: &SqlExpr, _: &mut SQLContext) -> PolarsResult<Self>
+    where
+        Self: Sized,
+    {
+        match expr {
+            SqlExpr::Value(v) => match v {
+                SqlValue::SingleQuotedString(s) => Ok(StrptimeOptions {
+                    format: Some(s.clone()),
+                    ..StrptimeOptions::default()
+                }),
+                _ => polars_bail!(ComputeError: "can't parse literal {:?}", v),
+            },
+            _ => polars_bail!(ComputeError: "can't parse literal {:?}", expr),
+        }
+    }
+}
+
 impl FromSqlExpr for Expr {
     fn from_sql_expr(expr: &SqlExpr, ctx: &mut SQLContext) -> PolarsResult<Self>
     where

diff --git a/py-polars/tests/unit/sql/test_sql.py b/py-polars/tests/unit/sql/test_sql.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import datetime
 import math
 from pathlib import Path
 
@@ -1208,3 +1209,27 @@ def test_sql_unary_ops_8890(match_float: bool) -> None:
             "c": [-3, -3],
             "d": [4, 4],
         }
+
+
+def test_sql_date() -> None:
+    df = pl.DataFrame(
+        {
+            "date": [
+                datetime.date(2021, 3, 15),
+                datetime.date(2021, 3, 28),
+                datetime.date(2021, 4, 4),
+            ],
+            "version": ["0.0.1", "0.7.3", "0.7.4"],
+        }
+    )
+
+    with pl.SQLContext(df=df, eager_execution=True) as ctx:
+        expected = pl.DataFrame({"date": [True, False, False]})
+        assert ctx.execute("SELECT date < DATE('2021-03-20') from df").frame_equal(
+            expected
+        )
+
+    expected = pl.DataFrame({"literal": ["2023-03-01"]})
+    assert pl.select(
+        pl.sql_expr("""CAST(DATE('2023-03', '%Y-%m') as STRING)""")
+    ).frame_equal(expected)