add more tests and don't run filter pushdown if no supported filter

Author: wengh

python/pyspark/sql/tests/test_python_datasource.py

@@ -254,8 +254,8 @@ def __init__(self):
                 self.has_filter = False
 
             def pushdownFilters(self, filters: List[Filter]) -> Iterable[Filter]:
-                assert len(filters) == 2
-                assert set(filters) == {EqualTo(("x",), 1), EqualTo(("y",), 2)}
+                assert len(filters) == 2, filters
+                assert set(filters) == {EqualTo(("x",), 1), EqualTo(("y",), 2)}, filters
                 self.has_filter = True
                 # pretend we support x = 1 filter but in fact we don't
                 # so we only return y = 2 filter
@@ -293,10 +293,10 @@ def pushdownFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                 yield EqualTo(("",), 1)
 
             def partitions(self):
-                ...
+                assert False
 
             def read(self, partition):
-                ...
+                assert False
 
         class TestDataSource(DataSource):
             @classmethod
@@ -313,6 +313,55 @@ def reader(self, schema) -> "DataSourceReader":
         with self.assertRaisesRegex(Exception, "DATA_SOURCE_EXTRANEOUS_FILTERS"):
             self.spark.read.format("test").load().filter("x = 1").show()
 
+    def test_filter_pushdown_error(self):
+        class TestDataSourceReader(DataSourceReader):
+            def pushdownFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+                raise Exception("dummy error")
+
+            def read(self, partition):
+                yield [1]
+
+        class TestDataSource(DataSource):
+            @classmethod
+            def name(cls):
+                return "test"
+
+            def schema(self):
+                return "x int"
+
+            def reader(self, schema) -> "DataSourceReader":
+                return TestDataSourceReader()
+
+        self.spark.dataSource.register(TestDataSource)
+        df = self.spark.read.format("test").load().filter("cos(x) > 0")
+        assertDataFrameEqual(df, [Row(x=1)])  # works when not pushing down filters
+        with self.assertRaisesRegex(Exception, "dummy error"):
+            df.filter("x = 1").show()
+
+    def test_unsupported_filter(self):
+        class TestDataSourceReader(DataSourceReader):
+            def pushdownFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+                assert filters == [EqualTo(("x",), 1)], filters
+                return filters
+
+            def read(self, partition):
+                yield [1, 2, 3]
+
+        class TestDataSource(DataSource):
+            @classmethod
+            def name(cls):
+                return "test"
+
+            def schema(self):
+                return "x int, y int, z int"
+
+            def reader(self, schema) -> "DataSourceReader":
+                return TestDataSourceReader()
+
+        self.spark.dataSource.register(TestDataSource)
+        df = self.spark.read.format("test").load().filter("x = 1 and y = z")
+        assertDataFrameEqual(df, [])
+
     def _get_test_json_data_source(self):
         import json
         import os

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/python/UserDefinedPythonDataSource.scala

@@ -80,10 +80,15 @@ case class UserDefinedPythonDataSource(dataSourceCls: PythonFunction) {
   def pushdownFiltersInPython(
       pythonResult: PythonDataSourceReader,
       filters: Array[Filter]): PythonFilterPushdownResult = {
-    new UserDefinedPythonDataSourceFilterPushdownRunner(
+    val runner = new UserDefinedPythonDataSourceFilterPushdownRunner(
       createPythonFunction(pythonResult.reader),
       filters
-    ).runInPython()
+    )
+    if (runner.isAnyFilterSupported) {
+      runner.runInPython()
+    } else {
+      PythonFilterPushdownResult(pythonResult, filters.map(_ => false))
+    }
   }
 
   /**
@@ -409,6 +414,8 @@ private class UserDefinedPythonDataSourceFilterPushdownRunner(
   // See the logic in `pyspark.sql.worker.data_source_pushdown_filters.py`.
   override val workerModule = "pyspark.sql.worker.data_source_pushdown_filters"
 
+  def isAnyFilterSupported: Boolean = !serializedFilters.isEmpty
+
   override protected def writeToPython(dataOut: DataOutputStream, pickler: Pickler): Unit = {
     // Send Python data source
     PythonWorkerUtils.writePythonFunction(reader, dataOut)