[Data] - Evaluate Expressions sequentially and upsert columns to block

python/ray/data/_internal/arrow_block.py

@@ -335,6 +335,19 @@ def _zip(self, acc: BlockAccessor) -> "Block":
             r = r.append_column(col_name, col)
         return r
 
+    def upsert_column(
+        self, column_name: str, column_data: BlockColumn
+    ) -> "pyarrow.Table":
+        assert isinstance(
+            column_data, (pyarrow.Array, pyarrow.ChunkedArray)
+        ), f"Expected either a pyarrow.Array or pyarrow.ChunkedArray, got: {type(column_data)}"
+
+        column_idx = self._table.schema.get_field_index(column_name)
+        if column_idx == -1:
+            return self._table.append_column(column_name, column_data)
+        else:
+            return self._table.set_column(column_idx, column_name, column_data)
+
     @staticmethod
     def builder() -> ArrowBlockBuilder:
         return ArrowBlockBuilder()

python/ray/data/_internal/pandas_block.py

@@ -317,6 +317,14 @@ def select(self, columns: List[str]) -> "pandas.DataFrame":
     def rename_columns(self, columns_rename: Dict[str, str]) -> "pandas.DataFrame":
         return self._table.rename(columns=columns_rename, inplace=False, copy=False)
 
+    def upsert_column(
+        self, column_name: str, column_data: BlockColumn
+    ) -> "pandas.DataFrame":
+        if isinstance(column_data, (pyarrow.Array, pyarrow.ChunkedArray)):
+            column_data = column_data.to_pandas()
+
+        return self._table.assign(**{column_name: column_data})
+
     def random_shuffle(self, random_seed: Optional[int]) -> "pandas.DataFrame":
         table = self._table.sample(frac=1, random_state=random_seed)
         table.reset_index(drop=True, inplace=True)

python/ray/data/_internal/planner/plan_udf_map_op.py

@@ -128,19 +128,14 @@ def fn(block: Block) -> Block:
             # 1. evaluate / add expressions
             if exprs:
                 block_accessor = BlockAccessor.for_block(block)
-                new_columns = {}
-                for col_name in block_accessor.column_names():
-                    # For Arrow blocks, block[col_name] gives us a ChunkedArray
-                    # For Pandas blocks, block[col_name] gives us a Series
-                    new_columns[col_name] = block[col_name]
-
                 # Add/update with expression results
+                result_block = block
                 for name, expr in exprs.items():
-                    result = eval_expr(expr, block)
-                    new_columns[name] = result
+                    result = eval_expr(expr, result_block)
+                    result_block_accessor = BlockAccessor.for_block(result_block)
+                    result_block = result_block_accessor.upsert_column(name, result)
 
-                # Create a new block from the combined columns and add it
-                block = BlockAccessor.batch_to_block(new_columns)
+                block = result_block
 
             # 2. (optional) column projection
             if columns:

python/ray/data/block.py

@@ -343,6 +343,19 @@ def rename_columns(self, columns_rename: Dict[str, str]) -> Block:
         """Return the block reflecting the renamed columns."""
         raise NotImplementedError
 
+    def upsert_column(self, column_name: str, column_data: BlockColumn) -> Block:
+        """
+        Upserts a column into the block. If the column already exists, it will be replaced.
+
+        Args:
+            column_name: The name of the column to upsert.
+            column_data: The data to upsert into the column. (Arrow Array/ChunkedArray for Arrow blocks, Series or array-like for Pandas blocks)
+
+        Returns:
+            The updated block.
+        """
+        raise NotImplementedError()
+
     def random_shuffle(self, random_seed: Optional[int]) -> Block:
         """Randomly shuffle this block."""
         raise NotImplementedError