You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "ueshin (via GitHub)" <gi...@apache.org> on 2023/07/11 23:34:04 UTC
[GitHub] [spark] ueshin opened a new pull request, #41948: [SPARK-44380][PYTHON] Support for UDTF to analyze in Python
ueshin opened a new pull request, #41948:
URL: https://github.com/apache/spark/pull/41948
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'core/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as the documentation fix.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267104749
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +138,123 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ *
+ * When the Python UDTF is defined without a static return type,
+ * the analyze will call this while resolving table-valued functions.
+ *
+ * This expects the Python UDTF to have `analyze` static method that take arguments:
+ *
+ * - The number and order of arguments are the same as the UDTF inputs
+ * - Each argument is an `AnalyzeArgument`, containing:
+ * - data_type: DataType
+ * - value: Any: if the argument is foldable; otherwise None
+ * - is_table: bool: True if the argument is TABLE
+ *
+ * and that return an `AnalyzeResult`.
+ *
+ * It serializes/deserializes the data types via JSON,
+ * and the values for the case the argument is foldable are pickled.
+ *
+ * `AnalysisException` with the error class "TABLE_VALUED_FUNCTION_FAILED_TO_ANALYZE_IN_PYTHON"
+ * will be thrown when an exception is raised in Python.
+ */
+ def analyzeInPython(
+ tableArgs: Seq[Boolean])(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
Review Comment:
That's a good point. Let's discuss it in a separate PR.
`PYSPARK_DRIVER_PYTHON` is also used to run PySpark on Jupyter, so we might want to add another variable like `driverPythonExec` getting from the `sys.executable` or something.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] dtenedor commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "dtenedor (via GitHub)" <gi...@apache.org>.
dtenedor commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264223020
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
+ class TestUDTF:
+ def analyze(self) -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_returning_non_struct(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze():
+ return StringType()
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException,
+ "Output of `analyze` static method of Python UDTFs expects a StructType "
+ "but got: <class 'pyspark.sql.types.StringType'>",
+ ):
+ func().collect()
+
+ def test_udtf_with_analyze_taking_wrong_number_of_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a):
+ yield a, a + 1
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException, r"analyze\(\) missing 1 required positional argument: 'b'"
Review Comment:
OK. We can also add other enforcement in the SQL analyzer that the number of provided arguments was equal to the required number.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] dtenedor commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "dtenedor (via GitHub)" <gi...@apache.org>.
dtenedor commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1268552267
##########
python/pyspark/sql/udtf.py:
##########
@@ -153,6 +175,19 @@ def _validate_udtf_handler(cls: Any) -> None:
error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": cls.__name__}
)
+ has_analyze_staticmethod = hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
Review Comment:
I would recommend to return an error in this case if it's easy to do so, in order to prevent confusion. This would be in addition to the existing error that @ueshin pointed out that we are already returning:
> It IS checking if the UDTF class has an analyze method and throw an error if it is not static, and the error message will be like:
```
The UDTF '<name>' is invalid. It does not specify its return type or implement the required 'analyze' static method. Please specify the return type or implement the 'analyze' static method in '<name>' and try again.
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] dtenedor commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "dtenedor (via GitHub)" <gi...@apache.org>.
dtenedor commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263116521
##########
python/pyspark/errors/error_classes.py:
##########
@@ -277,6 +277,11 @@
"The UDTF '<name>' is invalid. It does not implement the required 'eval' method. Please implement the 'eval' method in '<name>' and try again."
]
},
+ "INVALID_UDTF_RETURN_TYPE" : {
+ "message" : [
+ "The UDTF '<name>' is invalid. It does not specify its return type or implement the required 'analysis' static method. Please specify the return type or implement the 'analyze' static method in '<name>' and try again."
Review Comment:
```suggestion
"The UDTF '<name>' is invalid. It does not specify its return type or implement the required 'analyze' static method. Please specify the return type or implement the 'analyze' static method in '<name>' and try again."
```
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
Review Comment:
Can you also please add a test for each of these cases using SQL as well (via PySpark), to cover the end-to-end query processing steps?
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
+ class TestUDTF:
+ def analyze(self) -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_returning_non_struct(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze():
+ return StringType()
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException,
+ "Output of `analyze` static method of Python UDTFs expects a StructType "
+ "but got: <class 'pyspark.sql.types.StringType'>",
+ ):
+ func().collect()
+
+ def test_udtf_with_analyze_taking_wrong_number_of_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a):
+ yield a, a + 1
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException, r"analyze\(\) missing 1 required positional argument: 'b'"
Review Comment:
Can you update this error message to say something like "The table valued function <functionName>" instead of "analyze()" to help the user understand what to do?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
Review Comment:
```suggestion
val pickler = new Pickler(useMemo = true,
```
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ false)
+
+ val (worker: Socket, _) =
+ env.createPythonWorker(pythonExec, workerModule, envVars.asScala.toMap)
+ val releasedOrClosed = new AtomicBoolean(false)
+ try {
+ val dataOut =
+ new DataOutputStream(new BufferedOutputStream(worker.getOutputStream, bufferSize))
+ val dataIn = new DataInputStream(new BufferedInputStream(worker.getInputStream, bufferSize))
+
+ // Python version of driver
+ PythonRDD.writeUTF(pythonVer, dataOut)
+
+ // Send Python UDTF
+ dataOut.writeInt(func.command.length)
+ dataOut.write(func.command.toArray)
+
+ // Send arguments
+ dataOut.writeInt(exprs.length)
+ exprs.foreach { expr =>
+ PythonRDD.writeUTF(expr.dataType.json, dataOut)
+ if (expr.foldable) {
+ dataOut.writeBoolean(true)
+ val obj = pickler.dumps(EvaluatePython.toJava(expr.eval(), expr.dataType))
+ dataOut.writeInt(obj.length)
+ dataOut.write(obj)
+ } else {
+ dataOut.writeBoolean(false)
+ }
+ dataOut.writeBoolean(expr.isInstanceOf[FunctionTableSubqueryArgumentExpression])
+ }
+
+ dataOut.writeInt(SpecialLengths.END_OF_STREAM)
+ dataOut.flush()
+
+ // Receive the schema
+ val schema = dataIn.readInt() match {
+ case length if length >= 0 =>
+ val obj = new Array[Byte](length)
+ dataIn.readFully(obj)
+ DataType.fromJson(new String(obj, StandardCharsets.UTF_8)).asInstanceOf[StructType]
+
+ case SpecialLengths.PYTHON_EXCEPTION_THROWN =>
+ val exLength = dataIn.readInt()
+ val obj = new Array[Byte](exLength)
+ dataIn.readFully(obj)
+ val msg = new String(obj, StandardCharsets.UTF_8)
+ env.destroyPythonWorker(pythonExec, workerModule, envVars.asScala.toMap, worker)
+ throw new AnalysisException(msg)
Review Comment:
should we make a new error class for this?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ false)
Review Comment:
```suggestion
valueCompare = false)
```
##########
python/pyspark/sql/functions.py:
##########
@@ -15524,9 +15524,10 @@ def udtf(
----------
cls : class
the Python user-defined table function handler class.
- returnType : :class:`pyspark.sql.types.StructType` or str
+ returnType : :class:`pyspark.sql.types.StructType` or str, optional
the return type of the user-defined table function. The value can be either a
:class:`pyspark.sql.types.StructType` object or a DDL-formatted struct type string.
+ If None, the handler class must provide `analyze` static method.
Review Comment:
Can you please port the description of how the 'analyze' method works from the PR description to a new action right above "Notes" on L15580? This will be useful to have a place for documentation for how this works in the code.
##########
python/pyspark/sql/worker/analyze_udtf.py:
##########
@@ -0,0 +1,137 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import inspect
+import os
+import sys
+import traceback
+from typing import IO
+
+from pyspark.errors import PySparkRuntimeError, PySparkValueError
+from pyspark.java_gateway import local_connect_and_auth
+from pyspark.serializers import (
+ read_bool,
+ read_int,
+ write_int,
+ write_with_length,
+ CPickleSerializer,
+ SpecialLengths,
+ UTF8Deserializer,
+)
+from pyspark.sql.types import StructType, _parse_datatype_json_string
+from pyspark.util import try_simplify_traceback
+from pyspark.worker import read_command
+
+pickleSer = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
Review Comment:
can you put a high level comment here describing that this runs on the Python interpreter and is responsible for receiving the call to the 'analyze' method, deserializing the arguments, invoking the method, serializing the result, and sending it back?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
Review Comment:
please also mention:
* who is expected to call this
* what parts of the Python UDTF in 'func' we inspect (i.e. the arguments)
* how we serialize and deserialize the arguments, and the 'analyze' result type
* which error conditions may occur and which exceptions we raise under those circumstances
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
Review Comment:
this says "non-static", but the method looks static and always returns the same schema?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ false)
+
+ val (worker: Socket, _) =
+ env.createPythonWorker(pythonExec, workerModule, envVars.asScala.toMap)
+ val releasedOrClosed = new AtomicBoolean(false)
Review Comment:
please add a brief comment mentioning why we need an atomic boolean here? what is the concurrency?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ false)
+
+ val (worker: Socket, _) =
+ env.createPythonWorker(pythonExec, workerModule, envVars.asScala.toMap)
+ val releasedOrClosed = new AtomicBoolean(false)
+ try {
+ val dataOut =
+ new DataOutputStream(new BufferedOutputStream(worker.getOutputStream, bufferSize))
+ val dataIn = new DataInputStream(new BufferedInputStream(worker.getInputStream, bufferSize))
+
+ // Python version of driver
+ PythonRDD.writeUTF(pythonVer, dataOut)
+
+ // Send Python UDTF
+ dataOut.writeInt(func.command.length)
+ dataOut.write(func.command.toArray)
+
+ // Send arguments
+ dataOut.writeInt(exprs.length)
+ exprs.foreach { expr =>
+ PythonRDD.writeUTF(expr.dataType.json, dataOut)
+ if (expr.foldable) {
Review Comment:
here we 'eval' the expression if it's foldable but non-literal, and that is supported. Can you update the PR description and any class comments/docs in this PR to mention that we can provide the expression value like this for not only literal expressions, but also constant but non-literal expressions as well?
##########
python/pyspark/sql/worker/analyze_udtf.py:
##########
@@ -0,0 +1,137 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import inspect
+import os
+import sys
+import traceback
+from typing import IO
+
+from pyspark.errors import PySparkRuntimeError, PySparkValueError
+from pyspark.java_gateway import local_connect_and_auth
+from pyspark.serializers import (
+ read_bool,
+ read_int,
+ write_int,
+ write_with_length,
+ CPickleSerializer,
+ SpecialLengths,
+ UTF8Deserializer,
+)
+from pyspark.sql.types import StructType, _parse_datatype_json_string
+from pyspark.util import try_simplify_traceback
+from pyspark.worker import read_command
+
+pickleSer = CPickleSerializer()
+utf8_deserializer = UTF8Deserializer()
+
+
+def main(infile: IO, outfile: IO) -> None:
+ try:
+ # Check Python version
Review Comment:
optional: this is a big block of code in the `try` block, should we move it into a helper function to reduce indentation?
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
+ class TestUDTF:
+ def analyze(self) -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_returning_non_struct(self):
Review Comment:
other test ideas:
positive tests:
* call the UDTF with integer types, strings, array/struct/map types
* call the UDTF with a relation using the TABLE keyword, and the 'analyze' method returns a relation with the same schema + an additional column
* call the UDTF with a constant (but non-literal) scalar integer argument, and the 'analyze' method returns a relation with the same number of columns as the value of that integer (or returns an error unless the integer is between 1 and 10).
negative tests:
* the 'analyze' method throws an explicit exception
* the 'analyze' method gets called with literal `NULL`
* the 'analyze' method attempts to lookup a key from the dictionary that doesn't exist
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
Review Comment:
can you also add a case with both a return type and an 'analyze' method?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1266380603
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +138,123 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ *
+ * When the Python UDTF is defined without a static return type,
+ * the analyze will call this while resolving table-valued functions.
+ *
+ * This expects the Python UDTF to have `analyze` static method that take arguments:
+ *
+ * - The number and order of arguments are the same as the UDTF inputs
+ * - Each argument is an `AnalyzeArgument`, containing:
+ * - data_type: DataType
+ * - value: Any: if the argument is foldable; otherwise None
+ * - is_table: bool: True if the argument is TABLE
+ *
+ * and that return an `AnalyzeResult`.
+ *
+ * It serializes/deserializes the data types via JSON,
+ * and the values for the case the argument is foldable are pickled.
+ *
+ * `AnalysisException` with the error class "TABLE_VALUED_FUNCTION_FAILED_TO_ANALYZE_IN_PYTHON"
+ * will be thrown when an exception is raised in Python.
+ */
+ def analyzeInPython(
+ tableArgs: Seq[Boolean])(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
Review Comment:
hmmm seems like this doesn't respect `PYSPARK_DRIVER_PYTHON`. It's fine to don't address this in this PR but we should be aware of this.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhengruifeng commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1266684996
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +138,123 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ *
+ * When the Python UDTF is defined without a static return type,
+ * the analyze will call this while resolving table-valued functions.
+ *
+ * This expects the Python UDTF to have `analyze` static method that take arguments:
+ *
+ * - The number and order of arguments are the same as the UDTF inputs
+ * - Each argument is an `AnalyzeArgument`, containing:
+ * - data_type: DataType
+ * - value: Any: if the argument is foldable; otherwise None
+ * - is_table: bool: True if the argument is TABLE
+ *
+ * and that return an `AnalyzeResult`.
+ *
+ * It serializes/deserializes the data types via JSON,
+ * and the values for the case the argument is foldable are pickled.
+ *
+ * `AnalysisException` with the error class "TABLE_VALUED_FUNCTION_FAILED_TO_ANALYZE_IN_PYTHON"
+ * will be thrown when an exception is raised in Python.
+ */
+ def analyzeInPython(
+ tableArgs: Seq[Boolean])(func: PythonFunction, exprs: Seq[Expression]): StructType = {
Review Comment:
will we cache the analyzed results?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] dtenedor commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "dtenedor (via GitHub)" <gi...@apache.org>.
dtenedor commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264031515
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
Review Comment:
Good question, we might be best off just returning error if we attempt to register the UDTF with a static return type, and the UDTF class has an βanalyzeβ method, to make them mutually exclusive.
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
Review Comment:
Good question, we might be best off just returning error if we attempt to register the UDTF with a static return type, and the UDTF class has an βanalyzeβ method, to make them mutually exclusive.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] allisonwang-db commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "allisonwang-db (via GitHub)" <gi...@apache.org>.
allisonwang-db commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267359161
##########
core/src/main/scala/org/apache/spark/SparkEnv.scala:
##########
@@ -73,7 +73,10 @@ class SparkEnv (
val conf: SparkConf) extends Logging {
@volatile private[spark] var isStopped = false
- private val pythonWorkers = mutable.HashMap[(String, Map[String, String]), PythonWorkerFactory]()
+
+ private case class PythonWorkersKey(
+ pythonExec: String, workerModule: String, daemonModule: String, envVars: Map[String, String])
Review Comment:
Nice! This is so much more readable. Nit: can we add a comment to explain the parameters here (workerModule and daemonModule) briefly?
##########
common/utils/src/main/resources/error/error-classes.json:
##########
@@ -2446,6 +2446,11 @@
],
"sqlState" : "42P01"
},
+ "TABLE_VALUED_FUNCTION_FAILED_TO_ANALYZE_IN_PYTHON" : {
+ "message" : [
+ "The user defined Python table function failed to analyze in Python, <msg>"
Review Comment:
```suggestion
"Failed to analyze the Python user defined table function: <msg>"
```
Is it possible to add the function name?
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -748,6 +769,442 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(func().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, DataType)
+ assert a.value is not None
+ assert a.is_table is False
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
+ # another data type
+ (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
+ # array type
+ (
+ func(array(lit(1), lit(2), lit(3))),
+ StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
+ [Row(a=[1, 2, 3])],
+ ),
+ # map type
+ (
+ func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
+ ),
+ [Row(a={"x": 1, "y": 2})],
+ ),
+ # struct type
+ (
+ func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a",
+ StructType()
+ .add("x", IntegerType(), nullable=False)
+ .add("y", IntegerType(), nullable=False),
+ ),
+ [Row(a=Row(x=1, y=2))],
+ ),
+ # use SQL
+ (
+ self.spark.sql("SELECT * from test_udtf(1)"),
+ StructType().add("a", IntegerType()),
+ [Row(a=1)],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
Review Comment:
can we use the new test utility methods here? `assertSchemaEqual` and `assertDataFrameEqual`
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -748,6 +769,442 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(func().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, DataType)
+ assert a.value is not None
+ assert a.is_table is False
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
+ # another data type
+ (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
+ # array type
+ (
+ func(array(lit(1), lit(2), lit(3))),
+ StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
+ [Row(a=[1, 2, 3])],
+ ),
+ # map type
+ (
+ func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
+ ),
+ [Row(a={"x": 1, "y": 2})],
+ ),
+ # struct type
+ (
+ func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a",
+ StructType()
+ .add("x", IntegerType(), nullable=False)
+ .add("y", IntegerType(), nullable=False),
+ ),
+ [Row(a=Row(x=1, y=2))],
+ ),
+ # use SQL
+ (
+ self.spark.sql("SELECT * from test_udtf(1)"),
+ StructType().add("a", IntegerType()),
+ [Row(a=1)],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_decorator(self):
+ @udtf
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_decorator_parens(self):
+ @udtf()
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("a", a.data_type).add("b", b.data_type))
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1), lit("x")),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_arbitary_number_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(*args: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(
+ StructType([StructField(f"col{i}", a.data_type) for i, a in enumerate(args)])
+ )
+
+ def eval(self, *args):
+ yield args
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1)),
+ StructType().add("col0", IntegerType()),
+ [Row(a=1)],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("col0", IntegerType()).add("col1", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ # TODO(SPARK-44479): Support Python UDTFs with empty schema
+ # (func(), StructType(), [Row()]),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, StructType)
+ assert a.value is None
+ assert a.is_table is True
+ return AnalyzeResult(StructType().add("a", a.data_type[0].dataType))
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_analyze_table_argument_adding_columns(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a.data_type, StructType)
+ assert a.is_table is True
+ return AnalyzeResult(a.data_type.add("is_even", BooleanType()))
+
+ def eval(self, a: Row):
+ yield a["id"], a["id"] % 2 == 0
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(
+ df.schema,
+ StructType().add("id", LongType(), nullable=False).add("is_even", BooleanType()),
+ )
+ self.assertEqual(
+ df.collect(),
+ [
+ Row(a=0, is_even=True),
+ Row(a=1, is_even=False),
+ Row(a=2, is_even=True),
+ Row(a=3, is_even=False),
+ ],
+ )
+
+ def test_udtf_with_analyze_table_argument_repeating_rows(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(n, row) -> AnalyzeResult:
+ if n.value is None or not isinstance(n.value, int) or (n.value < 1 or n.value > 10):
+ raise Exception("The first argument must be a scalar integer between 1 and 10")
+
+ if row.is_table is False:
+ raise Exception("The second argument must be a table argument")
+
+ assert isinstance(row.data_type, StructType)
+ return AnalyzeResult(row.data_type)
+
+ def eval(self, n: int, row: Row):
+ for _ in range(n):
+ yield row
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected_schema = StructType().add("id", LongType(), nullable=False)
+ expected_results = [
+ Row(a=0),
+ Row(a=0),
+ Row(a=1),
+ Row(a=1),
+ Row(a=2),
+ Row(a=2),
+ Row(a=3),
+ Row(a=3),
+ ]
+ for i, df in enumerate(
+ [
+ self.spark.sql("SELECT * FROM test_udtf(2, TABLE (SELECT id FROM range(0, 4)))"),
+ self.spark.sql(
+ "SELECT * FROM test_udtf(1 + 1, TABLE (SELECT id FROM range(0, 4)))"
Review Comment:
Since we are accessing `.value` for the analyze argument, what if the input value here is not a foldable one? for example `SELECT * FROM values (0, 1) t(c1, c2), lateral test_udtf(c1, ...)`.
##########
python/pyspark/sql/udtf.py:
##########
@@ -27,20 +29,32 @@
from pyspark.rdd import PythonEvalType
from pyspark.sql.column import _to_java_column, _to_seq
from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version
-from pyspark.sql.types import StructType, _parse_datatype_string
+from pyspark.sql.types import DataType, StructType, _parse_datatype_string
from pyspark.sql.udf import _wrap_function
if TYPE_CHECKING:
from pyspark.sql._typing import ColumnOrName
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.session import SparkSession
-__all__ = ["UDTFRegistration"]
+__all__ = ["AnalyzeArgument", "AnalyzeResult", "UDTFRegistration"]
+
+
+@dataclass(frozen=True)
+class AnalyzeArgument:
Review Comment:
nit: let's add some comments for this class and variables.
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -748,6 +769,442 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(func().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, DataType)
+ assert a.value is not None
+ assert a.is_table is False
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
+ # another data type
+ (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
+ # array type
+ (
+ func(array(lit(1), lit(2), lit(3))),
+ StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
+ [Row(a=[1, 2, 3])],
+ ),
+ # map type
+ (
+ func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
+ ),
+ [Row(a={"x": 1, "y": 2})],
+ ),
+ # struct type
+ (
+ func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a",
+ StructType()
+ .add("x", IntegerType(), nullable=False)
+ .add("y", IntegerType(), nullable=False),
+ ),
+ [Row(a=Row(x=1, y=2))],
+ ),
+ # use SQL
+ (
+ self.spark.sql("SELECT * from test_udtf(1)"),
+ StructType().add("a", IntegerType()),
+ [Row(a=1)],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_decorator(self):
+ @udtf
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_decorator_parens(self):
+ @udtf()
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("a", a.data_type).add("b", b.data_type))
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1), lit("x")),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_arbitary_number_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(*args: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(
+ StructType([StructField(f"col{i}", a.data_type) for i, a in enumerate(args)])
+ )
+
+ def eval(self, *args):
+ yield args
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1)),
+ StructType().add("col0", IntegerType()),
+ [Row(a=1)],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("col0", IntegerType()).add("col1", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ # TODO(SPARK-44479): Support Python UDTFs with empty schema
+ # (func(), StructType(), [Row()]),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, StructType)
+ assert a.value is None
+ assert a.is_table is True
+ return AnalyzeResult(StructType().add("a", a.data_type[0].dataType))
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_analyze_table_argument_adding_columns(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a.data_type, StructType)
+ assert a.is_table is True
+ return AnalyzeResult(a.data_type.add("is_even", BooleanType()))
+
+ def eval(self, a: Row):
+ yield a["id"], a["id"] % 2 == 0
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(
+ df.schema,
+ StructType().add("id", LongType(), nullable=False).add("is_even", BooleanType()),
+ )
+ self.assertEqual(
+ df.collect(),
+ [
+ Row(a=0, is_even=True),
+ Row(a=1, is_even=False),
+ Row(a=2, is_even=True),
+ Row(a=3, is_even=False),
+ ],
+ )
+
+ def test_udtf_with_analyze_table_argument_repeating_rows(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(n, row) -> AnalyzeResult:
+ if n.value is None or not isinstance(n.value, int) or (n.value < 1 or n.value > 10):
+ raise Exception("The first argument must be a scalar integer between 1 and 10")
+
+ if row.is_table is False:
+ raise Exception("The second argument must be a table argument")
+
+ assert isinstance(row.data_type, StructType)
+ return AnalyzeResult(row.data_type)
+
+ def eval(self, n: int, row: Row):
+ for _ in range(n):
+ yield row
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected_schema = StructType().add("id", LongType(), nullable=False)
+ expected_results = [
+ Row(a=0),
+ Row(a=0),
+ Row(a=1),
+ Row(a=1),
+ Row(a=2),
+ Row(a=2),
+ Row(a=3),
+ Row(a=3),
+ ]
+ for i, df in enumerate(
+ [
+ self.spark.sql("SELECT * FROM test_udtf(2, TABLE (SELECT id FROM range(0, 4)))"),
+ self.spark.sql(
+ "SELECT * FROM test_udtf(1 + 1, TABLE (SELECT id FROM range(0, 4)))"
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ with self.assertRaisesRegex(
+ AnalysisException, "The first argument must be a scalar integer between 1 and 10"
+ ):
+ self.spark.sql(
+ "SELECT * FROM test_udtf(0, TABLE (SELECT id FROM range(0, 4)))"
+ ).collect()
+
+ with self.sql_conf(
+ {"spark.sql.tvf.allowMultipleTableArguments.enabled": True}
+ ), self.assertRaisesRegex(
+ AnalysisException, "The first argument must be a scalar integer between 1 and 10"
+ ):
+ self.spark.sql(
+ """
+ SELECT * FROM test_udtf(
+ TABLE (SELECT id FROM range(0, 1)),
+ TABLE (SELECT id FROM range(0, 4)))
+ """
+ ).collect()
+
+ with self.assertRaisesRegex(
+ AnalysisException, "The second argument must be a table argument"
+ ):
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')").collect()
+
+ def test_udtf_with_both_return_type_and_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF, returnType="c1: string, c2: string")
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_non_staticmethod(self):
+ class TestUDTF:
+ def analyze(self) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
Review Comment:
Instead of throwing this `INVALID_UDTF_RETURN_TYPE` error, shall we check if the UDTF class has an `analyze` method and throw an error if it is not static? Or we could update the error message for INVALID_UDTF_RETURN_TYPE to suggest making `analyze` a staticmethod.
##########
python/pyspark/sql/worker/analyze_udtf.py:
##########
@@ -0,0 +1,142 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import inspect
+import os
+import sys
+import traceback
+from typing import List, IO
+
+from pyspark.errors import PySparkRuntimeError, PySparkValueError
+from pyspark.java_gateway import local_connect_and_auth
+from pyspark.serializers import (
+ read_bool,
+ read_int,
+ write_int,
+ write_with_length,
+ SpecialLengths,
+)
+from pyspark.sql.types import _parse_datatype_json_string
+from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult
+from pyspark.util import try_simplify_traceback
+from pyspark.worker import check_python_version, read_command, pickleSer, utf8_deserializer
+
+
+def read_udtf(infile: IO) -> type:
+ """Reads the Python UDTF and checks if its valid or not."""
+ # Receive Python UDTF
+ handler = read_command(pickleSer, infile)
+ if not isinstance(handler, type):
+ raise PySparkRuntimeError(
+ f"Invalid UDTF handler type. Expected a class (type 'type'), but "
+ f"got an instance of {type(handler).__name__}."
+ )
+
+ if not hasattr(handler, "analyze") or not isinstance(
+ inspect.getattr_static(handler, "analyze"), staticmethod
+ ):
+ raise PySparkRuntimeError(
+ "Failed to execute the user defined table function because it has not "
+ "implemented the 'analyze' static method or specified a fixed "
+ "return type during registration time. "
+ "Please add the 'analyze' static method or specify the return type, "
+ "and try the query again."
+ )
+ return handler
+
+
+def read_arguments(infile: IO) -> List[AnalyzeArgument]:
+ """Reads the arguments for `analyze` static method."""
+ # Receive arguments
+ num_args = read_int(infile)
+ args: List[AnalyzeArgument] = []
+ for _ in range(num_args):
+ dt = _parse_datatype_json_string(utf8_deserializer.loads(infile))
+ if read_bool(infile): # is foldable
+ value = pickleSer._read_with_length(infile)
+ if dt.needConversion():
+ value = dt.fromInternal(value)
+ else:
+ value = None
Review Comment:
Let's add this as a comment in the AnalyzeArgument data class: The value field will be None when the input value is not foldable.
##########
python/pyspark/sql/udtf.py:
##########
@@ -27,20 +29,32 @@
from pyspark.rdd import PythonEvalType
from pyspark.sql.column import _to_java_column, _to_seq
from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version
-from pyspark.sql.types import StructType, _parse_datatype_string
+from pyspark.sql.types import DataType, StructType, _parse_datatype_string
from pyspark.sql.udf import _wrap_function
if TYPE_CHECKING:
from pyspark.sql._typing import ColumnOrName
from pyspark.sql.dataframe import DataFrame
from pyspark.sql.session import SparkSession
-__all__ = ["UDTFRegistration"]
+__all__ = ["AnalyzeArgument", "AnalyzeResult", "UDTFRegistration"]
+
+
+@dataclass(frozen=True)
+class AnalyzeArgument:
+ data_type: DataType
+ value: Optional[Any]
+ is_table: bool
+
+
+@dataclass(frozen=True)
+class AnalyzeResult:
Review Comment:
ditto
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala:
##########
@@ -176,6 +177,28 @@ case class PythonUDTF(
copy(children = newChildren)
}
+/**
+ * A place holder of a Polymorphic Python table-valued function.
Review Comment:
```suggestion
* A placeholder of a polymorphic Python table-valued function.
```
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
:+1:
##########
python/pyspark/sql/udtf.py:
##########
@@ -153,6 +175,19 @@ def _validate_udtf_handler(cls: Any) -> None:
error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": cls.__name__}
)
+ has_analyze_staticmethod = hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
Review Comment:
I wonder if we should throw a warning message if the UDTF handler class contains `analyze,` but it is not a static method.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -63,18 +77,51 @@ case class UserDefinedPythonFunction(
case class UserDefinedPythonTableFunction(
name: String,
func: PythonFunction,
- returnType: StructType,
+ returnType: Option[StructType],
pythonEvalType: Int,
udfDeterministic: Boolean) {
- def builder(e: Seq[Expression]): LogicalPlan = {
- val udtf = PythonUDTF(
- name = name,
- func = func,
- elementSchema = returnType,
- children = e,
- evalType = pythonEvalType,
- udfDeterministic = udfDeterministic)
+ def this(
+ name: String,
+ func: PythonFunction,
+ returnType: StructType,
+ pythonEvalType: Int,
+ udfDeterministic: Boolean) = {
+ this(name, func, Some(returnType), pythonEvalType, udfDeterministic)
+ }
+
+ def this(
+ name: String,
+ func: PythonFunction,
+ pythonEvalType: Int,
+ udfDeterministic: Boolean) = {
+ this(name, func, None, pythonEvalType, udfDeterministic)
+ }
+
+ def builder(exprs: Seq[Expression]): LogicalPlan = {
+ val udtf = returnType match {
+ case Some(rt) =>
+ PythonUDTF(
+ name = name,
+ func = func,
+ elementSchema = rt,
+ children = exprs,
+ evalType = pythonEvalType,
+ udfDeterministic = udfDeterministic)
+ case _ =>
+ UnresolvedPolymorphicPythonUDTF(
+ name = name,
+ func = func,
+ children = exprs,
+ evalType = pythonEvalType,
+ udfDeterministic = udfDeterministic,
+ resolveElementSchema = UserDefinedPythonTableFunction.analyzeInPython(
+ exprs.map {
+ case _: FunctionTableSubqueryArgumentExpression => true
+ case NamedArgumentExpression(_, _: FunctionTableSubqueryArgumentExpression) => true
+ case _ => false
Review Comment:
Can we add a comment to explain this code logic here?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267138618
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +138,123 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ *
+ * When the Python UDTF is defined without a static return type,
+ * the analyze will call this while resolving table-valued functions.
+ *
+ * This expects the Python UDTF to have `analyze` static method that take arguments:
+ *
+ * - The number and order of arguments are the same as the UDTF inputs
+ * - Each argument is an `AnalyzeArgument`, containing:
+ * - data_type: DataType
+ * - value: Any: if the argument is foldable; otherwise None
+ * - is_table: bool: True if the argument is TABLE
+ *
+ * and that return an `AnalyzeResult`.
+ *
+ * It serializes/deserializes the data types via JSON,
+ * and the values for the case the argument is foldable are pickled.
+ *
+ * `AnalysisException` with the error class "TABLE_VALUED_FUNCTION_FAILED_TO_ANALYZE_IN_PYTHON"
+ * will be thrown when an exception is raised in Python.
+ */
+ def analyzeInPython(
+ tableArgs: Seq[Boolean])(func: PythonFunction, exprs: Seq[Expression]): StructType = {
Review Comment:
Once it's analyzed, it will be `PythonUDTF` with the returned fixed schema.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon closed pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon closed pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
URL: https://github.com/apache/spark/pull/41948
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] HyukjinKwon commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1643063875
Merged to master.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] allisonwang-db commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "allisonwang-db (via GitHub)" <gi...@apache.org>.
allisonwang-db commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1644951738
Yes. Created a new task: [SPARK-44508](https://issues.apache.org/jira/browse/SPARK-44508)
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263165177
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
Review Comment:
We can't remove it because the class is written in Java that doesn't allow named arguments.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] dtenedor commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "dtenedor (via GitHub)" <gi...@apache.org>.
dtenedor commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264029412
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
Review Comment:
My mistake on this, carry on!
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ false)
Review Comment:
My mistake on this, carry on!
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263165177
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
Review Comment:
We can't use named arguments here because the class `Pickler` is written in Java that doesn't allow named arguments.
Or just `true`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264196125
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
+ class TestUDTF:
+ def analyze(self) -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_returning_non_struct(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze():
+ return StringType()
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException,
+ "Output of `analyze` static method of Python UDTFs expects a StructType "
+ "but got: <class 'pyspark.sql.types.StringType'>",
+ ):
+ func().collect()
+
+ def test_udtf_with_analyze_taking_wrong_number_of_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a):
+ yield a, a + 1
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException, r"analyze\(\) missing 1 required positional argument: 'b'"
Review Comment:
For this one, let me leave it as-is for now because the error message is from Python runtime using a general error type and message. If we capture and modify it, it could also modify an actual error from inside of the `analyze` method.
I'll revisit this later.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1636598237
Based on the offline discussion, I'll update this to:
- define a class for the arguments instead of using `dict`
- define a class for the return instead of using `StructType`. The class will have a field `schema: StructType`
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1641222671
@rangadi Could you take a look at the change 0e268d1?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhengruifeng commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1266684996
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +138,123 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ *
+ * When the Python UDTF is defined without a static return type,
+ * the analyze will call this while resolving table-valued functions.
+ *
+ * This expects the Python UDTF to have `analyze` static method that take arguments:
+ *
+ * - The number and order of arguments are the same as the UDTF inputs
+ * - Each argument is an `AnalyzeArgument`, containing:
+ * - data_type: DataType
+ * - value: Any: if the argument is foldable; otherwise None
+ * - is_table: bool: True if the argument is TABLE
+ *
+ * and that return an `AnalyzeResult`.
+ *
+ * It serializes/deserializes the data types via JSON,
+ * and the values for the case the argument is foldable are pickled.
+ *
+ * `AnalysisException` with the error class "TABLE_VALUED_FUNCTION_FAILED_TO_ANALYZE_IN_PYTHON"
+ * will be thrown when an exception is raised in Python.
+ */
+ def analyzeInPython(
+ tableArgs: Seq[Boolean])(func: PythonFunction, exprs: Seq[Expression]): StructType = {
Review Comment:
will we cache the analyzed results in the future?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1268593999
##########
python/pyspark/sql/udtf.py:
##########
@@ -153,6 +175,19 @@ def _validate_udtf_handler(cls: Any) -> None:
error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": cls.__name__}
)
+ has_analyze_staticmethod = hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
Review Comment:
So now at [4e2b6be](https://github.com/apache/spark/pull/41948/commits/4e2b6befbd015fcdad1ddf6ccecd79a630ac7d60),
- When `returnType` is not provided, the class must have `analyze` static method; otherwise raise an error `INVALID_UDTF_RETURN_TYPE`.
- When `returnType` is provided and also the class has any `analyze` attribute, raise an error `INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264182071
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
+ class TestUDTF:
+ def analyze(self) -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_returning_non_struct(self):
Review Comment:
Updated/added:
- test_udtf_with_analyze
- test_udtf_with_analyze_table_argument_adding_columns
- test_udtf_with_analyze_table_argument_repeating_rows
- test_udtf_with_analyze_raising_an_exception
- test_udtf_with_analyze_null_literal
- test_udtf_with_analyze_unknown_key
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264266121
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
Discussed offline and decided so far we don't support non-`@staticmethod` analyze, to avoid confusion for a case like:
```py
class MyUDTF:
def analyze(self, a):
self.some_boolean_value = ...
return StructType(...)
def eval(self, a):
if self.some_boolean_value:
yield a + self._foo
```
The instances between `analyze` and `eval` will be different, so it worn't work as expected.
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
Discussed offline and decided we don't support non-`@staticmethod` analyze for now, to avoid confusion for a case like:
```py
class MyUDTF:
def analyze(self, a):
self.some_boolean_value = ...
return StructType(...)
def eval(self, a):
if self.some_boolean_value:
yield a + self._foo
```
The instances between `analyze` and `eval` will be different, so it worn't work as expected.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263165177
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
Review Comment:
We can't remove it because the class is written in Java that doesn't allow named arguments.
Or just `true`?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263166608
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
Review Comment:
That's interesting. Which should be picked when there're both? Or should raise an exception?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] dtenedor commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "dtenedor (via GitHub)" <gi...@apache.org>.
dtenedor commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264230197
##########
python/pyspark/sql/worker/analyze_udtf.py:
##########
@@ -0,0 +1,140 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import inspect
+import os
+import sys
+import traceback
+from typing import Any, Dict, List, IO
+
+from pyspark.errors import PySparkRuntimeError, PySparkValueError
+from pyspark.java_gateway import local_connect_and_auth
+from pyspark.serializers import (
+ read_bool,
+ read_int,
+ write_int,
+ write_with_length,
+ SpecialLengths,
+)
+from pyspark.sql.types import StructType, _parse_datatype_json_string
+from pyspark.util import try_simplify_traceback
+from pyspark.worker import check_python_version, read_command, pickleSer, utf8_deserializer
+
+
+def read_udtf(infile: IO) -> type:
+ """Reads the Python UDTF and checks if its valid or not."""
+ # Receive Python UDTF
+ handler = read_command(pickleSer, infile)
+ if not isinstance(handler, type):
+ raise PySparkRuntimeError(
+ f"Invalid UDTF handler type. Expected a class (type 'type'), but "
+ f"got an instance of {type(handler).__name__}."
+ )
+
+ if not hasattr(handler, "analyze") or not isinstance(
+ inspect.getattr_static(handler, "analyze"), staticmethod
+ ):
+ raise PySparkRuntimeError(
+ "Failed to execute the user defined table function because it has not "
+ "implemented the 'analyze' static method. "
Review Comment:
```suggestion
"implemented the 'analyze' static method or specified a fixed "
"return type during registration time. "
```
##########
python/pyspark/sql/functions.py:
##########
@@ -50,6 +50,7 @@
# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409
from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401
from pyspark.sql.udtf import UserDefinedTableFunction, _create_py_udtf
+from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401
Review Comment:
please sort alphabetically?
Also, F401 is for "Module imported but unused", do we need to import the AnalyzeArgument and AnalyzeResult here, or can we just skip the import? Or do we need this here so users can import these symbols directly from this `functions.py` as well?
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -753,6 +773,376 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ self.assertEqual(func().collect(), [Row(c1="hello", c2="world")])
+ self.assertEqual(
+ self.spark.sql("SELECT * FROM test_udtf()").collect(), [Row(c1="hello", c2="world")]
+ )
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, DataType)
+ assert a.value is not None
+ assert a.is_table is False
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
+ # another data type
+ (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
+ # array type
+ (
+ func(array(lit(1), lit(2), lit(3))),
+ StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
+ [Row(a=[1, 2, 3])],
+ ),
+ # map type
+ (
+ func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
+ ),
+ [Row(a={"x": 1, "y": 2})],
+ ),
+ # struct type
+ (
+ func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a",
+ StructType()
+ .add("x", IntegerType(), nullable=False)
+ .add("y", IntegerType(), nullable=False),
+ ),
+ [Row(a=Row(x=1, y=2))],
+ ),
+ # use SQL
+ (
+ self.spark.sql("SELECT * from test_udtf(1)"),
+ StructType().add("a", IntegerType()),
+ [Row(a=1)],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_decorator(self):
+ @udtf
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ self.assertEqual(TestUDTF().collect(), [Row(c1="hello", c2="world")])
+ self.assertEqual(
+ self.spark.sql("SELECT * FROM test_udtf()").collect(), [Row(c1="hello", c2="world")]
+ )
+
+ def test_udtf_with_analyze_decorator_parens(self):
+ @udtf()
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ self.assertEqual(TestUDTF().collect(), [Row(c1="hello", c2="world")])
+ self.assertEqual(
+ self.spark.sql("SELECT * FROM test_udtf()").collect(), [Row(c1="hello", c2="world")]
+ )
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("a", a.data_type).add("b", b.data_type))
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1), lit("x")),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, StructType)
+ assert a.value is None
+ assert a.is_table is True
+ return AnalyzeResult(StructType().add("a", a.data_type[0].dataType))
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_analyze_table_argument_adding_columns(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a.data_type, StructType)
+ assert a.is_table is True
+ return AnalyzeResult(a.data_type.add("is_even", BooleanType()))
+
+ def eval(self, a: Row):
+ yield a["id"], a["id"] % 2 == 0
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(
+ df.schema,
+ StructType().add("id", LongType(), nullable=False).add("is_even", BooleanType()),
+ )
+ self.assertEqual(
+ df.collect(),
+ [
+ Row(a=0, is_even=True),
+ Row(a=1, is_even=False),
+ Row(a=2, is_even=True),
+ Row(a=3, is_even=False),
+ ],
+ )
+
+ def test_udtf_with_analyze_table_argument_repeating_rows(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(n, row) -> AnalyzeResult:
+ if n.value is None or not isinstance(n.value, int) or (n.value < 1 or n.value > 10):
+ raise Exception("The first argument must be a scalar integer between 1 and 10")
+
+ if row.is_table is False:
+ raise Exception("The second argument must be a table argument")
+
+ assert isinstance(row.data_type, StructType)
+ return AnalyzeResult(row.data_type)
+
+ def eval(self, n: int, row: Row):
+ for _ in range(n):
+ yield row
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df1 = self.spark.sql("SELECT * FROM test_udtf(2, TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(df1.schema, StructType().add("id", LongType(), nullable=False))
+ self.assertEqual(
+ df1.collect(),
+ [Row(a=0), Row(a=0), Row(a=1), Row(a=1), Row(a=2), Row(a=2), Row(a=3), Row(a=3)],
+ )
+
+ df2 = self.spark.sql("SELECT * FROM test_udtf(1 + 1, TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(df2.schema, StructType().add("id", LongType(), nullable=False))
+ self.assertEqual(
+ df2.collect(),
+ [Row(a=0), Row(a=0), Row(a=1), Row(a=1), Row(a=2), Row(a=2), Row(a=3), Row(a=3)],
+ )
+
+ with self.assertRaisesRegex(
+ AnalysisException, "The first argument must be a scalar integer between 1 and 10"
+ ):
+ self.spark.sql(
+ "SELECT * FROM test_udtf(0, TABLE (SELECT id FROM range(0, 4)))"
+ ).collect()
+
+ with self.sql_conf(
+ {"spark.sql.tvf.allowMultipleTableArguments.enabled": True}
+ ), self.assertRaisesRegex(
+ AnalysisException, "The first argument must be a scalar integer between 1 and 10"
+ ):
+ self.spark.sql(
+ """
+ SELECT * FROM test_udtf(
+ TABLE (SELECT id FROM range(0, 1)),
+ TABLE (SELECT id FROM range(0, 4)))
+ """
+ ).collect()
+
+ with self.assertRaisesRegex(
+ AnalysisException, "The second argument must be a table argument"
+ ):
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')").collect()
+
+ def test_udtf_with_both_return_type_and_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF, returnType="c1: string, c2: string")
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_non_staticmethod(self):
+ class TestUDTF:
+ def analyze(self) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_returning_non_struct(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze():
+ return StringType()
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(
+ AnalysisException,
+ "Output of `analyze` static method of Python UDTFs expects "
+ "a pyspark.sql.udtf.AnalyzeResult but got: <class 'pyspark.sql.types.StringType'>",
+ ):
+ func().collect()
+
+ def test_udtf_with_analyze_raising_an_exception(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ raise Exception("Failed to analyze.")
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+
+ with self.assertRaisesRegex(AnalysisException, "Failed to analyze."):
+ func().collect()
+
+ def test_udtf_with_analyze_null_literal(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(None))
+ self.assertEqual(df.schema, StructType().add("a", NullType()))
+ self.assertEqual(df.collect(), [Row(a=None)])
+
+ def test_udtf_with_analyze_taking_wrong_number_of_arguments(self):
Review Comment:
we talked offline, it seems possible to make a test UDTF where the `analyze` method accepts non-keyword arguments: `def analyze(*args) -> AnalyzeResult:`. Can we please add a test case for this?
We can add another test to exercise `def analyze(**kwargs) -> AnalyzeResult` as well? In the future, we could return some reasonable error message in this case like "Failed to evaluate table-valued function XXXXX because keyword arguments are not supported for the 'analyze' static method."
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,122 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ *
+ * When the Python UDTF is defined without a static return type,
+ * the analyze will call this while resolving table-valued functions.
Review Comment:
```suggestion
* the analyzer will call this while resolving table-valued functions.
```
##########
connector/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala:
##########
@@ -2482,20 +2482,24 @@ class SparkConnectPlanner(val sessionHolder: SessionHolder) extends Logging {
private def createPythonUserDefinedTableFunction(
fun: proto.CommonInlineUserDefinedTableFunction): UserDefinedPythonTableFunction = {
val udtf = fun.getPythonUdtf
- // Currently return type is required for Python UDTFs.
- // TODO(SPARK-44380): support `analyze` in Python UDTFs
- assert(udtf.hasReturnType)
- val returnType = transformDataType(udtf.getReturnType)
- if (!returnType.isInstanceOf[StructType]) {
- throw InvalidPlanInput(
- "Invalid Python user-defined table function return type. " +
- s"Expect a struct type, but got ${returnType.typeName}.")
+ val returnType = if (udtf.hasReturnType) {
+ val returnType = transformDataType(udtf.getReturnType)
+
+ if (!returnType.isInstanceOf[StructType]) {
Review Comment:
optional: should we do
```
returnType match {
case s: StructType => s
case _ => throw InvalidPlanInput(...)
}
```
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
π
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
π
##########
python/pyspark/errors/error_classes.py:
##########
@@ -268,6 +268,11 @@
"Eval type for UDF must be <eval_type>."
]
},
+ "INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD" : {
+ "message" : [
+ "The UDTF '<name>' is invalid. It has both its return type and the required 'analyze' static method. Please make it have one of either the return type or the 'analyze' static method in '<name>' and try again."
Review Comment:
π this seems helpful to avoid confusion.
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +730,394 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ self.assertEqual(func().collect(), [Row(c1="hello", c2="world")])
Review Comment:
maybe assign these expected result rows to a variable in each case, so we don't have to repeat it again for the SQL case?
##########
python/pyspark/sql/functions.py:
##########
@@ -15564,6 +15566,36 @@ def udtf(
| 1| 2|
+---+---+
+ UDTF can also have `analyze` static method instead of a static return type:
+
+ The `analyze` static method should take arguments:
+
+ - The number and order of arguments are the same as the UDTF inputs
+ - Each argument is a :class:`pyspark.sql.udtf.AnalyzeArgument`, containing:
+ - data_type: DataType
+ - value: Any: if the argument is foldable; otherwise None
+ - is_table: bool: True if the argument is TABLE
+
+ and return a :class:`pyspark.sql.udtf.AnalyzeResult`, containing.
Review Comment:
optional: maybe also add the "import" statement we'd need to include at the top of the file to gain access to `@udtf`, `AnalyzeArgument`, and `AnalyzeResult` symbols?
##########
python/pyspark/sql/connect/udtf.py:
##########
@@ -33,8 +33,8 @@
)
from pyspark.sql.connect.types import UnparsedDataType
from pyspark.sql.connect.utils import get_python_ver
-from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration
-from pyspark.sql.udtf import _validate_udtf_handler
+from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration, _validate_udtf_handler
+from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401
Review Comment:
F401 is for "Module imported but unused", do we need to import the `AnalyzeArgument` and `AnalyzeResult` here, or can we just skip the import?
##########
python/pyspark/errors/error_classes.py:
##########
@@ -268,6 +268,11 @@
"Eval type for UDF must be <eval_type>."
]
},
+ "INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD" : {
+ "message" : [
+ "The UDTF '<name>' is invalid. It has both its return type and the required 'analyze' static method. Please make it have one of either the return type or the 'analyze' static method in '<name>' and try again."
Review Comment:
π this seems helpful to avoid confusion.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267273028
##########
python/pyspark/sql/connect/udtf.py:
##########
@@ -33,8 +33,8 @@
)
from pyspark.sql.connect.types import UnparsedDataType
from pyspark.sql.connect.utils import get_python_ver
-from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration
-from pyspark.sql.udtf import _validate_udtf_handler
+from pyspark.sql.udtf import UDTFRegistration as PySparkUDTFRegistration, _validate_udtf_handler
+from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401
Review Comment:
I think it's still good to have this here for a case when a user wants to directly use `connect` package.
```py
from pyspark.sql.connect.udtf import AnalyzeArgument, AnalyzeResult
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267474580
##########
python/pyspark/sql/udtf.py:
##########
@@ -153,6 +175,19 @@ def _validate_udtf_handler(cls: Any) -> None:
error_class="INVALID_UDTF_NO_EVAL", message_parameters={"name": cls.__name__}
)
+ has_analyze_staticmethod = hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
Review Comment:
cc @dtenedor WDYT?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] cloud-fan commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "cloud-fan (via GitHub)" <gi...@apache.org>.
cloud-fan commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1644948542
Shall we add a doc page for Python UDTF? It does not have an interface like java so there is no place to put API doc.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263165949
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -719,6 +726,153 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> StructType:
+ return StructType().add("c1", StringType()).add("c2", StringType())
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ rows = func().collect()
+ self.assertEqual(rows, [Row(c1="hello", c2="world")])
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], DataType)
+ assert a["value"] is not None
+ assert a["is_table"] is False
+ return StructType().add("a", a["data_type"])
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+
+ df1 = func(lit(1))
+ self.assertEquals(df1.schema, StructType().add("a", IntegerType()))
+ self.assertEqual(df1.collect(), [Row(a=1)])
+
+ df2 = func(lit("x"))
+ self.assertEquals(df2.schema, StructType().add("a", StringType()))
+ self.assertEqual(df2.collect(), [Row(a="x")])
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a, b) -> StructType:
+ return StructType().add("a", a["data_type"]).add("b", b["data_type"])
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+
+ df = func(lit(1), lit("x"))
+ self.assertEquals(df.schema, StructType().add("a", IntegerType()).add("b", StringType()))
+ self.assertEqual(df.collect(), [Row(a=1, b="x")])
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a) -> StructType:
+ assert isinstance(a, dict)
+ assert isinstance(a["data_type"], StructType)
+ assert a["value"] is None
+ assert a["is_table"] is True
+ return StructType().add("a", a["data_type"][0].dataType)
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_non_static_analyze(self):
Review Comment:
I meant non-`@staticmethod`. will update the test name.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263165177
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
Review Comment:
We can't use named arguments here because the class is written in Java that doesn't allow named arguments.
Or just `true`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] allisonwang-db commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "allisonwang-db (via GitHub)" <gi...@apache.org>.
allisonwang-db commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264218509
##########
python/pyspark/errors/error_classes.py:
##########
@@ -277,6 +277,11 @@
"The UDTF '<name>' is invalid. It does not implement the required 'eval' method. Please implement the 'eval' method in '<name>' and try again."
]
},
+ "INVALID_UDTF_RETURN_TYPE" : {
Review Comment:
I have the same error class defined here: https://github.com/apache/spark/pull/41989/files#diff-2823e146fc0e6bddff3505b5bee6e2b855782d9f71e900e6f9099fc97d1fffa6R280. I will change the error class name in my PR.
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
Why do we make `analyze` a `staticmethod`?
For example, we could have a UDTF defined in this way:
```
class MyUDTF:
def __init__(self):
self.some_boolean_value = ...
def analyze(self, a, b):
if self.some_boolean_value:
return StructType(..)
...
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhengruifeng commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1266661826
##########
python/pyspark/errors/error_classes.py:
##########
@@ -268,6 +268,11 @@
"Eval type for UDF must be <eval_type>."
]
},
+ "INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD" : {
+ "message" : [
+ "The UDTF '<name>' is invalid. It has both its return type and the required 'analyze' static method. Please make it have one of either the return type or the 'analyze' static method in '<name>' and try again."
Review Comment:
just to confirm, it fails even if the user-specified return type == the analyzed result.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1264266121
##########
python/pyspark/sql/udtf.py:
##########
@@ -103,6 +104,14 @@ class VectorizedUDTF:
def __init__(self) -> None:
self.func = cls()
+ if hasattr(cls, "analyze") and isinstance(
+ inspect.getattr_static(cls, "analyze"), staticmethod
+ ):
+
+ @staticmethod
Review Comment:
Discussed offline and so far we don't support non-`@staticmethod` analyze, to avoid confusion for a case like:
```py
class MyUDTF:
def analyze(self, a):
self.some_boolean_value = ...
return StructType(...)
def eval(self, a):
if self.some_boolean_value:
yield a + self._foo
```
The instances between `analyze` and `eval` will be different, so it worn't work as expected.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1642586143
The failing test is [Spark on Kubernetes Integration test](https://github.com/ueshin/apache-spark/actions/runs/5596078914/jobs/10247102939#logs) that seems to be broken in `master` branch.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] zhengruifeng commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1266666757
##########
core/src/main/scala/org/apache/spark/SparkEnv.scala:
##########
@@ -73,7 +73,8 @@ class SparkEnv (
val conf: SparkConf) extends Logging {
@volatile private[spark] var isStopped = false
- private val pythonWorkers = mutable.HashMap[(String, Map[String, String]), PythonWorkerFactory]()
+ private val pythonWorkers =
+ mutable.HashMap[(String, String, String, Map[String, String]), PythonWorkerFactory]()
Review Comment:
nit: shall we define a class to replace `(String, String, String, Map[String, String])`?
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on PR #41948:
URL: https://github.com/apache/spark/pull/41948#issuecomment-1634799572
cc @dtenedor @allisonwang-db @zhengruifeng
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1263165512
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/UserDefinedPythonFunction.scala:
##########
@@ -91,3 +122,104 @@ case class UserDefinedPythonTableFunction(
Dataset.ofRows(session, udtf)
}
}
+
+object UserDefinedPythonTableFunction {
+
+ private[this] val workerModule = "pyspark.sql.worker.analyze_udtf"
+
+ /**
+ * Runs the Python UDTF's `analyze` static method.
+ */
+ def analyzeInPython(func: PythonFunction, exprs: Seq[Expression]): StructType = {
+ val env = SparkEnv.get
+ val bufferSize: Int = env.conf.get(BUFFER_SIZE)
+ val authSocketTimeout = env.conf.get(PYTHON_AUTH_SOCKET_TIMEOUT)
+ val reuseWorker = env.conf.get(PYTHON_WORKER_REUSE)
+ val simplifiedTraceback: Boolean = SQLConf.get.pysparkSimplifiedTraceback
+
+ val envVars = new HashMap[String, String](func.envVars)
+ val pythonExec = func.pythonExec
+ val pythonVer = func.pythonVer
+
+ if (reuseWorker) {
+ envVars.put("SPARK_REUSE_WORKER", "1")
+ }
+ if (simplifiedTraceback) {
+ envVars.put("SPARK_SIMPLIFIED_TRACEBACK", "1")
+ }
+ envVars.put("SPARK_AUTH_SOCKET_TIMEOUT", authSocketTimeout.toString)
+ envVars.put("SPARK_BUFFER_SIZE", bufferSize.toString)
+
+ EvaluatePython.registerPicklers()
+ val pickler = new Pickler(/* useMemo = */ true,
+ /* valueCompare = */ false)
+
+ val (worker: Socket, _) =
+ env.createPythonWorker(pythonExec, workerModule, envVars.asScala.toMap)
+ val releasedOrClosed = new AtomicBoolean(false)
Review Comment:
Good point. Actually it doesn't need to be `AtomicBoolean`.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267466437
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -748,6 +769,442 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(func().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, DataType)
+ assert a.value is not None
+ assert a.is_table is False
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
+ # another data type
+ (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
+ # array type
+ (
+ func(array(lit(1), lit(2), lit(3))),
+ StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
+ [Row(a=[1, 2, 3])],
+ ),
+ # map type
+ (
+ func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
+ ),
+ [Row(a={"x": 1, "y": 2})],
+ ),
+ # struct type
+ (
+ func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a",
+ StructType()
+ .add("x", IntegerType(), nullable=False)
+ .add("y", IntegerType(), nullable=False),
+ ),
+ [Row(a=Row(x=1, y=2))],
+ ),
+ # use SQL
+ (
+ self.spark.sql("SELECT * from test_udtf(1)"),
+ StructType().add("a", IntegerType()),
+ [Row(a=1)],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_decorator(self):
+ @udtf
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_decorator_parens(self):
+ @udtf()
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("a", a.data_type).add("b", b.data_type))
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1), lit("x")),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_arbitary_number_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(*args: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(
+ StructType([StructField(f"col{i}", a.data_type) for i, a in enumerate(args)])
+ )
+
+ def eval(self, *args):
+ yield args
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1)),
+ StructType().add("col0", IntegerType()),
+ [Row(a=1)],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("col0", IntegerType()).add("col1", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ # TODO(SPARK-44479): Support Python UDTFs with empty schema
+ # (func(), StructType(), [Row()]),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, StructType)
+ assert a.value is None
+ assert a.is_table is True
+ return AnalyzeResult(StructType().add("a", a.data_type[0].dataType))
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_analyze_table_argument_adding_columns(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a.data_type, StructType)
+ assert a.is_table is True
+ return AnalyzeResult(a.data_type.add("is_even", BooleanType()))
+
+ def eval(self, a: Row):
+ yield a["id"], a["id"] % 2 == 0
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(
+ df.schema,
+ StructType().add("id", LongType(), nullable=False).add("is_even", BooleanType()),
+ )
+ self.assertEqual(
+ df.collect(),
+ [
+ Row(a=0, is_even=True),
+ Row(a=1, is_even=False),
+ Row(a=2, is_even=True),
+ Row(a=3, is_even=False),
+ ],
+ )
+
+ def test_udtf_with_analyze_table_argument_repeating_rows(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(n, row) -> AnalyzeResult:
+ if n.value is None or not isinstance(n.value, int) or (n.value < 1 or n.value > 10):
+ raise Exception("The first argument must be a scalar integer between 1 and 10")
+
+ if row.is_table is False:
+ raise Exception("The second argument must be a table argument")
+
+ assert isinstance(row.data_type, StructType)
+ return AnalyzeResult(row.data_type)
+
+ def eval(self, n: int, row: Row):
+ for _ in range(n):
+ yield row
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected_schema = StructType().add("id", LongType(), nullable=False)
+ expected_results = [
+ Row(a=0),
+ Row(a=0),
+ Row(a=1),
+ Row(a=1),
+ Row(a=2),
+ Row(a=2),
+ Row(a=3),
+ Row(a=3),
+ ]
+ for i, df in enumerate(
+ [
+ self.spark.sql("SELECT * FROM test_udtf(2, TABLE (SELECT id FROM range(0, 4)))"),
+ self.spark.sql(
+ "SELECT * FROM test_udtf(1 + 1, TABLE (SELECT id FROM range(0, 4)))"
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ with self.assertRaisesRegex(
+ AnalysisException, "The first argument must be a scalar integer between 1 and 10"
+ ):
+ self.spark.sql(
+ "SELECT * FROM test_udtf(0, TABLE (SELECT id FROM range(0, 4)))"
+ ).collect()
+
+ with self.sql_conf(
+ {"spark.sql.tvf.allowMultipleTableArguments.enabled": True}
+ ), self.assertRaisesRegex(
+ AnalysisException, "The first argument must be a scalar integer between 1 and 10"
+ ):
+ self.spark.sql(
+ """
+ SELECT * FROM test_udtf(
+ TABLE (SELECT id FROM range(0, 1)),
+ TABLE (SELECT id FROM range(0, 4)))
+ """
+ ).collect()
+
+ with self.assertRaisesRegex(
+ AnalysisException, "The second argument must be a table argument"
+ ):
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')").collect()
+
+ def test_udtf_with_both_return_type_and_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF, returnType="c1: string, c2: string")
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_neither_return_type_nor_analyze(self):
+ class TestUDTF:
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
+ message_parameters={"name": "TestUDTF"},
+ )
+
+ def test_udtf_with_analyze_non_staticmethod(self):
+ class TestUDTF:
+ def analyze(self) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ with self.assertRaises(PySparkAttributeError) as e:
+ udtf(TestUDTF)
+
+ self.check_error(
+ exception=e.exception,
+ error_class="INVALID_UDTF_RETURN_TYPE",
Review Comment:
It IS checking if the UDTF class has an `analyze` method and throw an error if it is not static, and the error message will be like:
```
The UDTF '<name>' is invalid. It does not specify its return type or implement the required 'analyze' static method. Please specify the return type or implement the 'analyze' static method in '<name>' and try again.
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267467229
##########
python/pyspark/sql/tests/test_udtf.py:
##########
@@ -748,6 +769,442 @@ def terminate(self):
self.assertIn("Evaluate the input row", cls.eval.__doc__)
self.assertIn("Terminate the UDTF", cls.terminate.__doc__)
+ def test_simple_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(func().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, DataType)
+ assert a.value is not None
+ assert a.is_table is False
+ return AnalyzeResult(StructType().add("a", a.data_type))
+
+ def eval(self, a):
+ yield a,
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (func(lit(1)), StructType().add("a", IntegerType()), [Row(a=1)]),
+ # another data type
+ (func(lit("x")), StructType().add("a", StringType()), [Row(a="x")]),
+ # array type
+ (
+ func(array(lit(1), lit(2), lit(3))),
+ StructType().add("a", ArrayType(IntegerType(), containsNull=False)),
+ [Row(a=[1, 2, 3])],
+ ),
+ # map type
+ (
+ func(create_map(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a", MapType(StringType(), IntegerType(), valueContainsNull=False)
+ ),
+ [Row(a={"x": 1, "y": 2})],
+ ),
+ # struct type
+ (
+ func(named_struct(lit("x"), lit(1), lit("y"), lit(2))),
+ StructType().add(
+ "a",
+ StructType()
+ .add("x", IntegerType(), nullable=False)
+ .add("y", IntegerType(), nullable=False),
+ ),
+ [Row(a=Row(x=1, y=2))],
+ ),
+ # use SQL
+ (
+ self.spark.sql("SELECT * from test_udtf(1)"),
+ StructType().add("a", IntegerType()),
+ [Row(a=1)],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_decorator(self):
+ @udtf
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_decorator_parens(self):
+ @udtf()
+ class TestUDTF:
+ @staticmethod
+ def analyze() -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("c1", StringType()).add("c2", StringType()))
+
+ def eval(self):
+ yield "hello", "world"
+
+ self.spark.udtf.register("test_udtf", TestUDTF)
+
+ expected = [Row(c1="hello", c2="world")]
+ self.assertEqual(TestUDTF().collect(), expected)
+ self.assertEqual(self.spark.sql("SELECT * FROM test_udtf()").collect(), expected)
+
+ def test_udtf_with_analyze_multiple_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument, b: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(StructType().add("a", a.data_type).add("b", b.data_type))
+
+ def eval(self, a, b):
+ yield a, b
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1), lit("x")),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("a", IntegerType()).add("b", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_arbitary_number_arguments(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(*args: AnalyzeArgument) -> AnalyzeResult:
+ return AnalyzeResult(
+ StructType([StructField(f"col{i}", a.data_type) for i, a in enumerate(args)])
+ )
+
+ def eval(self, *args):
+ yield args
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ for i, (df, expected_schema, expected_results) in enumerate(
+ [
+ (
+ func(lit(1)),
+ StructType().add("col0", IntegerType()),
+ [Row(a=1)],
+ ),
+ (
+ self.spark.sql("SELECT * FROM test_udtf(1, 'x')"),
+ StructType().add("col0", IntegerType()).add("col1", StringType()),
+ [Row(a=1, b="x")],
+ ),
+ # TODO(SPARK-44479): Support Python UDTFs with empty schema
+ # (func(), StructType(), [Row()]),
+ ]
+ ):
+ with self.subTest(query_no=i):
+ self.assertEqual(df.schema, expected_schema)
+ self.assertEqual(df.collect(), expected_results)
+
+ def test_udtf_with_analyze_table_argument(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a, AnalyzeArgument)
+ assert isinstance(a.data_type, StructType)
+ assert a.value is None
+ assert a.is_table is True
+ return AnalyzeResult(StructType().add("a", a.data_type[0].dataType))
+
+ def eval(self, a: Row):
+ if a["id"] > 5:
+ yield a["id"],
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 8)))")
+ self.assertEqual(df.schema, StructType().add("a", LongType()))
+ self.assertEqual(df.collect(), [Row(a=6), Row(a=7)])
+
+ def test_udtf_with_analyze_table_argument_adding_columns(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(a: AnalyzeArgument) -> AnalyzeResult:
+ assert isinstance(a.data_type, StructType)
+ assert a.is_table is True
+ return AnalyzeResult(a.data_type.add("is_even", BooleanType()))
+
+ def eval(self, a: Row):
+ yield a["id"], a["id"] % 2 == 0
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ df = self.spark.sql("SELECT * FROM test_udtf(TABLE (SELECT id FROM range(0, 4)))")
+ self.assertEqual(
+ df.schema,
+ StructType().add("id", LongType(), nullable=False).add("is_even", BooleanType()),
+ )
+ self.assertEqual(
+ df.collect(),
+ [
+ Row(a=0, is_even=True),
+ Row(a=1, is_even=False),
+ Row(a=2, is_even=True),
+ Row(a=3, is_even=False),
+ ],
+ )
+
+ def test_udtf_with_analyze_table_argument_repeating_rows(self):
+ class TestUDTF:
+ @staticmethod
+ def analyze(n, row) -> AnalyzeResult:
+ if n.value is None or not isinstance(n.value, int) or (n.value < 1 or n.value > 10):
+ raise Exception("The first argument must be a scalar integer between 1 and 10")
+
+ if row.is_table is False:
+ raise Exception("The second argument must be a table argument")
+
+ assert isinstance(row.data_type, StructType)
+ return AnalyzeResult(row.data_type)
+
+ def eval(self, n: int, row: Row):
+ for _ in range(n):
+ yield row
+
+ func = udtf(TestUDTF)
+ self.spark.udtf.register("test_udtf", func)
+
+ expected_schema = StructType().add("id", LongType(), nullable=False)
+ expected_results = [
+ Row(a=0),
+ Row(a=0),
+ Row(a=1),
+ Row(a=1),
+ Row(a=2),
+ Row(a=2),
+ Row(a=3),
+ Row(a=3),
+ ]
+ for i, df in enumerate(
+ [
+ self.spark.sql("SELECT * FROM test_udtf(2, TABLE (SELECT id FROM range(0, 4)))"),
+ self.spark.sql(
+ "SELECT * FROM test_udtf(1 + 1, TABLE (SELECT id FROM range(0, 4)))"
Review Comment:
It's mentioned in the doc:
https://github.com/apache/spark/blob/0e268d149ae3b952a4cf2e97d5aeb389766106f0/python/pyspark/sql/functions.py#L15576
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267107024
##########
python/pyspark/errors/error_classes.py:
##########
@@ -268,6 +268,11 @@
"Eval type for UDF must be <eval_type>."
]
},
+ "INVALID_UDTF_BOTH_RETURN_TYPE_AND_ANALYZE_STATICMETHOD" : {
+ "message" : [
+ "The UDTF '<name>' is invalid. It has both its return type and the required 'analyze' static method. Please make it have one of either the return type or the 'analyze' static method in '<name>' and try again."
Review Comment:
Yes, it fails when a Python UDTF has both `returnType` and `analyze` static method, to avoid confusions.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267278691
##########
python/pyspark/sql/functions.py:
##########
@@ -15564,6 +15566,36 @@ def udtf(
| 1| 2|
+---+---+
+ UDTF can also have `analyze` static method instead of a static return type:
+
+ The `analyze` static method should take arguments:
+
+ - The number and order of arguments are the same as the UDTF inputs
+ - Each argument is a :class:`pyspark.sql.udtf.AnalyzeArgument`, containing:
+ - data_type: DataType
+ - value: Any: if the argument is foldable; otherwise None
+ - is_table: bool: True if the argument is TABLE
+
+ and return a :class:`pyspark.sql.udtf.AnalyzeResult`, containing.
Review Comment:
Added the import in the example.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
[GitHub] [spark] ueshin commented on a diff in pull request #41948: [SPARK-44380][SQL][PYTHON] Support for Python UDTF to analyze in Python
Posted by "ueshin (via GitHub)" <gi...@apache.org>.
ueshin commented on code in PR #41948:
URL: https://github.com/apache/spark/pull/41948#discussion_r1267274406
##########
python/pyspark/sql/functions.py:
##########
@@ -50,6 +50,7 @@
# Keep UserDefinedFunction import for backwards compatible import; moved in SPARK-22409
from pyspark.sql.udf import UserDefinedFunction, _create_py_udf # noqa: F401
from pyspark.sql.udtf import UserDefinedTableFunction, _create_py_udtf
+from pyspark.sql.udtf import AnalyzeArgument, AnalyzeResult # noqa: F401
Review Comment:
Sure, updated.
For F401, yes, by putting them here, users can import these symbols directly from `functions` pacakge.
```py
from pyspark.sql.functions import udtf, AnalyzeArgument, AnalyzeResult
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org