Add cast_map to scan_db for server-side column narrowing

scan_db gains an optional cast_map={column: dtype} that casts the named
columns server-side with a SQL CAST, wrapping the query in an
auto-enumerated projecting subquery so `select *` still flows every
column. The narrowed dtype is reported in the schema, so a filter on a
cast column (for example a datetime narrowed to date) is pushed down to
the database instead of stalling above a client-side cast.

Clauses that are illegal inside a derived table (a top-level ORDER BY
without TOP/OFFSET, or OPTION hints in the T-SQL dialect) are hoisted
onto the cast SELECT so the nested subquery stays valid. Unsupported
target dtypes are rejected rather than silently emitting VARCHAR, and
decimal precision and scale are preserved.

Extract a shared polars_dtype_to_sqlglot_type helper used by both the
predicate translator and the cast wrapping.

Signed-off-by: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com>

Author: ptomecek

docs/wiki/API-Reference.md

@@ -156,12 +156,15 @@ without materializing the whole frame.
 ### `scan_db`
 
 ```python
-scan_db(query, connection, fetch_size=10000, **kwargs) -> pl.LazyFrame
+scan_db(query, connection, fetch_size=10000, cast_map=None, **kwargs) -> pl.LazyFrame
 ```
 
 Run a SQL query over an ODBC `connection` string with predicate and projection pushdown.
 The SQL dialect is detected from the connection; filters become `WHERE` clauses and
-selections narrow the `SELECT`.
+selections narrow the `SELECT`. Pass `cast_map={column: dtype}` to cast columns
+server-side (a SQL `CAST`, keeping `select *`) and report the narrowed dtype in the
+schema, so a filter on a cast column still pushes down — for example narrowing a
+`datetime` column that is logically a `date`.
 
 ### `scan_clickhouse`
 

docs/wiki/Reading-and-Writing-Data.md

@@ -41,6 +41,25 @@ The dialect is detected from the ODBC connection, so the generated SQL matches y
 database. Pass `fetch_size=` to control the batch size used when Polars does not
 request one.
 
+If a column is stored as a wider type than it is used as — for example a `datetime`
+column that is logically a `date` — cast it server-side with `cast_map` so filters on
+it still push down:
+
+```python
+lf = scan_db(
+    "SELECT * FROM trades",
+    connection="Driver={PostgreSQL};Server=db.example.com;Database=mkt;Uid=reader;******",
+    cast_map={"ts": pl.Date},
+)
+
+# `ts` is narrowed to a date in the query, so this filter becomes a SQL WHERE clause
+# instead of being applied after a full scan.
+result = lf.filter(pl.col("ts") == date(2024, 1, 1)).collect()
+```
+
+`cast_map` wraps your query in a projecting subquery that casts the named columns and
+passes the rest through, so `select *` keeps flowing every column.
+
 ## Read from ClickHouse
 
 `scan_clickhouse` streams query results over ClickHouse's HTTP interface as Arrow IPC.

polars_io_tools/io_sources/lazy_sql_reader.py

@@ -10,6 +10,7 @@
 from .sql_utils import (
     apply_polars_io_source_exprs,
     fix_three_part_identifiers,
+    wrap_query_with_casts,
 )
 from .util import register_io_source_with_is_pure
 
