Skip to content

[SPARK-52249][PS] Enable divide-by-zero for numeric truediv with ANSI enabled #50972

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 10 commits into from
Closed

[SPARK-52249][PS] Enable divide-by-zero for numeric truediv with ANSI enabled #50972

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f9bc4e0
init
xinrong-meng May 19, 2025
8d85e3e
fix
xinrong-meng May 21, 2025
a1685d4
test
xinrong-meng May 21, 2025
329fdd2
enable test
xinrong-meng May 21, 2025
4764db4
able to switch
xinrong-meng May 22, 2025
3449164
fix
xinrong-meng May 22, 2025
720892c
fix
xinrong-meng May 23, 2025
b1ef242
dup test
xinrong-meng May 27, 2025
802f601
is_ansi_mode_enabled
xinrong-meng May 28, 2025
5ddc6c7
get sparksession from indexop
xinrong-meng May 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 32 additions & 13 deletions python/pyspark/pandas/data_type_ops/num_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
_is_boolean_type,
)
from pyspark.pandas.typedef.typehints import extension_dtypes, pandas_on_spark_type
from pyspark.pandas.utils import is_ansi_mode_enabled
from pyspark.sql import functions as F, Column as PySparkColumn
from pyspark.sql.types import (
BooleanType,
Expand Down Expand Up @@ -247,14 +248,23 @@ def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
_sanitize_list_like(right)
if not is_valid_operand_for_numeric_arithmetic(right):
raise TypeError("True division can not be applied to given types.")
spark_session = left._internal.spark_frame.sparkSession
right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)

def truediv(left: PySparkColumn, right: Any) -> PySparkColumn:
return F.when(
F.lit(right != 0) | F.lit(right).isNull(),
left.__div__(right),
).otherwise(F.lit(np.inf).__div__(left))
if is_ansi_mode_enabled(spark_session):
return F.when(
F.lit(right == 0),
F.when(left < 0, F.lit(float("-inf")))
.when(left > 0, F.lit(float("inf")))
.otherwise(F.lit(np.nan)),
).otherwise(left / right)
else:
return F.when(
F.lit(right != 0) | F.lit(right).isNull(),
left.__div__(right),
).otherwise(F.lit(np.inf).__div__(left))

right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
return numpy_column_op(truediv)(left, right)

def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
Expand Down Expand Up @@ -332,18 +342,27 @@ def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
_sanitize_list_like(right)
if not is_valid_operand_for_numeric_arithmetic(right):
raise TypeError("True division can not be applied to given types.")
spark_session = left._internal.spark_frame.sparkSession
right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)

def truediv(left: PySparkColumn, right: Any) -> PySparkColumn:
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 is_ansi_mode_enabled(spark_session):
return F.when(
F.lit(right == 0),
F.when(left < 0, F.lit(float("-inf")))
.when(left > 0, F.lit(float("inf")))
.otherwise(F.lit(np.nan)),
).otherwise(left / right)
else:
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)
)
)
)

right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type)
return numpy_column_op(truediv)(left, right)

def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
Expand Down
1 change: 0 additions & 1 deletion python/pyspark/pandas/tests/computation/test_binary_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,6 @@ def test_binary_operator_sub(self):
psdf = ps.DataFrame({"a": ["x"], "b": ["y"]})
self.assertRaisesRegex(TypeError, ks_err_msg, lambda: psdf["a"] - psdf["b"])

@unittest.skipIf(is_ansi_mode_test, ansi_mode_not_supported_message)
Copy link
Member

@ueshin ueshin May 28, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you also try to find where we can remove skipping caused by the division error?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can I follow up with https://issues.apache.org/jira/browse/SPARK-52349 on that if you don't mind, in order to unblock the other pr?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, that's fine. 👍

def test_divide_by_zero_behavior(self):
# float / float
# np.float32
Expand Down
7 changes: 7 additions & 0 deletions python/pyspark/pandas/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1070,6 +1070,13 @@ def xor(df1: PySparkDataFrame, df2: PySparkDataFrame) -> PySparkDataFrame:
)


def is_ansi_mode_enabled(spark: SparkSession) -> bool:
return (
ps.get_option("compute.ansi_mode_support", spark_session=spark)
and spark.conf.get("spark.sql.ansi.enabled") == "true"
)


def _test() -> None:
import os
import doctest
Expand Down