Add cast_map to scan_db for server-side column narrowing

docs/wiki/API-Reference.md

@@ -176,12 +176,17 @@ 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`, or a `float` id that should be an integer.
+A predicate on a cast column pushes down over its `CAST(...)`, which can affect index use;
+the impact on the query plan is backend dependent (see Reading and Writing Data).
 
 ### `scan_clickhouse`
 

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

@@ -41,6 +41,32 @@ 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 different type than it is used as — for example a `datetime`
+column that is logically a `date`, or a numeric id delivered as `float` that should be an
+integer — 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") == pl.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.
+
+A predicate on a cast column is pushed down as `CAST(col AS ...) <op> value`. Wrapping the
+column in a function can stop the database from using an index on it, so the effect on the
+query plan is **backend dependent** — some engines optimize particular conversions (for
+example SQL Server can still seek on `CAST(datetime AS date)`) while others fall back to a
+full scan. When a filtered column is indexed and on a hot path, prefer filtering the
+physical column directly over its cast form.
+
 ## 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,19 @@ 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``, \
+            or a numeric id delivered as ``float`` that should be an integer) without abandoning \
+            ``select *`` — remaining columns pass through untouched. \
+            Caveat: a predicate on a cast column is pushed down as ``CAST(col AS ...) <op> value``. \
+            Wrapping the column in a function can prevent the database from using an index on it \
+            (SARGability), so the effect on the query plan is backend dependent — some engines \
+            optimize specific conversions (for example SQL Server seeks on ``CAST(datetime AS date)``) \
+            while others fall back to a scan. For a hot path on a large indexed table, prefer a \
+            filter expressed directly on the physical column instead of the cast one.
         **kwargs: Additional arguments for the database connector
 
     Returns:
@@ -132,6 +146,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

@@ -17,6 +17,7 @@
     "convert_predicate_to_sql",
     "create_sqlglot_literal",
     "fix_three_part_identifiers",
+    "polars_dtype_to_sqlglot_type",
 )
 
 
@@ -34,6 +35,59 @@
 DEFAULT_MAX_IN_PREDICATE_SIZE = 4096
 
 
+def polars_dtype_to_sqlglot_type(dtype: pl.DataType | type[pl.DataType], *, strict: bool = False) -> sqlglot.exp.DataType:
+    """
+    Convert a Polars dtype (class or instance) to a SQLGlot ``DataType``.
+
+    Decimal precision and scale are preserved so a decimal cast does not silently
+    truncate the fractional part.
+
+    Args:
+        dtype (pl.DataType): The Polars dtype to translate.
+        strict (bool): When True, raise for a dtype with no dedicated SQL mapping
+            instead of falling back to ``VARCHAR``.
+
+    Returns:
+        sqlglot.exp.DataType: The SQL type to cast to.
+
+    Raises:
+        ValueError: If ``strict`` is set and ``dtype`` has no dedicated SQL mapping.
+    """
+    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.
 
@@ -435,35 +489,7 @@ 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")
+        sql_type = polars_dtype_to_sqlglot_type(node.dtype)
 
         # Create the CAST expression
         self.result = sqlglot.exp.Cast(this=input_expr, to=sql_type)
@@ -715,6 +741,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,