Fix validating pyspark dataframes with regex columns

pandera/api/pyspark/container.py

@@ -234,7 +234,9 @@ def get_dtypes(self, dataframe: DataFrame) -> Dict[str, DataType]:
                 regex_dtype.update(
                     {
                         c: column.dtype
-                        for c in column.BACKEND.get_regex_columns(
+                        for c in column.get_backend(
+                            dataframe
+                        ).get_regex_columns(
                             column,
                             dataframe.columns,
                         )

pandera/backends/pyspark/container.py

@@ -252,7 +252,7 @@
             if col_schema.regex:
                 try:
                     column_names.extend(
-                        col_schema.BACKEND.get_regex_columns(
+                        col_schema.get_backend(check_obj).get_regex_columns(
                             col_schema, check_obj.columns
                         )
                     )
@@ -457,9 +457,9 @@
         for colname, col_schema in schema.columns.items():
             if col_schema.regex:
                 try:
-                    matched_columns = col_schema.BACKEND.get_regex_columns(
-                        col_schema, obj.columns
-                    )
+                    matched_columns = col_schema.get_backend(
+                        obj
+                    ).get_regex_columns(col_schema, obj.columns)
                 except SchemaError:
                     matched_columns = None
 

tests/pyspark/test_pyspark_container.py

@@ -157,3 +157,31 @@ def test_pyspark_sample():
     df_out = schema.validate(df, sample=0.5)
 
     assert isinstance(df_out, DataFrame)
+
+
+def test_pyspark_regex_column():
+    """
+    Test creating a pyspark DataFrameSchema object with regex columns
+    """
+
+    schema = DataFrameSchema(
+        {
+            # Columns with all caps names must have string values
+            "[A-Z]+": Column(T.StringType(), regex=True),
+        }
+    )
+
+    data = [("Neeraj", 35), ("Jask", 30)]
+
+    df = spark.createDataFrame(data=data, schema=["NAME", "AGE"])
+    df_out = schema.validate(df)
+
+    assert df_out.pandera.errors is not None
+
+    data = [("Neeraj", "35"), ("Jask", "a")]
+
+    df2 = spark.createDataFrame(data=data, schema=["NAME", "AGE"])
+
+    df_out = schema.validate(df2)
+
+    assert not df_out.pandera.errors