Skip to content

Commit 2c45a71

Browse files
authored
[data] Fix MapTransformFn __eq__ (#51434)
## Why are these changes needed? <!-- Please give a short summary of the change and the problem this solves. --> Fix `MapTransformFn` `__eq__` for below by also testing for `_is_udf`. * `RowMapTransformFn` * `BatchMapTransformFn` * `RowToBlockMapTransformFn` Signed-off-by: Srinath Krishnamachari <srinath.krishnamachari@anyscale.com>
1 parent a3d66da commit 2c45a71

File tree

1 file changed

+9
-2
lines changed

1 file changed

+9
-2
lines changed

python/ray/data/_internal/execution/operators/map_transformer.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -305,7 +305,11 @@ def __repr__(self) -> str:
305305
return f"RowMapTransformFn({self._row_fn})"
306306

307307
def __eq__(self, other):
308-
return isinstance(other, RowMapTransformFn) and self._row_fn == other._row_fn
308+
return (
309+
isinstance(other, RowMapTransformFn)
310+
and self._row_fn == other._row_fn
311+
and self._is_udf == other._is_udf
312+
)
309313

310314

311315
class BatchMapTransformFn(MapTransformFn):
@@ -332,7 +336,9 @@ def __repr__(self) -> str:
332336

333337
def __eq__(self, other):
334338
return (
335-
isinstance(other, BatchMapTransformFn) and self._batch_fn == other._batch_fn
339+
isinstance(other, BatchMapTransformFn)
340+
and self._batch_fn == other._batch_fn
341+
and self._is_udf == other._is_udf
336342
)
337343

338344

@@ -357,6 +363,7 @@ def __eq__(self, other):
357363
return (
358364
isinstance(other, RowToBlockMapTransformFn)
359365
and self._transform_fn == other._transform_fn
366+
and self._is_udf == other._is_udf
360367
)
361368

362369

0 commit comments

Comments
 (0)