xinrong-databricks commented on a change in pull request #33373:
URL: https://github.com/apache/spark/pull/33373#discussion_r671584402
##########
File path: python/pyspark/pandas/data_type_ops/categorical_ops.py
##########
@@ -71,28 +72,58 @@ def astype(self, index_ops: IndexOpsLike, dtype: Union[str,
type, Dtype]) -> Ind
scol = map_scol.getItem(index_ops.spark.column)
return index_ops._with_new_scol(scol).astype(dtype)
+ def eq(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
+ return _compare(left, right, Column.__eq__,
is_equality_comparison=True)
+
+ def ne(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
+ return _compare(left, right, Column.__ne__,
is_equality_comparison=True)
+
def lt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
- _non_equality_comparison_input_check(left, right)
- return column_op(Column.__lt__)(left, right)
+ return _compare(left, right, Column.__lt__)
def le(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
- _non_equality_comparison_input_check(left, right)
- return column_op(Column.__le__)(left, right)
+ return _compare(left, right, Column.__le__)
def gt(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
- _non_equality_comparison_input_check(left, right)
- return column_op(Column.__gt__)(left, right)
+ return _compare(left, right, Column.__gt__)
def ge(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
- _non_equality_comparison_input_check(left, right)
- return column_op(Column.__ge__)(left, right)
+ return _compare(left, right, Column.__ge__)
-def _non_equality_comparison_input_check(left: IndexOpsLike, right: Any) ->
None:
- if not cast(CategoricalDtype, left.dtype).ordered:
- raise TypeError("Unordered Categoricals can only compare equality or
not.")
+def _compare(
+ left: IndexOpsLike,
+ right: Any,
+ f: Callable[..., Column],
+ *,
+ is_equality_comparison: bool = False
+) -> SeriesOrIndex:
+ """
+ Compare a Categorical operand `left` to `right` with the given Spark
Column function.
+
+ Parameters
+ ----------
+ left: A Categorical operand
+ right: The other operand to compare with
+ f : The Spark Column function to apply
+ is_equality_comparison: True if it is equality comparison, ie. == or !=.
False by default.
+
+ Returns
+ -------
+ SeriesOrIndex
+ """
if isinstance(right, IndexOpsMixin) and isinstance(right.dtype,
CategoricalDtype):
+ if not is_equality_comparison:
+ if not cast(CategoricalDtype, left.dtype).ordered:
+ raise TypeError("Unordered Categoricals can only compare
equality or not.")
if hash(left.dtype) != hash(right.dtype):
Review comment:
I referenced
https://github.com/pandas-dev/pandas/blob/master/pandas/core/arrays/categorical.py#L145.
`hash` can tell the same dtype, same categories, and same ordered.
--
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]
