fix: Ignore dy.Any columns in Schema.cast (#315)

Co-authored-by: gabriel <gabriel.g.robin@airbus.com>
Co-authored-by: Andreas Albert <103571926+AndreasAlbertQC@users.noreply.github.com>
Co-authored-by: Andreas Albert <andreas.albert@quantco.com>

Author: 3 people

dataframely/schema.py

@@ -813,9 +813,7 @@ def cast(
             the lazy frame's schema but also means that a call to :meth:`polars.LazyFrame.collect`
             further down the line might fail because of the cast and/or missing columns.
         """
-        lf = df.lazy().select(
-            pl.col(name).cast(col.dtype) for name, col in cls.columns().items()
-        )
+        lf = match_to_schema(df.lazy(), cls, casting="strict")
         if isinstance(df, pl.DataFrame):
             return lf.collect()  # type: ignore
         return lf  # type: ignore

tests/collection/test_cast.py

@@ -2,10 +2,10 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 import polars as pl
-import polars.exceptions as plexc
 import pytest
 
 import dataframely as dy
+from dataframely.exc import SchemaError
 
 
 class FirstSchema(dy.Schema):
@@ -48,12 +48,12 @@ def test_cast_invalid_members(df_type: type[pl.DataFrame] | type[pl.LazyFrame])
 
 def test_cast_invalid_member_schema_eager() -> None:
     first = pl.DataFrame({"b": [3]})
-    with pytest.raises(plexc.ColumnNotFoundError):
+    with pytest.raises(SchemaError):
         Collection.cast({"first": first})
 
 
 def test_cast_invalid_member_schema_lazy() -> None:
     first = pl.LazyFrame({"b": [3]})
     collection = Collection.cast({"first": first})
-    with pytest.raises(plexc.ColumnNotFoundError):
+    with pytest.raises(SchemaError):
         collection.collect_all()

tests/column_types/test_any.py

@@ -20,3 +20,9 @@ class AnySchema(dy.Schema):
 def test_any_dtype_passes(data: dict[str, Any]) -> None:
     df = pl.DataFrame(data)
     assert AnySchema.is_valid(df)
+
+
+def test_any_cast() -> None:
+    df = pl.DataFrame({"a": 0})
+    result = AnySchema.cast(df)
+    assert result["a"].dtype == pl.Int64

tests/schema/test_cast.py

@@ -1,13 +1,12 @@
 # Copyright (c) QuantCo 2025-2026
 # SPDX-License-Identifier: BSD-3-Clause
-
 from typing import Any
 
 import polars as pl
-import polars.exceptions as plexc
 import pytest
 
 import dataframely as dy
+from dataframely.exc import SchemaError
 
 
 class MySchema(dy.Schema):
@@ -34,12 +33,27 @@ def test_cast_valid(
 
 def test_cast_invalid_schema_eager() -> None:
     df = pl.DataFrame({"a": [1]})
-    with pytest.raises(plexc.ColumnNotFoundError):
+    with pytest.raises(SchemaError):
         MySchema.cast(df)
 
 
 def test_cast_invalid_schema_lazy() -> None:
     lf = pl.LazyFrame({"a": [1]})
     lf = MySchema.cast(lf)
-    with pytest.raises(plexc.ColumnNotFoundError):
+    with pytest.raises(SchemaError):
         lf.collect()
+
+
+class IntegerSchema(dy.Schema):
+    a = dy.Integer()
+
+
+@pytest.mark.parametrize("df_type", [pl.DataFrame, pl.LazyFrame])
+def test_cast_preserves_valid_dtype(
+    df_type: type[pl.DataFrame] | type[pl.LazyFrame],
+) -> None:
+    """Test that cast doesn't change already valid dtypes (issue #318)."""
+    df = df_type({"a": [1, 2, 3]}, schema={"a": pl.Int32})
+    result = IntegerSchema.cast(df)
+    # Int32 is valid for dy.Integer, so it should NOT be cast to Int64
+    assert result.lazy().collect_schema()["a"] == pl.Int32