|
43 | 43 | _is_boolean_type,
|
44 | 44 | )
|
45 | 45 | from pyspark.pandas.typedef.typehints import extension_dtypes, pandas_on_spark_type
|
| 46 | +from pyspark.pandas.utils import is_ansi_mode_enabled |
46 | 47 | from pyspark.sql import functions as F, Column as PySparkColumn
|
47 | 48 | from pyspark.sql.types import (
|
48 | 49 | BooleanType,
|
@@ -247,14 +248,23 @@ def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
|
247 | 248 | _sanitize_list_like(right)
|
248 | 249 | if not is_valid_operand_for_numeric_arithmetic(right):
|
249 | 250 | raise TypeError("True division can not be applied to given types.")
|
| 251 | + spark_session = left._internal.spark_frame.sparkSession |
| 252 | + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) |
250 | 253 |
|
251 | 254 | def truediv(left: PySparkColumn, right: Any) -> PySparkColumn:
|
252 |
| - return F.when( |
253 |
| - F.lit(right != 0) | F.lit(right).isNull(), |
254 |
| - left.__div__(right), |
255 |
| - ).otherwise(F.lit(np.inf).__div__(left)) |
| 255 | + if is_ansi_mode_enabled(spark_session): |
| 256 | + return F.when( |
| 257 | + F.lit(right == 0), |
| 258 | + F.when(left < 0, F.lit(float("-inf"))) |
| 259 | + .when(left > 0, F.lit(float("inf"))) |
| 260 | + .otherwise(F.lit(np.nan)), |
| 261 | + ).otherwise(left / right) |
| 262 | + else: |
| 263 | + return F.when( |
| 264 | + F.lit(right != 0) | F.lit(right).isNull(), |
| 265 | + left.__div__(right), |
| 266 | + ).otherwise(F.lit(np.inf).__div__(left)) |
256 | 267 |
|
257 |
| - right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) |
258 | 268 | return numpy_column_op(truediv)(left, right)
|
259 | 269 |
|
260 | 270 | def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
|
@@ -332,18 +342,27 @@ def truediv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
|
332 | 342 | _sanitize_list_like(right)
|
333 | 343 | if not is_valid_operand_for_numeric_arithmetic(right):
|
334 | 344 | raise TypeError("True division can not be applied to given types.")
|
| 345 | + spark_session = left._internal.spark_frame.sparkSession |
| 346 | + right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) |
335 | 347 |
|
336 | 348 | def truediv(left: PySparkColumn, right: Any) -> PySparkColumn:
|
337 |
| - return F.when( |
338 |
| - F.lit(right != 0) | F.lit(right).isNull(), |
339 |
| - left.__div__(right), |
340 |
| - ).otherwise( |
341 |
| - F.when(F.lit(left == np.inf) | F.lit(left == -np.inf), left).otherwise( |
342 |
| - F.lit(np.inf).__div__(left) |
| 349 | + if is_ansi_mode_enabled(spark_session): |
| 350 | + return F.when( |
| 351 | + F.lit(right == 0), |
| 352 | + F.when(left < 0, F.lit(float("-inf"))) |
| 353 | + .when(left > 0, F.lit(float("inf"))) |
| 354 | + .otherwise(F.lit(np.nan)), |
| 355 | + ).otherwise(left / right) |
| 356 | + else: |
| 357 | + return F.when( |
| 358 | + F.lit(right != 0) | F.lit(right).isNull(), |
| 359 | + left.__div__(right), |
| 360 | + ).otherwise( |
| 361 | + F.when(F.lit(left == np.inf) | F.lit(left == -np.inf), left).otherwise( |
| 362 | + F.lit(np.inf).__div__(left) |
| 363 | + ) |
343 | 364 | )
|
344 |
| - ) |
345 | 365 |
|
346 |
| - right = transform_boolean_operand_to_numeric(right, spark_type=left.spark.data_type) |
347 | 366 | return numpy_column_op(truediv)(left, right)
|
348 | 367 |
|
349 | 368 | def floordiv(self, left: IndexOpsLike, right: Any) -> SeriesOrIndex:
|
|
0 commit comments