Skip to content

Commit 61a79ca

Browse files
committed
[SPARK-52622][PS] Avoid CAST_INVALID_INPUT of DataFrame.melt in ANSI mode
### What changes were proposed in this pull request? Avoid CAST_INVALID_INPUT of `DataFrame.melt` in ANSI mode ### Why are the changes needed? Ensure pandas on Spark works well with ANSI mode on. Part of https://issues.apache.org/jira/browse/SPARK-52556. ### Does this PR introduce _any_ user-facing change? Yes, `DataFrame.melt` work well with ANSI. ```py >>> import pandas as pd >>> import numpy as np >>> >>> ps.set_option("compute.fail_on_ansi_mode", False) >>> ps.set_option("compute.ansi_mode_support", True) >>> df = ps.DataFrame({'A': {0: 'a', 1: 'b', 2: 'c'}, ... 'B': {0: 1, 1: 3, 2: 5}, ... 'C': {0: 2, 1: 4, 2: 6}}, ... columns=['A', 'B', 'C']) ``` BEFORE ```py >>> ps.melt(df) ... pyspark.errors.exceptions.captured.NumberFormatException: [CAST_INVALID_INPUT] The value 'a' of the type "STRING" cannot be cast to "BIGINT" because it is malformed. Correct the value as per the syntax, or change its target type. Use `try_cast` to tolerate malformed input and return NULL instead. SQLSTATE: 22018 ``` AFTER ```py >>> ps.melt(df) variable value 0 A a 1 B 1 2 C 2 3 A b 4 B 3 5 C 4 6 A c 7 B 5 8 C 6 ``` ### How was this patch tested? Unit tests. Commands below pass ``` SPARK_ANSI_SQL_MODE=true ./python/run-tests --python-executables=python3.11 --testnames "pyspark.pandas.tests.computation.test_melt FrameMeltTests.test_melt" SPARK_ANSI_SQL_MODE=false ./python/run-tests --python-executables=python3.11 --testnames "pyspark.pandas.tests.computation.test_melt FrameMeltTests.test_melt" ``` ### Was this patch authored or co-authored using generative AI tooling? No Closes #51326 from xinrong-meng/melt. Authored-by: Xinrong Meng <xinrong@apache.org> Signed-off-by: Xinrong Meng <xinrong@apache.org>
1 parent dc8fba6 commit 61a79ca

1 file changed

Lines changed: 19 additions & 1 deletion

File tree

python/pyspark/pandas/frame.py

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10617,12 +10617,30 @@ def melt(
1061710617
else:
1061810618
var_name = [var_name] # type: ignore[list-item]
1061910619

10620+
value_col_types = [
10621+
self._internal.spark_column_for(label).expr.dataType for label in value_vars
10622+
]
10623+
# If any value column is of StringType, cast all value columns to StringType to avoid
10624+
# ANSI mode errors during explode - mixing strings and integers.
10625+
string_cast_required_type = (
10626+
StringType() if any(isinstance(t, StringType) for t in value_col_types) else None
10627+
)
10628+
use_cast = is_ansi_mode_enabled(self._internal.spark_frame.sparkSession)
10629+
1062010630
pairs = F.explode(
1062110631
F.array(
1062210632
*[
1062310633
F.struct(
1062410634
*[F.lit(c).alias(name) for c, name in zip(label, var_name)],
10625-
*[self._internal.spark_column_for(label).alias(value_name)],
10635+
*[
10636+
(
10637+
self._internal.spark_column_for(label).cast(
10638+
string_cast_required_type
10639+
)
10640+
if use_cast and string_cast_required_type is not None
10641+
else self._internal.spark_column_for(label)
10642+
).alias(value_name)
10643+
],
1062610644
)
1062710645
for label in column_labels
1062810646
if label in value_vars

0 commit comments

Comments
 (0)