Skip to content

Commit

Permalink
[SPARK-45577][PYTHON] Fix UserDefinedPythonTableFunctionAnalyzeRunner…
Browse files Browse the repository at this point in the history
… to pass folded values from named arguments

### What changes were proposed in this pull request?

Fix `UserDefinedPythonTableFunctionAnalyzeRunner` to pass folded values from named arguments.

### Why are the changes needed?

In the analysis in Python for Python UDTF, the folded values are not passed from named arguments.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Modified the related tests.

### Was this patch authored or co-authored using generative AI tooling?

No.

Closes apache#43407 from ueshin/issues/SPARK-45577/analyze.

Authored-by: Takuya UESHIN <[email protected]>
Signed-off-by: Takuya UESHIN <[email protected]>
  • Loading branch information
ueshin committed Oct 17, 2023
1 parent b10fea9 commit 4a0ed9c
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 6 deletions.
17 changes: 15 additions & 2 deletions python/pyspark/sql/tests/test_udtf.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
import tempfile
import unittest
from dataclasses import dataclass
from typing import Iterator
from typing import Iterator, Optional

from py4j.protocol import Py4JJavaError

Expand Down Expand Up @@ -1967,6 +1967,12 @@ def test_udtf_with_analyze_kwargs(self):
class TestUDTF:
@staticmethod
def analyze(**kwargs: AnalyzeArgument) -> AnalyzeResult:
assert isinstance(kwargs["a"].data_type, IntegerType)
assert kwargs["a"].value == 10
assert not kwargs["a"].is_table
assert isinstance(kwargs["b"].data_type, StringType)
assert kwargs["b"].value == "x"
assert not kwargs["b"].is_table
return AnalyzeResult(
StructType(
[StructField(key, arg.data_type) for key, arg in sorted(kwargs.items())]
Expand Down Expand Up @@ -2021,7 +2027,14 @@ def test_udtf_with_named_arguments_and_defaults(self):
@udtf
class TestUDTF:
@staticmethod
def analyze(a, b=None):
def analyze(a: AnalyzeArgument, b: Optional[AnalyzeArgument] = None):
assert isinstance(a.data_type, IntegerType)
assert a.value == 10
assert not a.is_table
if b is not None:
assert isinstance(b.data_type, StringType)
assert b.value == "z"
assert not b.is_table
schema = StructType().add("a", a.data_type)
if b is None:
return AnalyzeResult(schema.add("b", IntegerType()))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -198,17 +198,21 @@ class UserDefinedPythonTableFunctionAnalyzeRunner(
dataOut.writeInt(exprs.length)
exprs.zip(tableArgs).foreach { case (expr, is_table) =>
PythonWorkerUtils.writeUTF(expr.dataType.json, dataOut)
if (expr.foldable) {
val (key, value) = expr match {
case NamedArgumentExpression(k, v) => (Some(k), v)
case _ => (None, expr)
}
if (value.foldable) {
dataOut.writeBoolean(true)
val obj = pickler.dumps(EvaluatePython.toJava(expr.eval(), expr.dataType))
val obj = pickler.dumps(EvaluatePython.toJava(value.eval(), value.dataType))
PythonWorkerUtils.writeBytes(obj, dataOut)
} else {
dataOut.writeBoolean(false)
}
dataOut.writeBoolean(is_table)
// If the expr is NamedArgumentExpression, send its name.
expr match {
case NamedArgumentExpression(key, _) =>
key match {
case Some(key) =>
dataOut.writeBoolean(true)
PythonWorkerUtils.writeUTF(key, dataOut)
case _ =>
Expand Down

0 comments on commit 4a0ed9c

Please sign in to comment.