@@ -90,7 +91,7 @@ def get_schema_from_query_odbc(
         raise ValueError(f"Could not determine schema for query: {query}, with error: {e}") from e
 
 
-def scan_db(query: str, connection: str, fetch_size: int = 10000, **kwargs) -> pl.LazyFrame:
+def scan_db(query: str, connection: str, fetch_size: int = 10000, cast_map: dict[str, Any] | None = None, **kwargs) -> pl.LazyFrame:
     """
     Create a LazyFrame from a SQL query with predicate pushdown support.
 
@@ -107,6 +108,12 @@ def scan_db(query: str, connection: str, fetch_size: int = 10000, **kwargs) -> p
             source generator function that scan_db wraps (because it is required \
             by the Polars IO plugins API). This value will only be used if Polars \
             does not pass a value for batch size; if it does, that will be used instead.
+        cast_map (dict[str, pl.DataType] | None, default None): Optional mapping of output column \
+            name to a Polars dtype to cast that column to *server-side*. The narrowing is emitted \
+            as a SQL ``CAST`` inside the query, and the reported schema reflects the target dtype, \
+            so filters on the cast column push down to the database. Use this to correct a \
+            mis-declared source type (e.g. a column stored as ``datetime`` that should be ``date``) \
+            without abandoning ``select *`` — remaining columns pass through untouched.
         **kwargs: Additional arguments for the database connector
 
     Returns:
@@ -132,6 +139,15 @@ def _fetch_info_needing_connection() -> tuple[
 
     schema, parsed_query, dialect = _fetch_info_needing_connection()
 
+    if cast_map:
+        unknown = [name for name in cast_map if name not in schema]
+        if unknown:
+            raise ValueError(f"cast_map references column(s) not in the query schema {list(schema)}: {unknown}")
+        # Narrow the columns server-side and report the target dtypes, so predicates on
+        # a cast column push down to the database instead of stalling above a client cast.
+        parsed_query = wrap_query_with_casts(parsed_query, dialect, list(schema), cast_map)
+        schema = {name: cast_map.get(name, dtype) for name, dtype in schema.items()}
+
     # Create the generator function for our custom IO source
     def source_generator(
         with_columns: list[str] | None,

polars_io_tools/io_sources/sql_utils.py

@@ -16,13 +16,66 @@
     "convert_predicate_to_sql",
     "create_sqlglot_literal",
     "fix_three_part_identifiers",
+    "polars_dtype_to_sqlglot_type",
 )
 
 
 # Configure logging
 log = logging.getLogger(__name__)
 
 
+def polars_dtype_to_sqlglot_type(dtype: pl.DataType | type[pl.DataType], *, strict: bool = False) -> sqlglot.exp.DataType:
+    """Map a Polars dtype (instance or class) to a sqlglot ``DataType``.
+
+    Shared by the predicate translator (``SQLExpressionVisitor.visit_cast``) and the
+    server-side cast wrapping in ``scan_db`` so both emit identical SQL types.
+
+    Decimal precision and scale are preserved so a decimal cast does not silently
+    truncate the fractional part.
+
+    Args:
+        dtype: The Polars dtype to translate, as a class or an instance.
+        strict: When True, raise ``ValueError`` for a dtype without a dedicated SQL
+            mapping instead of falling back to ``VARCHAR``. Callers that report the
+            requested dtype as the resulting schema (so the emitted SQL must actually
+            produce it) should set this; the predicate translator leaves it False and
+            lets the database coerce an unmapped cast.
+    """
+    type_map = {
+        pl.Int8: sqlglot.exp.DataType.Type.TINYINT,
+        pl.Int16: sqlglot.exp.DataType.Type.SMALLINT,
+        pl.Int32: sqlglot.exp.DataType.Type.INT,
+        pl.Int64: sqlglot.exp.DataType.Type.BIGINT,
+        pl.UInt8: getattr(sqlglot.exp.DataType.Type, "UTINYINT", sqlglot.exp.DataType.Type.TINYINT),
+        pl.UInt16: getattr(sqlglot.exp.DataType.Type, "USMALLINT", sqlglot.exp.DataType.Type.SMALLINT),
+        pl.UInt32: getattr(sqlglot.exp.DataType.Type, "UINT", sqlglot.exp.DataType.Type.INT),
+        pl.UInt64: getattr(sqlglot.exp.DataType.Type, "UBIGINT", sqlglot.exp.DataType.Type.BIGINT),
+        pl.Float32: sqlglot.exp.DataType.Type.FLOAT,
+        pl.Float64: sqlglot.exp.DataType.Type.DOUBLE,
+        pl.Utf8: sqlglot.exp.DataType.Type.VARCHAR,
+        pl.Date: sqlglot.exp.DataType.Type.DATE,
+        pl.Datetime: sqlglot.exp.DataType.Type.TIMESTAMP,
+        pl.Time: sqlglot.exp.DataType.Type.TIME,
+        pl.Decimal: sqlglot.exp.DataType.Type.DECIMAL,
+    }
+    key = dtype if isinstance(dtype, type) else type(dtype)
+    base = type_map.get(key)
+    if base is None:
+        if strict:
+            supported = ", ".join(sorted(t.__name__ for t in type_map))
+            raise ValueError(f"No SQL type mapping for Polars dtype {dtype!r}. Supported dtypes: {supported}.")
+        return sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.VARCHAR)
+
+    # Preserve decimal precision/scale so the cast keeps the fractional part.
+    if key is pl.Decimal and not isinstance(dtype, type):
+        precision = getattr(dtype, "precision", None)
+        if precision is not None:
+            scale = getattr(dtype, "scale", None) or 0
+            return sqlglot.exp.DataType.build(f"DECIMAL({precision}, {scale})")
+
+    return sqlglot.exp.DataType(this=base)
+
+
 def create_sqlglot_literal(value: Any) -> sqlglot.exp.Expression:
     """Create a sqlglot literal from a raw value.
 
@@ -405,35 +458,8 @@ def visit_cast(self, node: CastNode) -> None:
             self.result = input_expr
             return
 
-        # Map Polars types to SQL types
-        type_map = {
-            pl.Int8: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.TINYINT),
-            pl.Int16: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.SMALLINT),
-            pl.Int32: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.INT),
-            pl.Int64: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.BIGINT),
-            pl.UInt8: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.UTINYINT)
-            if hasattr(sqlglot.exp.DataType.Type, "UTINYINT")
-            else sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.TINYINT),
-            pl.UInt16: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.USMALLINT)
-            if hasattr(sqlglot.exp.DataType.Type, "USMALLINT")
-            else sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.SMALLINT),
-            pl.UInt32: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.UINT)
-            if hasattr(sqlglot.exp.DataType.Type, "UINT")
-            else sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.INT),
-            pl.UInt64: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.UBIGINT)
-            if hasattr(sqlglot.exp.DataType.Type, "UBIGINT")
-            else sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.BIGINT),
-            pl.Float32: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.FLOAT),
-            pl.Float64: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.DOUBLE),
-            pl.Utf8: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.VARCHAR),
-            pl.Date: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.DATE),
-            pl.Datetime: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.TIMESTAMP),
-            pl.Time: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.TIME),
-            pl.Decimal: sqlglot.exp.DataType(this=sqlglot.exp.DataType.Type.DECIMAL),
-        }
-
-        # Get the SQL type
-        sql_type = type_map.get(type(node.dtype), "VARCHAR")
+        # Map the Polars type to a SQL type (shared with scan_db's cast_map)
+        sql_type = polars_dtype_to_sqlglot_type(node.dtype)
 
         # Create the CAST expression
         self.result = sqlglot.exp.Cast(this=input_expr, to=sql_type)
