Fixing bug related to using coerce on a non-existent non-required col…

…umn (#1871) * Fixing bug related to using coerce on a non-existent non-required column Signed-off-by: mata <[email protected]> * Fixing test Signed-off-by: mata <[email protected]> * Fixes broken tests Signed-off-by: mata <[email protected]> --------- Signed-off-by: mata <[email protected]>
unionai-oss · Dec 10, 2024 · fccb7ac · fccb7ac
1 parent 9f667b4
commit fccb7ac
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/pandera/backends/polars/container.py b/pandera/backends/polars/container.py
@@ -446,11 +446,19 @@ def _coerce_dtype_helper(
             else "coerce"
         )
 
+        lf_columns = get_lazyframe_column_names(obj)
+
         try:
             if schema.dtype is not None:
                 obj = getattr(schema.dtype, coerce_fn)(obj)
             else:
                 for col_schema in schema.columns.values():
+                    if (
+                        not col_schema.required
+                        and col_schema.name not in lf_columns
+                    ):
+                        continue
+
                     if schema.coerce or col_schema.coerce:
                         obj = getattr(col_schema.dtype, coerce_fn)(
                             PolarsData(obj, col_schema.selector)

diff --git a/tests/polars/test_polars_container.py b/tests/polars/test_polars_container.py
@@ -177,6 +177,15 @@ def test_coerce_df_dtype_error(ldf_basic, ldf_schema_basic):
         modified_ldf.pipe(ldf_schema_basic.validate)
 
 
+def test_coerce_df_missing_column(ldf_basic, ldf_schema_basic):
+    """Test coerce dtype when column does not exist and not required."""
+    ldf_schema_basic.columns["string_col"].required = False
+    ldf_schema_basic.columns["string_col"].coerce = True
+    modified_ldf = ldf_basic.drop("string_col")
+    # should not raise an error
+    modified_ldf.pipe(ldf_schema_basic.validate)
+
+
 def test_strict_filter(ldf_basic, ldf_schema_basic):
     """Test strictness and filtering schema logic."""
     # by default, strict is False, so by default it should pass