zhengruifeng commented on code in PR #38956:
URL: https://github.com/apache/spark/pull/38956#discussion_r1041717834
##########
python/pyspark/sql/connect/column.py:
##########
@@ -129,6 +140,53 @@ def name(self) -> str:
...
+class CaseWhen(Expression):
+ def __init__(
+ self, branches: Sequence[Tuple[Expression, Expression]], else_value:
Optional[Expression]
+ ):
+
+ assert isinstance(branches, list)
+ for branch in branches:
+ assert (
+ isinstance(branch, tuple)
+ and len(branch) == 2
+ and all(isinstance(expr, Expression) for expr in branch)
+ )
+ self._branches = branches
+
+ if else_value is not None:
+ assert isinstance(else_value, Expression)
+
+ self._else_value = else_value
+
+ def to_plan(self, session: "SparkConnectClient") -> "proto.Expression":
+ # fun = proto.Expression()
+ # fun.unresolved_function.function_name = "when"
+ # for condition, value in self._branches:
+ #
fun.unresolved_function.arguments.extend(condition.to_plan(session))
+ # fun.unresolved_function.arguments.extend(value.to_plan(session))
+ # if self._else_value is not None:
+ #
fun.unresolved_function.arguments.extend(self._else_value.to_plan(session))
+ # return fun
+
+ args: Sequence[Expression] = []
+ for condition, value in self._branches:
+ args.append(condition)
+ args.append(value)
+
+ if self._else_value is not None:
+ args.append(self._else_value)
+
+ unresolved_function = UnresolvedFunction(name="when", args=args)
+
+ return unresolved_function.to_plan(session)
Review Comment:
let me add another constructor
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
