[Data] - Fix expression mapping for Pandas (#57868)

Author: goutamvenkat-anyscale

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

@@ -53,11 +53,11 @@ def _pa_is_in(left: Any, right: Any) -> Any:
     Operation.NE: operator.ne,
     Operation.AND: operator.and_,
     Operation.OR: operator.or_,
-    Operation.NOT: operator.not_,
+    Operation.NOT: operator.invert,
     Operation.IS_NULL: pd.isna,
     Operation.IS_NOT_NULL: pd.notna,
-    Operation.IN: lambda left, right: left.is_in(right),
-    Operation.NOT_IN: lambda left, right: ~left.is_in(right),
+    Operation.IN: lambda left, right: left.isin(right),
+    Operation.NOT_IN: lambda left, right: ~left.isin(right),
 }
 
 

python/ray/data/tests/test_filter.py

@@ -123,6 +123,13 @@ def test_filter_with_invalid_expression(ray_start_regular_shared, tmp_path):
     get_pyarrow_version() < parse_version("20.0.0"),
     reason="predicate expressions require PyArrow >= 20.0.0",
 )
+@pytest.mark.parametrize(
+    "data_source",
+    [
+        pytest.param("from_items", id="arrow_blocks"),
+        pytest.param("from_pandas", id="pandas_blocks"),
+    ],
+)
 @pytest.mark.parametrize(
     "predicate_expr, test_data, expected_indices, test_description",
     [
@@ -230,6 +237,38 @@ def test_filter_with_invalid_expression(ray_start_regular_shared, tmp_path):
             [0, 3],
             "string_exclusion_filter",
         ),
+        # Additional comparison operations
+        pytest.param(
+            col("age") > 25,
+            [
+                {"age": 20, "name": "Alice"},
+                {"age": 25, "name": "Bob"},
+                {"age": 30, "name": "Charlie"},
+                {"age": 35, "name": "David"},
+            ],
+            [2, 3],
+            "greater_than_filter",
+        ),
+        pytest.param(
+            col("age") < 25,
+            [
+                {"age": 20, "name": "Alice"},
+                {"age": 25, "name": "Bob"},
+                {"age": 30, "name": "Charlie"},
+            ],
+            [0],
+            "less_than_filter",
+        ),
+        pytest.param(
+            col("age") <= 25,
+            [
+                {"age": 20, "name": "Alice"},
+                {"age": 25, "name": "Bob"},
+                {"age": 30, "name": "Charlie"},
+            ],
+            [0, 1],
+            "less_than_equal_filter",
+        ),
         # Membership operations
         pytest.param(
             col("category").is_in(["A", "B"]),
@@ -241,7 +280,18 @@ def test_filter_with_invalid_expression(ray_start_regular_shared, tmp_path):
                 {"category": "A", "value": 5},
             ],
             [0, 1, 4],
-            "membership_filter",
+            "is_in_filter",
+        ),
+        pytest.param(
+            col("category").not_in(["A", "B"]),
+            [
+                {"category": "A", "value": 1},
+                {"category": "B", "value": 2},
+                {"category": "C", "value": 3},
+                {"category": "D", "value": 4},
+            ],
+            [2, 3],  # These are indices not the actual values
+            "not_in_filter",
         ),
         # Negation operations
         pytest.param(
@@ -271,14 +321,18 @@ def test_filter_with_invalid_expression(ray_start_regular_shared, tmp_path):
 )
 def test_filter_with_predicate_expressions(
     ray_start_regular_shared,
+    data_source,
     predicate_expr,
     test_data,
     expected_indices,
     test_description,
 ):
-    """Test filter() with Ray Data predicate expressions."""
-    # Create dataset from test data
-    ds = ray.data.from_items(test_data)
+    """Test filter() with Ray Data predicate expressions on both Arrow and pandas blocks."""
+    # Create dataset based on data_source parameter
+    if data_source == "from_items":
+        ds = ray.data.from_items(test_data)
+    else:  # from_pandas
+        ds = ray.data.from_pandas([pd.DataFrame(test_data)])
 
     # Apply filter with predicate expression
     filtered_ds = ds.filter(expr=predicate_expr)