wengh commented on code in PR #49961:
URL: https://github.com/apache/spark/pull/49961#discussion_r2003793967
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/PythonDataSourceSuite.scala:
##########
@@ -213,6 +220,66 @@ class PythonDataSourceSuite extends
PythonDataSourceSuiteBase {
parameters = Map("inputSchema" -> "INT", "dataType" -> "\"INT\""))
}
+ test("data source reader with filter pushdown") {
+ assume(shouldTestPandasUDFs)
+ val dataSourceScript =
+ s"""
+ |from pyspark.sql.datasource import (
+ | DataSource,
+ | DataSourceReader,
+ | EqualTo,
+ | InputPartition,
+ |)
+ |
+ |class SimpleDataSourceReader(DataSourceReader):
+ | def partitions(self):
+ | return [InputPartition(i) for i in range(2)]
+ |
+ | def pushFilters(self, filters):
+ | yield filters[filters.index(EqualTo(("id",), 1))]
+ |
+ | def read(self, partition):
+ | yield (0, partition.value)
+ | yield (1, partition.value)
+ | yield (2, partition.value)
+ |
+ |class SimpleDataSource(DataSource):
+ | def schema(self):
+ | return "id int, partition int"
+ |
+ | def reader(self, schema):
+ | return SimpleDataSourceReader()
+ |""".stripMargin
+ val schema = StructType.fromDDL("id INT, partition INT")
+ val dataSource =
+ createUserDefinedPythonDataSource(name = dataSourceName, pythonScript =
dataSourceScript)
+ spark.conf.set(SQLConf.PYTHON_FILTER_PUSHDOWN_ENABLED, true)
+ spark.dataSource.registerPython(dataSourceName, dataSource)
+ val df =
+ spark.read.format(dataSourceName).schema(schema).load().filter("id = 1
and partition = 0")
+ val plan = df.queryExecution.executedPlan
+
+ val filter = collectFirst(df.queryExecution.executedPlan) {
Review Comment:
```
== Physical Plan ==
*(1) Project [id#261, partition#262]
+- *(1) Filter ((isnotnull(id#261) AND isnotnull(partition#262)) AND (id#261
= 1))
+- BatchScan SimpleDataSource[id#261, partition#262] (Python)
PushedFilters: [EqualTo(partition,0)],
ReadSchema: struct<id:int,partition:int> RuntimeFilters: []
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
