You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2022/11/30 00:56:27 UTC
[spark] branch master updated: [SPARK-41268][CONNECT][PYTHON] Refactor "Column" for API Compatibility
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new fa21bc52f1d [SPARK-41268][CONNECT][PYTHON] Refactor "Column" for API Compatibility
fa21bc52f1d is described below
commit fa21bc52f1d079aac4be38a4809a96317913c7e8
Author: Rui Wang <ru...@databricks.com>
AuthorDate: Wed Nov 30 09:56:11 2022 +0900
[SPARK-41268][CONNECT][PYTHON] Refactor "Column" for API Compatibility
### What changes were proposed in this pull request?
In the past, the expression system is pretty different between Connect and PySpark. `Column` is the first class API for all the expression in PySpark while `Expression` is the top level API in Connect. To maintain API Compatibility we need to promote `Column` in Connect. Also Connect has its special needs to maintain a series of `to_proto` methods to convert expressions to proto representation.
This PR:
1. Promote `Column` to be the first class citizen.
2. Wrap `Expression` into `Column` and do not expose it into DataFrame.
**The main idea of this PR**
We wrap `Expression` into `Column` and only expose `Column` to DataFrame API surface. In this concept, `Column` is the public API that user access (along with functions.py which offer a rich set of ways to create columns). `Expression` serve as internal API to differentiate different types of columns and also build a bridge between Python and Connect protobuf.
It is not decided yet if `plan.py` or Logical Plan should deal with Expression or Column, and same for `str` versus `Column` or `Expression`.
Examples as the following:
```
Class Column:
def __init__(self, expr: Expression):
self._expr = expr;
def like(self: "Column", other: str) -> "Column":
# UnresolvedFunction extends Expression for `like` operation
# LiteralExpression as expression to carry `str` which is the function paramter.
return Column(UnresolvedFunction("like", self, Column(LiteralExpression(str)))
df.select(df.data, df.data.like("secret")).collect()
data | bool
test | False
secret | True
```
**Caveats:**
TBD
**Alternative approach:**
We keep `Column` as a base class and other expressions inherits it. Still only use `Column` in DataFrame API surface.
```
Class Column:
def __init__(self, expr: Expression):
self._expr = expr;
def like(self: "Column", other: str) -> "Column":
return cast(Column, UnresolvedFunction("like", self, LiteralExpression(str)))
df.select(df.data, df.data.like("secret")).collect()
```
**Caveats:**
TBD
### Why are the changes needed?
API Compatibility
### Does this PR introduce _any_ user-facing change?
NO
### How was this patch tested?
Existing UT
Closes #38806 from amaliujia/refactor_column.
Authored-by: Rui Wang <ru...@databricks.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/sql/connect/_typing.py | 9 +-
python/pyspark/sql/connect/column.py | 180 ++++++++++++++-------
python/pyspark/sql/connect/dataframe.py | 59 +++----
python/pyspark/sql/connect/function_builder.py | 18 +--
python/pyspark/sql/connect/functions.py | 8 +-
python/pyspark/sql/connect/plan.py | 36 ++---
.../connect/test_connect_column_expressions.py | 10 +-
.../sql/tests/connect/test_connect_plan_only.py | 4 +-
8 files changed, 188 insertions(+), 136 deletions(-)
diff --git a/python/pyspark/sql/connect/_typing.py b/python/pyspark/sql/connect/_typing.py
index 262d71fcea1..8629d1c23cc 100644
--- a/python/pyspark/sql/connect/_typing.py
+++ b/python/pyspark/sql/connect/_typing.py
@@ -26,10 +26,7 @@ from typing import Union, Optional
import datetime
import decimal
-from pyspark.sql.connect.column import ScalarFunctionExpression, Expression, Column
-from pyspark.sql.connect.function_builder import UserDefinedFunction
-
-ExpressionOrString = Union[Expression, str]
+from pyspark.sql.connect.column import ScalarFunctionExpression, Column
ColumnOrName = Union[Column, str]
@@ -45,10 +42,10 @@ DateTimeLiteral = Union[datetime.datetime, datetime.date]
class FunctionBuilderCallable(Protocol):
- def __call__(self, *_: ExpressionOrString) -> ScalarFunctionExpression:
+ def __call__(self, *_: ColumnOrName) -> ScalarFunctionExpression:
...
class UserDefinedFunctionCallable(Protocol):
- def __call__(self, *_: ColumnOrName) -> UserDefinedFunction:
+ def __call__(self, *_: ColumnOrName) -> Column:
...
diff --git a/python/pyspark/sql/connect/column.py b/python/pyspark/sql/connect/column.py
index 663af074925..83e8b28da0f 100644
--- a/python/pyspark/sql/connect/column.py
+++ b/python/pyspark/sql/connect/column.py
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from typing import get_args, TYPE_CHECKING, Callable, Any
+
+from typing import get_args, TYPE_CHECKING, Callable, Any, Union
import json
import decimal
@@ -31,53 +32,34 @@ if TYPE_CHECKING:
def _bin_op(
name: str, doc: str = "binary function", reverse: bool = False
-) -> Callable[["Column", Any], "Expression"]:
- def _(self: "Column", other: Any) -> "Expression":
+) -> Callable[["Column", Any], "Column"]:
+ def _(self: "Column", other: Any) -> "Column":
from pyspark.sql.connect._typing import PrimitiveType
+ from pyspark.sql.connect.functions import lit
if isinstance(other, get_args(PrimitiveType)):
- other = LiteralExpression(other)
+ other = lit(other)
if not reverse:
- return ScalarFunctionExpression(name, self, other)
+ return scalar_function(name, self, other)
else:
- return ScalarFunctionExpression(name, other, self)
+ return scalar_function(name, other, self)
return _
+def scalar_function(op: str, *args: "Column") -> "Column":
+ return Column(ScalarFunctionExpression(op, *args))
+
+
+def sql_expression(expr: str) -> "Column":
+ return Column(SQLExpression(expr))
+
+
class Expression(object):
"""
Expression base class.
"""
- __gt__ = _bin_op(">")
- __lt__ = _bin_op("<")
- __add__ = _bin_op("+")
- __sub__ = _bin_op("-")
- __mul__ = _bin_op("*")
- __div__ = _bin_op("/")
- __truediv__ = _bin_op("/")
- __mod__ = _bin_op("%")
- __radd__ = _bin_op("+", reverse=True)
- __rsub__ = _bin_op("-", reverse=True)
- __rmul__ = _bin_op("*", reverse=True)
- __rdiv__ = _bin_op("/", reverse=True)
- __rtruediv__ = _bin_op("/", reverse=True)
- __pow__ = _bin_op("pow")
- __rpow__ = _bin_op("pow", reverse=True)
- __ge__ = _bin_op(">=")
- __le__ = _bin_op("<=")
-
- def __eq__(self, other: Any) -> "Expression": # type: ignore[override]
- """Returns a binary expression with the current column as the left
- side and the other expression as the right side.
- """
- from pyspark.sql.connect._typing import PrimitiveType
-
- if isinstance(other, get_args(PrimitiveType)):
- other = LiteralExpression(other)
- return ScalarFunctionExpression("==", self, other)
-
def __init__(self) -> None:
pass
@@ -124,6 +106,21 @@ class Expression(object):
assert not kwargs, "Unexpected kwargs where passed: %s" % kwargs
return ColumnAlias(self, list(alias), metadata)
+ def desc(self) -> "SortOrder":
+ ...
+
+ def asc(self) -> "SortOrder":
+ ...
+
+ def ascending(self) -> bool:
+ ...
+
+ def nullsLast(self) -> bool:
+ ...
+
+ def name(self) -> str:
+ ...
+
class ColumnAlias(Expression):
def __init__(self, parent: Expression, alias: list[str], metadata: Any):
@@ -168,6 +165,9 @@ class LiteralExpression(Expression):
TODO(SPARK-40533) This method always assumes the largest type and can thus
create weird interpretations of the literal."""
+
+ from pyspark.sql.connect.functions import lit
+
expr = proto.Expression()
if self._value is None:
expr.literal.null = True
@@ -195,33 +195,29 @@ class LiteralExpression(Expression):
elif isinstance(self._value, list):
expr.literal.array.SetInParent()
for item in list(self._value):
- if isinstance(item, LiteralExpression):
+ if isinstance(item, Column):
expr.literal.array.values.append(item.to_plan(session).literal)
else:
- expr.literal.array.values.append(
- LiteralExpression(item).to_plan(session).literal
- )
+ expr.literal.array.values.append(lit(item).to_plan(session).literal)
elif isinstance(self._value, tuple):
expr.literal.struct.SetInParent()
for item in list(self._value):
- if isinstance(item, LiteralExpression):
+ if isinstance(item, Column):
expr.literal.struct.fields.append(item.to_plan(session).literal)
else:
- expr.literal.struct.fields.append(
- LiteralExpression(item).to_plan(session).literal
- )
+ expr.literal.struct.fields.append(lit(item).to_plan(session).literal)
elif isinstance(self._value, dict):
expr.literal.map.SetInParent()
for key, value in dict(self._value).items():
pair = proto.Expression.Literal.Map.Pair()
- if isinstance(key, LiteralExpression):
+ if isinstance(key, Column):
pair.key.CopyFrom(key.to_plan(session).literal)
else:
- pair.key.CopyFrom(LiteralExpression(key).to_plan(session).literal)
- if isinstance(value, LiteralExpression):
+ pair.key.CopyFrom(lit(key).to_plan(session).literal)
+ if isinstance(value, Column):
pair.value.CopyFrom(value.to_plan(session).literal)
else:
- pair.value.CopyFrom(LiteralExpression(value).to_plan(session).literal)
+ pair.value.CopyFrom(lit(value).to_plan(session).literal)
expr.literal.map.pairs.append(pair)
else:
raise ValueError(f"Could not convert literal for type {type(self._value)}")
@@ -232,19 +228,22 @@ class LiteralExpression(Expression):
return f"Literal({self._value})"
-class Column(Expression):
+class ColumnReference(Expression):
"""Represents a column reference. There is no guarantee that this column
actually exists. In the context of this project, we refer by its name and
treat it as an unresolved attribute. Attributes that have the same fully
qualified name are identical"""
@classmethod
- def from_qualified_name(cls, name: str) -> "Column":
- return Column(name)
+ def from_qualified_name(cls, name: str) -> "ColumnReference":
+ return ColumnReference(name)
- def __init__(self, name: str) -> None:
+ def __init__(self, name: Union[str, "Column"]) -> None:
super().__init__()
- self._unparsed_identifier: str = name
+ if isinstance(name, str):
+ self._unparsed_identifier = name
+ else:
+ self._unparsed_identifier = name.name()
def name(self) -> str:
"""Returns the qualified name of the column reference."""
@@ -263,7 +262,7 @@ class Column(Expression):
return SortOrder(self, ascending=True)
def __str__(self) -> str:
- return f"Column({self._unparsed_identifier})"
+ return f"ColumnReference({self._unparsed_identifier})"
class SQLExpression(Expression):
@@ -283,11 +282,13 @@ class SQLExpression(Expression):
class SortOrder(Expression):
- def __init__(self, col: Column, ascending: bool = True, nullsLast: bool = True) -> None:
+ def __init__(
+ self, col: ColumnReference, ascending: bool = True, nullsLast: bool = True
+ ) -> None:
super().__init__()
self.ref = col
- self.ascending = ascending
- self.nullsLast = nullsLast
+ self._ascending = ascending
+ self._nullsLast = nullsLast
def __str__(self) -> str:
return str(self.ref) + " ASC" if self.ascending else " DESC"
@@ -295,12 +296,18 @@ class SortOrder(Expression):
def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
return self.ref.to_plan(session)
+ def ascending(self) -> bool:
+ return self._ascending
+
+ def nullsLast(self) -> bool:
+ return self._nullsLast
+
class ScalarFunctionExpression(Expression):
def __init__(
self,
op: str,
- *args: Expression,
+ *args: "Column",
) -> None:
super().__init__()
self._args = args
@@ -314,3 +321,64 @@ class ScalarFunctionExpression(Expression):
def __str__(self) -> str:
return f"({self._op} ({', '.join([str(x) for x in self._args])}))"
+
+
+class Column(object):
+ """
+ A column in a DataFrame. Column can refer to different things based on the
+ wrapped expression. Some common examples include attribute references, functions,
+ literals, etc.
+
+ .. versionadded:: 3.4.0
+ """
+
+ def __init__(self, expr: Expression) -> None:
+ self._expr = expr
+
+ __gt__ = _bin_op(">")
+ __lt__ = _bin_op("<")
+ __add__ = _bin_op("+")
+ __sub__ = _bin_op("-")
+ __mul__ = _bin_op("*")
+ __div__ = _bin_op("/")
+ __truediv__ = _bin_op("/")
+ __mod__ = _bin_op("%")
+ __radd__ = _bin_op("+", reverse=True)
+ __rsub__ = _bin_op("-", reverse=True)
+ __rmul__ = _bin_op("*", reverse=True)
+ __rdiv__ = _bin_op("/", reverse=True)
+ __rtruediv__ = _bin_op("/", reverse=True)
+ __pow__ = _bin_op("pow")
+ __rpow__ = _bin_op("pow", reverse=True)
+ __ge__ = _bin_op(">=")
+ __le__ = _bin_op("<=")
+ # __eq__ = _bin_op("==") # ignore [assignment]
+
+ def __eq__(self, other: Any) -> "Column": # type: ignore[override]
+ """Returns a binary expression with the current column as the left
+ side and the other expression as the right side.
+ """
+ from pyspark.sql.connect._typing import PrimitiveType
+ from pyspark.sql.connect.functions import lit
+
+ if isinstance(other, get_args(PrimitiveType)):
+ other = lit(other)
+ return scalar_function("==", self, other)
+
+ def to_plan(self, session: "SparkConnectClient") -> proto.Expression:
+ return self._expr.to_plan(session)
+
+ def alias(self, *alias: str, **kwargs: Any) -> "Column":
+ return Column(self._expr.alias(*alias, **kwargs))
+
+ def desc(self) -> "Column":
+ return Column(self._expr.desc())
+
+ def asc(self) -> "Column":
+ return Column(self._expr.asc())
+
+ def name(self) -> str:
+ return self._expr.name()
+
+ def __str__(self) -> str:
+ return self._expr.__str__()
diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py
index 51ea1dace6d..08ad9954672 100644
--- a/python/pyspark/sql/connect/dataframe.py
+++ b/python/pyspark/sql/connect/dataframe.py
@@ -33,20 +33,15 @@ import pandas
import pyspark.sql.connect.plan as plan
from pyspark.sql.connect.readwriter import DataFrameWriter
-from pyspark.sql.connect.column import (
- Column,
- Expression,
- LiteralExpression,
- SQLExpression,
- ScalarFunctionExpression,
-)
+from pyspark.sql.connect.column import Column, scalar_function, sql_expression
+from pyspark.sql.connect.functions import col, lit
from pyspark.sql.types import (
StructType,
Row,
)
if TYPE_CHECKING:
- from pyspark.sql.connect._typing import ColumnOrName, ExpressionOrString, LiteralType
+ from pyspark.sql.connect._typing import ColumnOrName, LiteralType
from pyspark.sql.connect.session import SparkSession
@@ -55,7 +50,7 @@ class GroupedData(object):
self._df = df
self._grouping_cols = [x if isinstance(x, Column) else df[x] for x in grouping_cols]
- def agg(self, measures: Sequence[Expression]) -> "DataFrame":
+ def agg(self, measures: Sequence[Column]) -> "DataFrame":
assert len(measures) > 0, "exprs should not be empty"
res = DataFrame.withPlan(
plan.Aggregate(
@@ -67,27 +62,25 @@ class GroupedData(object):
)
return res
- def _map_cols_to_expression(
- self, fun: str, col: Union[Expression, str]
- ) -> Sequence[Expression]:
+ def _map_cols_to_expression(self, fun: str, param: Union[Column, str]) -> Sequence[Column]:
return [
- ScalarFunctionExpression(fun, Column(col)) if isinstance(col, str) else col,
+ scalar_function(fun, col(param)) if isinstance(param, str) else param,
]
- def min(self, col: Union[Expression, str]) -> "DataFrame":
+ def min(self, col: Union[Column, str]) -> "DataFrame":
expr = self._map_cols_to_expression("min", col)
return self.agg(expr)
- def max(self, col: Union[Expression, str]) -> "DataFrame":
+ def max(self, col: Union[Column, str]) -> "DataFrame":
expr = self._map_cols_to_expression("max", col)
return self.agg(expr)
- def sum(self, col: Union[Expression, str]) -> "DataFrame":
+ def sum(self, col: Union[Column, str]) -> "DataFrame":
expr = self._map_cols_to_expression("sum", col)
return self.agg(expr)
def count(self) -> "DataFrame":
- return self.agg([ScalarFunctionExpression("count", LiteralExpression(1))])
+ return self.agg([scalar_function("count", lit(1))])
class DataFrame(object):
@@ -134,7 +127,7 @@ class DataFrame(object):
"""
return len(self.take(1)) == 0
- def select(self, *cols: "ExpressionOrString") -> "DataFrame":
+ def select(self, *cols: "ColumnOrName") -> "DataFrame":
return DataFrame.withPlan(plan.Project(self._plan, *cols), session=self._session)
def selectExpr(self, *expr: Union[str, List[str]]) -> "DataFrame":
@@ -154,23 +147,23 @@ class DataFrame(object):
expr = expr[0] # type: ignore[assignment]
for element in expr:
if isinstance(element, str):
- sql_expr.append(SQLExpression(element))
+ sql_expr.append(sql_expression(element))
else:
- sql_expr.extend([SQLExpression(e) for e in element])
+ sql_expr.extend([sql_expression(e) for e in element])
return DataFrame.withPlan(plan.Project(self._plan, *sql_expr), session=self._session)
- def agg(self, *exprs: Union[Expression, Dict[str, str]]) -> "DataFrame":
+ def agg(self, *exprs: Union[Column, Dict[str, str]]) -> "DataFrame":
if not exprs:
raise ValueError("Argument 'exprs' must not be empty")
if len(exprs) == 1 and isinstance(exprs[0], dict):
- measures = [ScalarFunctionExpression(f, Column(e)) for e, f in exprs[0].items()]
+ measures = [scalar_function(f, col(e)) for e, f in exprs[0].items()]
return self.groupBy().agg(measures)
else:
# other expressions
- assert all(isinstance(c, Expression) for c in exprs), "all exprs should be Expression"
- exprs = cast(Tuple[Expression, ...], exprs)
+ assert all(isinstance(c, Column) for c in exprs), "all exprs should be Expression"
+ exprs = cast(Tuple[Column, ...], exprs)
return self.groupBy().agg(exprs)
def alias(self, alias: str) -> "DataFrame":
@@ -224,7 +217,7 @@ class DataFrame(object):
int
Number of rows.
"""
- pdd = self.agg(ScalarFunctionExpression("count", LiteralExpression(1))).toPandas()
+ pdd = self.agg(scalar_function("count", lit(1))).toPandas()
return pdd.iloc[0, 0]
def crossJoin(self, other: "DataFrame") -> "DataFrame":
@@ -352,7 +345,7 @@ class DataFrame(object):
session=self._session,
)
- def filter(self, condition: Union[Expression, str]) -> "DataFrame":
+ def filter(self, condition: Union[Column, str]) -> "DataFrame":
"""Filters rows using the given condition.
:func:`where` is an alias for :func:`filter`.
@@ -371,7 +364,7 @@ class DataFrame(object):
Filtered DataFrame.
"""
if isinstance(condition, str):
- expr = cast(Expression, SQLExpression(condition))
+ expr = sql_expression(condition)
else:
expr = condition
return DataFrame.withPlan(plan.Filter(child=self._plan, filter=expr), session=self._session)
@@ -652,7 +645,7 @@ class DataFrame(object):
assert pdf is not None
return pdf["show_string"][0]
- def withColumns(self, colsMap: Dict[str, Expression]) -> "DataFrame":
+ def withColumns(self, colsMap: Dict[str, Column]) -> "DataFrame":
"""
Returns a new :class:`DataFrame` by adding multiple columns or replacing the
existing columns that have the same names.
@@ -680,7 +673,7 @@ class DataFrame(object):
session=self._session,
)
- def withColumn(self, colName: str, col: Expression) -> "DataFrame":
+ def withColumn(self, colName: str, col: Column) -> "DataFrame":
"""
Returns a new :class:`DataFrame` by adding a column or replacing the
existing column that has the same name.
@@ -702,7 +695,7 @@ class DataFrame(object):
:class:`DataFrame`
DataFrame with new or replaced column.
"""
- if not isinstance(col, Expression):
+ if not isinstance(col, Column):
raise TypeError("col should be Column")
return DataFrame.withPlan(
plan.WithColumns(self._plan, {colName: col}),
@@ -895,7 +888,7 @@ class DataFrame(object):
session=self._session,
)
- def where(self, condition: Union[Expression, str]) -> "DataFrame":
+ def where(self, condition: Union[Column, str]) -> "DataFrame":
return self.filter(condition)
@property
@@ -1129,9 +1122,9 @@ class DataFrame(object):
# Check for alias
alias = self._get_alias()
if alias is not None:
- return Column(alias)
+ return col(alias)
else:
- return Column(name)
+ return col(name)
def _print_plan(self) -> str:
if self._plan:
diff --git a/python/pyspark/sql/connect/function_builder.py b/python/pyspark/sql/connect/function_builder.py
index 8df5e56b452..b65348c6862 100644
--- a/python/pyspark/sql/connect/function_builder.py
+++ b/python/pyspark/sql/connect/function_builder.py
@@ -20,24 +20,20 @@ from typing import TYPE_CHECKING, Optional, Any, Iterable, Union
import pyspark.sql.connect.proto as proto
import pyspark.sql.types
-from pyspark.sql.connect.column import (
- Column,
- Expression,
- ScalarFunctionExpression,
-)
+from pyspark.sql.connect.column import Expression, ScalarFunctionExpression, Column
+from pyspark.sql.connect.functions import col
if TYPE_CHECKING:
from pyspark.sql.connect._typing import (
ColumnOrName,
- ExpressionOrString,
FunctionBuilderCallable,
UserDefinedFunctionCallable,
)
from pyspark.sql.connect.client import SparkConnectClient
-def _build(name: str, *args: "ExpressionOrString") -> ScalarFunctionExpression:
+def _build(name: str, *args: "ColumnOrName") -> ScalarFunctionExpression:
"""
Simple wrapper function that converts the arguments into the appropriate types.
Parameters
@@ -49,7 +45,7 @@ def _build(name: str, *args: "ExpressionOrString") -> ScalarFunctionExpression:
-------
:class:`ScalarFunctionExpression`
"""
- cols = [x if isinstance(x, Expression) else Column.from_qualified_name(x) for x in args]
+ cols = [x if isinstance(x, Column) else col(x) for x in args]
return ScalarFunctionExpression(name, *cols)
@@ -57,7 +53,7 @@ class FunctionBuilder:
"""This class is used to build arbitrary functions used in expressions"""
def __getattr__(self, name: str) -> "FunctionBuilderCallable":
- def _(*args: "ExpressionOrString") -> ScalarFunctionExpression:
+ def _(*args: "ColumnOrName") -> ScalarFunctionExpression:
return _build(name, *args)
_.__doc__ = f"""Function to apply {name}"""
@@ -107,8 +103,8 @@ class UserDefinedFunction(Expression):
def _create_udf(
function: Any, return_type: Union[str, pyspark.sql.types.DataType]
) -> "UserDefinedFunctionCallable":
- def wrapper(*cols: "ColumnOrName") -> UserDefinedFunction:
- return UserDefinedFunction(func=function, return_type=return_type, args=cols)
+ def wrapper(*cols: "ColumnOrName") -> "Column":
+ return Column(UserDefinedFunction(func=function, return_type=return_type, args=cols))
return wrapper
diff --git a/python/pyspark/sql/connect/functions.py b/python/pyspark/sql/connect/functions.py
index 00d0a56aedb..f7eb47eab59 100644
--- a/python/pyspark/sql/connect/functions.py
+++ b/python/pyspark/sql/connect/functions.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from pyspark.sql.connect.column import Column, LiteralExpression
+from pyspark.sql.connect.column import Column, LiteralExpression, ColumnReference
from typing import Any
@@ -22,8 +22,8 @@ from typing import Any
def col(x: str) -> Column:
- return Column(x)
+ return Column(ColumnReference(x))
-def lit(x: Any) -> LiteralExpression:
- return LiteralExpression(x)
+def lit(x: Any) -> Column:
+ return Column(LiteralExpression(x))
diff --git a/python/pyspark/sql/connect/plan.py b/python/pyspark/sql/connect/plan.py
index 3b9d9c7e604..ed4080e4230 100644
--- a/python/pyspark/sql/connect/plan.py
+++ b/python/pyspark/sql/connect/plan.py
@@ -28,15 +28,11 @@ from typing import (
import pandas
import pyarrow as pa
import pyspark.sql.connect.proto as proto
-from pyspark.sql.connect.column import (
- Column,
- Expression,
- SortOrder,
-)
+from pyspark.sql.connect.column import Column
if TYPE_CHECKING:
- from pyspark.sql.connect._typing import ColumnOrName, ExpressionOrString
+ from pyspark.sql.connect._typing import ColumnOrName
from pyspark.sql.connect.client import SparkConnectClient
@@ -258,7 +254,7 @@ class Project(LogicalPlan):
"""
- def __init__(self, child: Optional["LogicalPlan"], *columns: "ExpressionOrString") -> None:
+ def __init__(self, child: Optional["LogicalPlan"], *columns: "ColumnOrName") -> None:
super().__init__(child)
self._raw_columns = list(columns)
self.alias: Optional[str] = None
@@ -267,16 +263,16 @@ class Project(LogicalPlan):
def _verify_expressions(self) -> None:
"""Ensures that all input arguments are instances of Expression or String."""
for c in self._raw_columns:
- if not isinstance(c, (Expression, str)):
+ if not isinstance(c, (Column, str)):
raise InputValidationError(
- f"Only Expressions or String can be used for projections: '{c}'."
+ f"Only Column or String can be used for projections: '{c}'."
)
def plan(self, session: "SparkConnectClient") -> proto.Relation:
assert self._child is not None
proj_exprs = []
for c in self._raw_columns:
- if isinstance(c, Expression):
+ if isinstance(c, Column):
proj_exprs.append(c.to_plan(session))
elif c == "*":
exp = proto.Expression()
@@ -309,7 +305,7 @@ class Project(LogicalPlan):
class WithColumns(LogicalPlan):
"""Logical plan object for a withColumns operation."""
- def __init__(self, child: Optional["LogicalPlan"], cols_map: Mapping[str, Expression]) -> None:
+ def __init__(self, child: Optional["LogicalPlan"], cols_map: Mapping[str, Column]) -> None:
super().__init__(child)
self._cols_map = cols_map
@@ -341,7 +337,7 @@ class WithColumns(LogicalPlan):
class Filter(LogicalPlan):
- def __init__(self, child: Optional["LogicalPlan"], filter: Expression) -> None:
+ def __init__(self, child: Optional["LogicalPlan"], filter: Column) -> None:
super().__init__(child)
self.filter = filter
@@ -495,7 +491,7 @@ class Sort(LogicalPlan):
def __init__(
self,
child: Optional["LogicalPlan"],
- columns: List[Union[SortOrder, Column, str]],
+ columns: List[Union[Column, str]],
is_global: bool,
) -> None:
super().__init__(child)
@@ -503,19 +499,19 @@ class Sort(LogicalPlan):
self.is_global = is_global
def col_to_sort_field(
- self, col: Union[SortOrder, Column, str], session: "SparkConnectClient"
+ self, col: Union[Column, str], session: "SparkConnectClient"
) -> proto.Sort.SortField:
- if isinstance(col, SortOrder):
+ if isinstance(col, Column):
sf = proto.Sort.SortField()
- sf.expression.CopyFrom(col.ref.to_plan(session))
+ sf.expression.CopyFrom(col.to_plan(session))
sf.direction = (
proto.Sort.SortDirection.SORT_DIRECTION_ASCENDING
- if col.ascending
+ if col._expr.ascending
else proto.Sort.SortDirection.SORT_DIRECTION_DESCENDING
)
sf.nulls = (
proto.Sort.SortNulls.SORT_NULLS_FIRST
- if not col.nullsLast
+ if not col._expr.nullsLast
else proto.Sort.SortNulls.SORT_NULLS_LAST
)
return sf
@@ -653,13 +649,13 @@ class Aggregate(LogicalPlan):
self,
child: Optional["LogicalPlan"],
grouping_cols: List[Column],
- measures: Sequence[Expression],
+ measures: Sequence[Column],
) -> None:
super().__init__(child)
self.grouping_cols = grouping_cols
self.measures = measures
- def _convert_measure(self, m: Expression, session: "SparkConnectClient") -> proto.Expression:
+ def _convert_measure(self, m: Column, session: "SparkConnectClient") -> proto.Expression:
proto_expr = proto.Expression()
proto_expr.CopyFrom(m.to_plan(session))
return proto_expr
diff --git a/python/pyspark/sql/tests/connect/test_connect_column_expressions.py b/python/pyspark/sql/tests/connect/test_connect_column_expressions.py
index 9e2ff8290d6..6a6e0f84e21 100644
--- a/python/pyspark/sql/tests/connect/test_connect_column_expressions.py
+++ b/python/pyspark/sql/tests/connect/test_connect_column_expressions.py
@@ -26,7 +26,7 @@ from pyspark.testing.sqlutils import have_pandas, pandas_requirement_message
if have_pandas:
from pyspark.sql.connect.proto import Expression as ProtoExpression
import pyspark.sql.connect.plan as p
- import pyspark.sql.connect.column as col
+ from pyspark.sql.connect.column import Column
import pyspark.sql.connect.functions as fun
@@ -36,11 +36,11 @@ class SparkConnectColumnExpressionSuite(PlanOnlyTestFixture):
df = self.connect.with_plan(p.Read("table"))
c1 = df.col_name
- self.assertIsInstance(c1, col.Column)
+ self.assertIsInstance(c1, Column)
c2 = df["col_name"]
- self.assertIsInstance(c2, col.Column)
+ self.assertIsInstance(c2, Column)
c3 = fun.col("col_name")
- self.assertIsInstance(c3, col.Column)
+ self.assertIsInstance(c3, Column)
# All Protos should be identical
cp1 = c1.to_plan(None)
@@ -182,7 +182,7 @@ class SparkConnectColumnExpressionSuite(PlanOnlyTestFixture):
def test_column_alias(self) -> None:
# SPARK-40809: Support for Column Aliases
col0 = fun.col("a").alias("martin")
- self.assertEqual("Alias(Column(a), (martin))", str(col0))
+ self.assertEqual("Alias(ColumnReference(a), (martin))", str(col0))
col0 = fun.col("a").alias("martin", metadata={"pii": True})
plan = col0.to_plan(self.session.client)
diff --git a/python/pyspark/sql/tests/connect/test_connect_plan_only.py b/python/pyspark/sql/tests/connect/test_connect_plan_only.py
index e7a14c4bb05..fe005cd6365 100644
--- a/python/pyspark/sql/tests/connect/test_connect_plan_only.py
+++ b/python/pyspark/sql/tests/connect/test_connect_plan_only.py
@@ -23,6 +23,7 @@ from pyspark.testing.sqlutils import have_pandas, pandas_requirement_message
if have_pandas:
import pyspark.sql.connect.proto as proto
+ from pyspark.sql.connect.column import Column
from pyspark.sql.connect.readwriter import DataFrameReader
from pyspark.sql.connect.function_builder import UserDefinedFunction, udf
from pyspark.sql.types import StringType
@@ -278,7 +279,8 @@ class SparkConnectTestsPlanOnly(PlanOnlyTestFixture):
u = udf(lambda x: "Martin", StringType())
self.assertIsNotNone(u)
expr = u("ThisCol", "ThatCol", "OtherCol")
- self.assertTrue(isinstance(expr, UserDefinedFunction))
+ self.assertTrue(isinstance(expr, Column))
+ self.assertTrue(isinstance(cast(Column, expr)._expr, UserDefinedFunction))
u_plan = expr.to_plan(self.connect)
self.assertIsNotNone(u_plan)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org