[Data] - Unify Project Operator to use Expressions

python/ray/data/_internal/arrow_block.py

@@ -160,6 +160,24 @@ def _table_from_pydict(columns: Dict[str, List[Any]]) -> Block:
     @staticmethod
     def _combine_tables(tables: List[Block]) -> Block:
         if len(tables) > 1:
+            # Check if we have 0-column tables to avoid losing rows during concat
+            # PyArrow's concat on 0-column tables returns a 0-row table
+            if all(table.num_columns == 0 for table in tables):
+                # Add stub column to preserve rows during concatenation
+                import pyarrow as pa
+
+                tables_with_stub = []
+                for table in tables:
+                    if table.num_rows > 0:
+                        table = table.append_column(
+                            "__concat_stub", pa.nulls(table.num_rows)
+                        )
+                    tables_with_stub.append(table)
+                result = transform_pyarrow.concat(tables_with_stub, promote_types=True)
+                # Remove stub column after concatenation
+                if result.num_columns > 0:
+                    result = result.select([])
+                return result
             return transform_pyarrow.concat(tables, promote_types=True)
         else:
             return tables[0]
@@ -221,7 +239,7 @@ def fill_column(self, name: str, value: Any) -> Block:
 
             array = pyarrow.nulls(len(self._table), type=type)
             array = pc.fill_null(array, value)
-            return self._table.append_column(name, array)
+            return self.upsert_column(name, array)
 
     @classmethod
     def from_bytes(cls, data: bytes) -> "ArrowBlockAccessor":
@@ -321,9 +339,8 @@ def to_arrow(self) -> "pyarrow.Table":
         return self._table
 
     def num_rows(self) -> int:
-        # Arrow may represent an empty table via an N > 0 row, 0-column table, e.g. when
-        # slicing an empty table, so we return 0 if num_columns == 0.
-        return self._table.num_rows if self._table.num_columns > 0 else 0
+        # Arrow may represent an empty table via an N > 0 row, 0-column table
+        return self._table.num_rows
 
     def size_bytes(self) -> int:
         return self._table.nbytes
@@ -469,7 +486,9 @@ def filter(self, predicate_expr: "Expr") -> "pyarrow.Table":
         if self._table.num_rows == 0:
             return self._table
 
-        from ray.data._expression_evaluator import eval_expr
+        from ray.data._internal.planner.plan_expression.expression_evaluator import (
+            eval_expr,
+        )
 
         # Evaluate the expression to get a boolean mask
         mask = eval_expr(predicate_expr, self._table)

python/ray/data/_internal/logical/operators/map_operator.py

@@ -268,16 +268,13 @@ def can_modify_num_rows(self) -> bool:
 
 
 class Project(AbstractMap):
-    """Logical operator for select_columns."""
+    """Logical operator for select_columns, with_column, rename_columns."""
 
     def __init__(
         self,
         input_op: LogicalOperator,
-        cols: Optional[List[str]] = None,
-        cols_rename: Optional[Dict[str, str]] = None,
-        exprs: Optional[
-            Dict[str, "Expr"]
-        ] = None,  # TODO Remove cols and cols_rename and replace them with corresponding exprs
+        exprs: List["Expr"],
+        preserve_existing: bool = False,
         compute: Optional[ComputeStrategy] = None,
         ray_remote_args: Optional[Dict[str, Any]] = None,
     ):
@@ -288,30 +285,25 @@ def __init__(
             compute=compute,
         )
         self._batch_size = None
-        self._cols = cols
-        self._cols_rename = cols_rename
         self._exprs = exprs
         self._batch_format = "pyarrow"
         self._zero_copy_batch = True
+        self._preserve_existing = preserve_existing
 
-        if exprs is not None:
-            # Validate that all values are expressions
-            for name, expr in exprs.items():
-                if not isinstance(expr, Expr):
-                    raise TypeError(
-                        f"Expected Expr for column '{name}', got {type(expr)}"
-                    )
+        for expr in self._exprs:
+            if not isinstance(expr, Expr):
+                raise TypeError(f"Expected Expr got {expr} with type: {type(expr)}")
+            if expr.name is None:
+                raise TypeError(
+                    "All Project expressions must be named; use .alias(name) or col(name)."
+                )
 
     @property
-    def cols(self) -> Optional[List[str]]:
-        return self._cols
+    def preserve_existing(self) -> bool:
+        return self._preserve_existing
 
     @property
-    def cols_rename(self) -> Optional[Dict[str, str]]:
-        return self._cols_rename
-
-    @property
-    def exprs(self) -> Optional[Dict[str, "Expr"]]:
+    def exprs(self) -> List["Expr"]:
         return self._exprs
 
     def can_modify_num_rows(self) -> bool: