feat(rust,python,cli): add SQL support for timestamp precision modi…

…fier

crates/polars-sql/src/sql_expr.rs

@@ -14,8 +14,8 @@ use sqlparser::ast::ExactNumberInfo;
 use sqlparser::ast::{
     ArrayAgg, ArrayElemTypeDef, BinaryOperator as SQLBinaryOperator, BinaryOperator, CastFormat,
     DataType as SQLDataType, DateTimeField, Expr as SQLExpr, Function as SQLFunction, Ident,
-    JoinConstraint, OrderByExpr, Query as Subquery, SelectItem, TrimWhereField, UnaryOperator,
-    Value as SQLValue,
+    JoinConstraint, OrderByExpr, Query as Subquery, SelectItem, TimezoneInfo, TrimWhereField,
+    UnaryOperator, Value as SQLValue,
 };
 use sqlparser::dialect::GenericDialect;
 use sqlparser::parser::{Parser, ParserOptions};
@@ -62,11 +62,30 @@ pub(crate) fn map_sql_polars_datatype(data_type: &SQLDataType) -> PolarsResult<D
         SQLDataType::Int2(_) => DataType::Int16,
         SQLDataType::Int4(_) => DataType::Int32,
         SQLDataType::Int8(_) => DataType::Int64,
-        SQLDataType::Interval => DataType::Duration(TimeUnit::Milliseconds),
+        SQLDataType::Interval => DataType::Duration(TimeUnit::Microseconds),
         SQLDataType::Real => DataType::Float32,
         SQLDataType::SmallInt(_) => DataType::Int16,
-        SQLDataType::Time { .. } => DataType::Time,
-        SQLDataType::Timestamp { .. } => DataType::Datetime(TimeUnit::Milliseconds, None),
+        SQLDataType::Time(_, tz) => match tz {
+            TimezoneInfo::None => DataType::Time,
+            _ => polars_bail!(ComputeError: "Time with timezone is not supported; found tz={}", tz),
+        },
+        SQLDataType::Timestamp(prec, tz) => {
+            let tu = match prec {
+                None => TimeUnit::Microseconds,
+                Some(3) => TimeUnit::Milliseconds,
+                Some(6) => TimeUnit::Microseconds,
+                Some(9) => TimeUnit::Nanoseconds,
+                Some(n) => {
+                    polars_bail!(ComputeError: "Unsupported timestamp precision; expected 3, 6 or 9, found prec={}", n)
+                },
+            };
+            match tz {
+                TimezoneInfo::None => DataType::Datetime(tu, None),
+                _ => {
+                    polars_bail!(ComputeError: "Timestamp with timezone is not (yet) supported; found tz={}", tz)
+                },
+            }
+        },
         SQLDataType::TinyInt(_) => DataType::Int8,
         SQLDataType::UnsignedBigInt(_) => DataType::UInt64,
         SQLDataType::UnsignedInt(_) | SQLDataType::UnsignedInteger(_) => DataType::UInt32,

py-polars/tests/unit/sql/test_temporal.py

@@ -6,6 +6,7 @@
 import pytest
 
 import polars as pl
+from polars.exceptions import ComputeError
 from polars.testing import assert_frame_equal
 
 
@@ -121,3 +122,43 @@ def test_extract_century_millennium(dt: date, expected: list[int]) -> None:
                 schema=["c1", "c2", "c3", "c4"],
             ).cast(pl.Int32),
         )
+
+
+@pytest.mark.parametrize(
+    ("unit", "expected"),
+    [
+        ("ms", [1704589323123, 1609324245987, 1136159999555]),
+        ("us", [1704589323123456, 1609324245987654, 1136159999555555]),
+        ("ns", [1704589323123456000, 1609324245987654000, 1136159999555555000]),
+        (None, [1704589323123456, 1609324245987654, 1136159999555555]),
+    ],
+)
+def test_timestamp_time_unit(unit: str | None, expected: list[int]) -> None:
+    df = pl.DataFrame(
+        {
+            "ts": [
+                datetime(2024, 1, 7, 1, 2, 3, 123456),
+                datetime(2020, 12, 30, 10, 30, 45, 987654),
+                datetime(2006, 1, 1, 23, 59, 59, 555555),
+            ],
+        }
+    )
+    precision = {"ms": 3, "us": 6, "ns": 9}
+
+    with pl.SQLContext(frame_data=df, eager_execution=True) as ctx:
+        prec = f"({precision[unit]})" if unit else ""
+        res = ctx.execute(f"SELECT ts::timestamp{prec} FROM frame_data").to_series()
+
+        assert res.dtype == pl.Datetime(time_unit=unit)  # type: ignore[arg-type]
+        assert res.to_physical().to_list() == expected
+
+
+def test_timestamp_time_unit_errors() -> None:
+    df = pl.DataFrame({"ts": [datetime(2024, 1, 7, 1, 2, 3, 123456)]})
+
+    with pl.SQLContext(frame_data=df, eager_execution=True) as ctx:
+        for prec in (0, 4, 15):
+            with pytest.raises(
+                ComputeError, match=f"Unsupported timestamp precision; .* prec={prec}"
+            ):
+                ctx.execute(f"SELECT ts::timestamp({prec}) FROM frame_data")