xinrong-databricks commented on a change in pull request #32611:
URL: https://github.com/apache/spark/pull/32611#discussion_r637299344
##########
File path: python/pyspark/pandas/data_type_ops/boolean_ops.py
##########
@@ -26,3 +41,183 @@ class BooleanOps(DataTypeOps):
@property
def pretty_name(self) -> str:
return 'booleans'
+
+ def add(self, left, right):
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return column_op(Column.__add__)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return column_op(Column.__add__)(left, right)
+ else:
+ raise TypeError(
+ "Addition can not be applied to %s and the given type." %
self.pretty_name)
+
+ def sub(self, left, right):
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return column_op(Column.__sub__)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return column_op(Column.__sub__)(left, right)
+ else:
+ raise TypeError(
+ "Subtraction can not be applied to %s and the given type." %
self.pretty_name)
+
+ def mul(self, left, right):
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return column_op(Column.__mul__)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return column_op(Column.__mul__)(left, right)
+ else:
+ raise TypeError(
+ "Multiplication can not be applied to %s and the given type."
% self.pretty_name)
+
+ def truediv(self, left, right):
+ def truediv(left, right):
+ return F.when(F.lit(right != 0) | F.lit(right).isNull(),
left.__div__(right)).otherwise(
+ F.when(F.lit(left == np.inf) | F.lit(left == -np.inf),
left).otherwise(
+ F.lit(np.inf).__div__(left)
+ )
+ )
+
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return numpy_column_op(truediv)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return numpy_column_op(truediv)(left, right)
+ else:
+ raise TypeError(
+ "True division can not be applied to %s and the given type." %
self.pretty_name)
+
+ def floordiv(self, left, right):
+ def floordiv(left, right):
+ return F.when(F.lit(right is np.nan), np.nan).otherwise(
+ F.when(
+ F.lit(right != 0) | F.lit(right).isNull(),
F.floor(left.__div__(right))
+ ).otherwise(
+ F.when(F.lit(left == np.inf) | F.lit(left == -np.inf),
left).otherwise(
+ F.lit(np.inf).__div__(left)
+ )
+ )
+ )
+
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return numpy_column_op(floordiv)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return numpy_column_op(floordiv)(left, right)
+ else:
+ raise TypeError(
+ "Floor division can not be applied to %s and the given type."
% self.pretty_name)
+
+ def mod(self, left, right):
+ def mod(left, right):
+ return ((left % right) + right) % right
+
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return column_op(mod)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return column_op(mod)(left, right)
+ else:
+ raise TypeError(
+ "Modulo can not be applied to %s and the given type." %
self.pretty_name)
+
+ def pow(self, left, right):
+ def pow_func(left, right):
+ return F.when(left == 1, left).otherwise(Column.__pow__(left,
right))
+
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return column_op(pow_func)(left, right)
+ elif isinstance(right, IndexOpsMixin) and is_numeric_index_ops(right):
+ left = transform_boolean_operand_to_numeric(left,
right.spark.data_type)
+ return column_op(pow_func)(left, right)
+
+ else:
+ raise TypeError(
+ "Exponentiation can not be applied to %s and the given type."
% self.pretty_name)
+
+ def radd(self, left, right=None):
+ if isinstance(right, numbers.Number):
+ left = left.spark.apply(lambda scol:
scol.cast(as_spark_type(type(right))))
+ return column_op(Column.__radd__)(left, right)
+ else:
+ raise TypeError(
+ "Addition can not be applied to %s and the given type." %
self.pretty_name)
+
+ def rsub(self, left, right=None):
Review comment:
`=None`s are removed.
--
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.
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
