zhengruifeng commented on code in PR #37235:
URL: https://github.com/apache/spark/pull/37235#discussion_r926474342
##########
python/pyspark/pandas/indexes/base.py:
##########
@@ -2666,6 +2666,93 @@ def inferred_type(self) -> str:
"""
return lib.infer_dtype([self.to_series().head(1).item()])
+ def putmask(
+ self,
+ mask: "Index",
+ value: Any
+ ) -> "Index":
+ """
+ Return a new Index of the values set with the mask.
+
+ Returns
+ -------
+ Index
+
+ See Also
+ --------
+ pandas.index.putmask : Changes elements of an array
+ based on conditional and input values.
+ """
+ # validate mask
+ mask = np.asarray(mask.tolist(), dtype=bool) # type:
ignore[assignment]
+ if mask.shape != self.values.shape:
+ raise ValueError("cond and data must be the same size")
+
+ noop = not mask.any()
+
+ if noop:
+ return self.copy()
+
+ # convert the insert value, check whether it can be inserted.
+ converted_other = value
+ if value is not np.nan:
+ try:
+ converted_other = np.dtype(self.dtype).type(value) # type:
ignore[arg-type]
+ except ValueError:
+ raise ValueError("The inserted value should be in the same
dtype.")
+
+ # reconstruct a new index
+ values = self.values.copy()
+ new_values = []
+ if isinstance(values, np.ndarray):
+ for val, is_need in zip(values, mask):
+ if is_need:
+ new_values.append(converted_other)
+ else:
+ new_values.append(val)
+ else:
+ new_values.append(converted_other)
+
+ return Index(new_values)
+
+ def where(
+ self,
+ cond: "Index",
+ other: Any = np.nan
+ ) -> "Index":
+ """
+ Replace values where the condition is False.
+
Review Comment:
ditto
--
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]
