[Data] - Unify Project Operator to use Expressions

doc/source/data/api/expressions.rst

@@ -17,6 +17,7 @@ Public API
     :nosignatures:
     :toctree: doc/
 
+    star
     col
     lit
     udf
@@ -37,4 +38,5 @@ instantiate them directly, but you may encounter them when working with expressi
     LiteralExpr
     BinaryExpr
     UnaryExpr
-    UDFExpr
+    UDFExpr
+    StarColumnsExpr

python/ray/data/_internal/arrow_block.py

@@ -58,6 +58,7 @@
 
 
 _MIN_PYARROW_VERSION_TO_NUMPY_ZERO_COPY_ONLY = parse_version("13.0.0")
+_BATCH_SIZE_PRESERVING_STUB_COL_NAME = "__bsp_stub"
 
 
 # Set the max chunk size in bytes for Arrow to Batches conversion in
@@ -221,7 +222,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":
@@ -453,7 +454,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/datasource/parquet_datasource.py

@@ -23,7 +23,10 @@
 
 import ray
 from ray._private.arrow_utils import get_pyarrow_version
-from ray.data._internal.arrow_block import ArrowBlockAccessor
+from ray.data._internal.arrow_block import (
+    _BATCH_SIZE_PRESERVING_STUB_COL_NAME,
+    ArrowBlockAccessor,
+)
 from ray.data._internal.progress_bar import ProgressBar
 from ray.data._internal.remote_fn import cached_remote_fn
 from ray.data._internal.util import (
@@ -104,9 +107,6 @@
 PARQUET_ENCODING_RATIO_ESTIMATE_NUM_ROWS = 1024
 
 
-_BATCH_SIZE_PRESERVING_STUB_COL_NAME = "__bsp_stub"
-
-
 class _ParquetFragment:
     """This wrapper class is created to avoid utilizing `ParquetFileFragment` original
     serialization protocol that actually does network RPCs during serialization

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

@@ -7,7 +7,7 @@
 from ray.data._internal.logical.interfaces import LogicalOperator
 from ray.data._internal.logical.operators.one_to_one_operator import AbstractOneToOne
 from ray.data.block import UserDefinedFunction
-from ray.data.expressions import Expr
+from ray.data.expressions import Expr, StarColumnsExpr
 from ray.data.preprocessor import Preprocessor
 
 logger = logging.getLogger(__name__)
@@ -268,16 +268,12 @@ 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"],
         compute: Optional[ComputeStrategy] = None,
         ray_remote_args: Optional[Dict[str, Any]] = None,
     ):
@@ -288,30 +284,23 @@ 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
 
-        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 expr.name is None and not isinstance(expr, StarColumnsExpr):
+                raise TypeError(
+                    "All Project expressions must be named (use .alias(name) or col(name)), "
+                    "or be a star() expression."
+                )
 
-    @property
-    def cols(self) -> Optional[List[str]]:
-        return self._cols
-
-    @property
-    def cols_rename(self) -> Optional[Dict[str, str]]:
-        return self._cols_rename
+    def has_all_columns_expr(self) -> bool:
+        """Check if this projection contains a star() expression."""
+        return any(isinstance(expr, StarColumnsExpr) for expr in self._exprs)
 
     @property
-    def exprs(self) -> Optional[Dict[str, "Expr"]]:
+    def exprs(self) -> List["Expr"]:
         return self._exprs
 
     def can_modify_num_rows(self) -> bool: