WIP filter serialization

Author: wengh

python/pyspark/sql/datasource.py

@@ -312,6 +312,70 @@ class EqualTo(Filter):
     value: Any
 
 
+@dataclass(frozen=True)
+class EqualNullSafe(Filter):
+    attribute: ColumnPath
+    value: Any
+
+
+@dataclass(frozen=True)
+class GreaterThan(Filter):
+    attribute: ColumnPath
+    value: Any
+
+
+@dataclass(frozen=True)
+class GreaterThanOrEqual(Filter):
+    attribute: ColumnPath
+    value: Any
+
+
+@dataclass(frozen=True)
+class LessThan(Filter):
+    attribute: ColumnPath
+    value: Any
+
+
+@dataclass(frozen=True)
+class LessThanOrEqual(Filter):
+    attribute: ColumnPath
+    value: Any
+
+
+@dataclass(frozen=True)
+class In(Filter):
+    attribute: ColumnPath
+    value: Tuple[Any, ...]
+
+
+@dataclass(frozen=True)
+class IsNull(Filter):
+    attribute: ColumnPath
+
+
+@dataclass(frozen=True)
+class IsNotNull(Filter):
+    attribute: ColumnPath
+
+
+@dataclass(frozen=True)
+class StringStartsWith(Filter):
+    attribute: ColumnPath
+    value: str
+
+
+@dataclass(frozen=True)
+class StringEndsWith(Filter):
+    attribute: ColumnPath
+    value: str
+
+
+@dataclass(frozen=True)
+class StringContains(Filter):
+    attribute: ColumnPath
+    value: str
+
+
 class InputPartition:
     """
     A base class representing an input partition returned by the `partitions()`

python/pyspark/sql/tests/test_python_datasource.py

@@ -18,29 +18,43 @@
 import platform
 import tempfile
 import unittest
+from datetime import datetime
+from decimal import Decimal
 from typing import Callable, Iterable, List, Union
 
-from pyspark.errors import PythonException, AnalysisException
+from pyspark.errors import AnalysisException, PythonException
 from pyspark.sql.datasource import (
+    CaseInsensitiveDict,
     DataSource,
+    DataSourceArrowWriter,
     DataSourceReader,
+    DataSourceWriter,
+    EqualNullSafe,
     EqualTo,
     Filter,
+    GreaterThan,
+    GreaterThanOrEqual,
+    In,
     InputPartition,
-    DataSourceWriter,
-    DataSourceArrowWriter,
+    IsNotNull,
+    IsNull,
+    LessThan,
+    LessThanOrEqual,
+    StringContains,
+    StringEndsWith,
+    StringStartsWith,
     WriterCommitMessage,
-    CaseInsensitiveDict,
 )
 from pyspark.sql.functions import spark_partition_id
 from pyspark.sql.session import SparkSession
 from pyspark.sql.types import Row, StructType
+from pyspark.testing import assertDataFrameEqual
 from pyspark.testing.sqlutils import (
+    SPARK_HOME,
+    ReusedSQLTestCase,
     have_pyarrow,
     pyarrow_requirement_message,
 )
-from pyspark.testing import assertDataFrameEqual
-from pyspark.testing.sqlutils import ReusedSQLTestCase, SPARK_HOME
 
 
 @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
@@ -258,6 +272,8 @@ def __init__(self):
 
             def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
                 assert set(filters) == {
+                    IsNotNull(("x",)),
+                    IsNotNull(("y",)),
                     EqualTo(("x",), 1),
                     EqualTo(("y",), 2),
                 }, filters
@@ -376,8 +392,9 @@ def _check_filters(self, sql_type, sql_filter, python_filters):
 
         class TestDataSourceReader(DataSourceReader):
             def pushFilters(self, filters: List[Filter]) -> Iterable[Filter]:
+                actual = [f for f in filters if not isinstance(f, IsNotNull)]
                 expected = python_filters
-                assert filters == expected, (filters, expected)
+                assert actual == expected, (actual, expected)
                 return filters
 
             def read(self, partition):
@@ -406,6 +423,45 @@ def test_unsupported_filter(self):
         self._check_filters("boolean", "not x", [])
         self._check_filters("array<int>", "x[0] = 1", [])
 
+    def test_filter_value_type(self):
+        self._check_filters("int", "x = 1", [EqualTo(("x",), 1)])
+        self._check_filters("int", "x = null", [EqualTo(("x",), None)])
+        self._check_filters("float", "x = 3 / 2", [EqualTo(("x",), 1.5)])
+        self._check_filters("string", "x = '1'", [EqualTo(("x",), "1")])
+        self._check_filters("array<int>", "x = array(1, 2)", [EqualTo(("x",), [1, 2])])
+        self._check_filters(
+            "struct<x:int>", "x = named_struct('x', 1)", [EqualTo(("x",), {"x": 1})]
+        )
+        self._check_filters(
+            "decimal", "x in (1.1, 2.1)", [In(("x",), [Decimal(1.1), Decimal(2.1)])]
+        )
+        self._check_filters(
+            "timestamp_ntz",
+            "x = timestamp_ntz '2020-01-01 00:00:00'",
+            [EqualTo(("x",), datetime.strptime("2020-01-01 00:00:00", "%Y-%m-%d %H:%M:%S"))],
+        )
+        self._check_filters(
+            "interval second",
+            "x = interval '2' second",
+            [],  # intervals are not supported
+        )
+
+    def test_filter_type(self):
+        self._check_filters("boolean", "x", [EqualTo(("x",), True)])
+        self._check_filters("int", "x is null", [IsNull(("x",))])
+        self._check_filters("int", "x <=> 1", [EqualNullSafe(("x",), 1)])
+        self._check_filters("int", "1 < x", [GreaterThan(("x",), 1)])
+        self._check_filters("int", "1 <= x", [GreaterThanOrEqual(("x",), 1)])
+        self._check_filters("int", "x < 1", [LessThan(("x",), 1)])
+        self._check_filters("int", "x <= 1", [LessThanOrEqual(("x",), 1)])
+        self._check_filters("string", "startswith(x, 'a')", [StringStartsWith(("x",), "a")])
+        self._check_filters("string", "endswith(x, 'a')", [StringEndsWith(("x",), "a")])
+        self._check_filters("string", "contains(x, 'a')", [StringContains(("x",), "a")])
+        self._check_filters("int", "x in (1, 2)", [In(("x",), [1, 2])])
+
+    def test_filter_nested_column(self):
+        self._check_filters("struct<y:int>", "x.y = 1", [EqualTo(("x", "y"), 1)])
+
     def _get_test_json_data_source(self):
         import json
         import os

python/pyspark/sql/worker/data_source_pushdown_filters.py

@@ -15,17 +15,37 @@
 # limitations under the License.
 #
 
+import base64
 import faulthandler
+import json
 import os
 import sys
+import typing
 from dataclasses import dataclass, field
-from typing import IO, List
+from typing import IO, Type, Union
 
 from pyspark.accumulators import _accumulatorRegistry
 from pyspark.errors import PySparkAssertionError, PySparkValueError
+from pyspark.errors.exceptions.base import PySparkNotImplementedError
 from pyspark.serializers import SpecialLengths, UTF8Deserializer, read_int, write_int
-from pyspark.sql.datasource import DataSource, DataSourceReader, EqualTo, Filter
-from pyspark.sql.types import StructType, _parse_datatype_json_string
+from pyspark.sql.datasource import (
+    DataSource,
+    DataSourceReader,
+    EqualNullSafe,
+    EqualTo,
+    Filter,
+    GreaterThan,
+    GreaterThanOrEqual,
+    In,
+    IsNotNull,
+    IsNull,
+    LessThan,
+    LessThanOrEqual,
+    StringContains,
+    StringEndsWith,
+    StringStartsWith,
+)
+from pyspark.sql.types import StructType, VariantVal, _parse_datatype_json_string
 from pyspark.util import handle_worker_exception, local_connect_and_auth
 from pyspark.worker_util import (
     check_python_version,
@@ -39,6 +59,25 @@
 
 utf8_deserializer = UTF8Deserializer()
 
+BinaryFilter = Union[
+    EqualTo,
+    EqualNullSafe,
+    GreaterThan,
+    GreaterThanOrEqual,
+    LessThan,
+    LessThanOrEqual,
+    In,
+    StringStartsWith,
+    StringEndsWith,
+    StringContains,
+]
+
+binary_filters = {cls.__name__: cls for cls in typing.get_args(BinaryFilter)}
+
+UnaryFilter = Union[IsNotNull, IsNull]
+
+unary_filters = {cls.__name__: cls for cls in typing.get_args(UnaryFilter)}
+
 
 @dataclass(frozen=True)
 class FilterRef:
@@ -49,6 +88,30 @@ def __post_init__(self) -> None:
         object.__setattr__(self, "id", id(self.filter))
 
 
+def deserializeVariant(variantDict: dict) -> VariantVal:
+    value = base64.b64decode(variantDict["value"])
+    metadata = base64.b64decode(variantDict["metadata"])
+    return VariantVal(value, metadata)
+
+
+def deserializeFilter(jsonDict: dict) -> Filter:
+    name = jsonDict["name"]
+    if name in binary_filters:
+        binary_filter_cls: Type[BinaryFilter] = binary_filters[name]
+        return binary_filter_cls(
+            attribute=tuple(jsonDict["columnPath"]),
+            value=deserializeVariant(jsonDict["value"]).toPython(),
+        )
+    elif name in unary_filters:
+        unary_filter_cls: Type[UnaryFilter] = unary_filters[name]
+        return unary_filter_cls(attribute=tuple(jsonDict["columnPath"]))
+    else:
+        raise PySparkNotImplementedError(
+            errorClass="UNSUPPORTED_FILTER",
+            messageParameters={"name": name},
+        )
+
+
 def main(infile: IO, outfile: IO) -> None:
     """
     Main method for planning a data source read with filter pushdown.
@@ -126,22 +189,9 @@ def main(infile: IO, outfile: IO) -> None:
             )
 
         # Receive the pushdown filters.
-        num_filters = read_int(infile)
-        filters: List[FilterRef] = []
-        for _ in range(num_filters):
-            name = utf8_deserializer.loads(infile)
-            if name == "EqualTo":
-                num_parts = read_int(infile)
-                column_path = tuple(utf8_deserializer.loads(infile) for _ in range(num_parts))
-                value = read_int(infile)
-                filters.append(FilterRef(EqualTo(column_path, value)))
-            else:
-                raise PySparkAssertionError(
-                    errorClass="DATA_SOURCE_UNSUPPORTED_FILTER",
-                    messageParameters={
-                        "name": name,
-                    },
-                )
+        json_str = utf8_deserializer.loads(infile)
+        filter_dicts = json.loads(json_str)
+        filters = [FilterRef(deserializeFilter(f)) for f in filter_dicts]
 
         # Push down the filters and get the indices of the unsupported filters.
         unsupported_filters = set(