@@ -678,6 +704,60 @@ def _prepare_inner_for_subquery(
     return inner, order_to_hoist, options_to_hoist
 
 
+def wrap_query_with_casts(
+    query: sqlglot.exp.Expression,
+    dialect: str | type[Dialect] | None,
+    ordered_columns: list[str],
+    cast_map: dict[str, pl.DataType | type[pl.DataType]],
+) -> sqlglot.exp.Expression:
+    """Wrap ``query`` in a projecting subquery that casts selected columns server-side.
+
+    Every column in ``ordered_columns`` is re-projected (preserving order); columns
+    present in ``cast_map`` are wrapped in a SQL ``CAST`` to the mapped type, the rest
+    pass through untouched. The original query is left intact as the inner relation,
+    so ``select *`` keeps flowing new upstream columns.
+
+    The narrowing lands one subquery level *below* the predicate/projection pushdown
+    added later by :func:`apply_polars_io_source_exprs`, so a filter on a cast column
+    (e.g. ``datetime`` narrowed to ``date``) resolves against the already-cast output
+    column and is pushed to the database rather than evaluated client-side.
+
+    Clauses that are illegal inside a derived table (an MSSQL top-level ``ORDER BY``
+    without ``TOP``/``OFFSET``, or ``OPTION`` hints) are hoisted onto the cast
+    ``SELECT`` so the nested subquery stays valid; the later wrapping in
+    :func:`apply_polars_io_source_exprs` hoists them again to statement level.
+    """
+    quote = _is_known_dialect(dialect)
+
+    def _ident(name: str) -> sqlglot.exp.Identifier:
+        return sqlglot.exp.Identifier(this=name, quoted=quote)
+
+    inner_query, order_to_hoist, options_to_hoist = _prepare_inner_for_subquery(query, dialect=dialect)
+
+    projections: list[sqlglot.exp.Expression] = []
+    for name in ordered_columns:
+        col = sqlglot.exp.Column(this=_ident(name))
+        if name in cast_map:
+            sql_type = polars_dtype_to_sqlglot_type(cast_map[name], strict=True)
+            projections.append(sqlglot.exp.Cast(this=col, to=sql_type).as_(name, quoted=quote))
+        else:
+            projections.append(col)
+
+    inner = sqlglot.exp.Subquery(
+        this=inner_query,
+        alias=sqlglot.exp.TableAlias(this=sqlglot.exp.Identifier(this="__cpl_cast")),
+    )
+    cast_select = sqlglot.exp.Select().from_(inner, dialect=dialect).select(*projections, append=False, dialect=dialect)
+
+    # Re-apply the clauses hoisted out of the (now nested) original query.
+    if order_to_hoist is not None:
+        cast_select.set("order", order_to_hoist)
+    for opt in options_to_hoist:
+        cast_select.args.setdefault("options", []).append(opt)
+
+    return cast_select
+
+
 def apply_polars_io_source_exprs(
     query: sqlglot.exp.Expression,
     dialect: str | type[Dialect] | None,

polars_io_tools/tests/io_sources/test_lazy_sql_reader.py

@@ -2821,3 +2821,141 @@ def test_mssql_parenthesized_set_operations_only_stabilize_leftmost_output_branc
         "[g]",
         "COUNT(*) AS except_count",
     ]
+
+
+# Tests for scan_db's server-side cast_map (datetime -> date narrowing pushdown)
+
+
+def _make_record_table(conn):
+    """Create a table whose RecordDate is a TIMESTAMP (the mis-declared-type case)."""
+    conn.execute(
+        """
+        CREATE TABLE RecordTbl AS SELECT * FROM (
+            SELECT CAST('2025-05-30 09:15:00' AS TIMESTAMP) AS RecordDate, 1 AS PointID, 1.5 AS Rate
+            UNION ALL
+            SELECT CAST('2025-05-30 17:45:00' AS TIMESTAMP) AS RecordDate, 2 AS PointID, 2.5 AS Rate
+            UNION ALL
+            SELECT CAST('2025-05-31 00:00:00' AS TIMESTAMP) AS RecordDate, 3 AS PointID, 3.5 AS Rate
+        )
+        """
+    )
+
+
+def test_cast_map_reports_target_dtype_in_schema(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    lf = cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string", cast_map={"RecordDate": pl.Date})
+    schema = lf.collect_schema()
+    assert schema["RecordDate"] == pl.Date
+    # Untouched columns keep their native dtypes and are still present (select * preserved).
+    assert set(schema.names()) == {"RecordDate", "PointID", "Rate"}
+
+
+def test_cast_map_narrows_column_server_side(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    out = cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string", cast_map={"RecordDate": pl.Date}).sort("PointID").collect()
+    assert out["RecordDate"].dtype == pl.Date
+    assert out["RecordDate"].to_list() == [date(2025, 5, 30), date(2025, 5, 30), date(2025, 5, 31)]
+
+
+def test_cast_map_pushes_date_predicate_to_sql(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    lf = cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string", cast_map={"RecordDate": pl.Date})
+
+    import arrow_odbc
+
+    captured: list[str] = []
+    original = arrow_odbc.read_arrow_batches_from_odbc
+
+    def capturing(*args, **kwargs):
+        captured.append(args[0] if args else kwargs["query"])
+        return original(*args, **kwargs)
+
+    arrow_odbc.read_arrow_batches_from_odbc = capturing
+    try:
+        out = lf.filter(pl.col("RecordDate") == date(2025, 5, 30)).sort("PointID").collect()
+    finally:
+        arrow_odbc.read_arrow_batches_from_odbc = original
+
+    # Two rows on 2025-05-30 survive the narrowing; the 05-31 row is filtered out.
+    assert out["PointID"].to_list() == [1, 2]
+    assert out["RecordDate"].to_list() == [date(2025, 5, 30), date(2025, 5, 30)]
+
+    # The predicate reached SQL: the executed query casts RecordDate to DATE and filters on it.
+    assert captured, "No SQL captured"
+    sql = captured[-1]
+    assert "CAST" in sql.upper() and "DATE" in sql.upper()
+    assert "RecordDate" in sql
+    assert "WHERE" in sql.upper(), f"date predicate was not pushed to SQL: {sql}"
+
+
+def test_cast_map_projection_pushdown_still_prunes(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    lf = cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string", cast_map={"RecordDate": pl.Date})
+
+    import arrow_odbc
+
+    captured: list[str] = []
+    original = arrow_odbc.read_arrow_batches_from_odbc
+
+    def capturing(*args, **kwargs):
+        captured.append(args[0] if args else kwargs["query"])
+        return original(*args, **kwargs)
+
+    arrow_odbc.read_arrow_batches_from_odbc = capturing
+    try:
+        out = lf.select(["RecordDate", "PointID"]).sort("PointID").collect()
+    finally:
+        arrow_odbc.read_arrow_batches_from_odbc = original
+
+    assert out.columns == ["RecordDate", "PointID"]
+    assert out["RecordDate"].dtype == pl.Date
+    # The outer projection is pruned to the selected columns (the inner cast layer
+    # still enumerates every column so the DB prunes Rate via the outer select).
+    sql = captured[-1]
+    assert sql.strip().upper().startswith('SELECT "RECORDDATE", "POINTID" FROM')
+
+
+def test_cast_map_unknown_column_raises(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    with pytest.raises(ValueError, match="cast_map references column"):
+        cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string", cast_map={"NotAColumn": pl.Date})
+
+
+def test_cast_map_none_is_unchanged(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    lf = cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string")
+    assert lf.collect_schema()["RecordDate"] == pl.Datetime("us")
+
+
+def test_cast_map_hoists_mssql_order_by_out_of_derived_table():
+    """A top-level MSSQL ORDER BY must not end up inside the cast derived table.
+
+    MSSQL rejects ORDER BY in a derived table without TOP/OFFSET (error 1033), so the
+    cast wrapper must hoist it to statement level.
+    """
+    from polars_io_tools.io_sources.sql_utils import apply_polars_io_source_exprs, wrap_query_with_casts
+
+    query = parse_one("SELECT * FROM dbo.t ORDER BY x", dialect=MSSQL)
+    wrapped = wrap_query_with_casts(query, MSSQL, ["x", "y"], {"x": pl.Date})
+    final = apply_polars_io_source_exprs(wrapped.copy(), MSSQL, None, pl.col("x") == date(2025, 5, 30), None, None)
+    sql = final.sql(dialect=MSSQL)
+
+    # ORDER BY sits at the end, after the outermost subquery closes, not inside __cpl_cast.
+    assert "AS __cpl_cast) AS __cpl_subq" in sql
+    assert "ORDER BY [x]" in sql
+    assert sql.rstrip().endswith("ORDER BY [x]")
+    assert "ORDER BY [x]) AS __cpl_cast" not in sql
+
+
+def test_cast_map_preserves_decimal_scale():
+    from polars_io_tools.io_sources.sql_utils import polars_dtype_to_sqlglot_type
+
+    assert polars_dtype_to_sqlglot_type(pl.Decimal(10, 2)).sql(dialect=MSSQL) == "NUMERIC(10, 2)"
+
+
+def test_cast_map_rejects_unsupported_dtype(duckdb_connection):
+    _make_record_table(duckdb_connection)
+    # Boolean has no server-side narrowing mapping; reporting it as the schema while
+    # emitting VARCHAR would misrepresent the data, so it must be rejected.
+    with pytest.raises(ValueError, match="No SQL type mapping"):
+        cpl.scan_db("SELECT * FROM RecordTbl", "fake_connection_string", cast_map={"RecordDate": pl.Boolean})