Skip to content

[SPARK-24324][PYTHON] Pandas Grouped Map UDF should assign result columns by name #21427

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Closed
BryanCutler wants to merge 15 commits into from
+226 −55
Closed

[SPARK-24324][PYTHON] Pandas Grouped Map UDF should assign result columns by name #21427

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
21e0c3e
added test for diff column order
BryanCutler May 24, 2018
bbe3587
needed to adjust expected values to compare results
BryanCutler May 24, 2018
a653e9b
for grouped map results, get columns based on name instead of position
BryanCutler May 24, 2018
88e2aa3
added test for positional columns
BryanCutler May 25, 2018
7cc0c49
adjust positional test
BryanCutler May 25, 2018
d4b5da1
add fallback to positional assignment
BryanCutler May 25, 2018
63c3963
added test for col name typo
BryanCutler May 25, 2018
9bbf014
fix test for py3
BryanCutler May 25, 2018
74c5d8e
added tests for integer index, incorrect field type with range index
BryanCutler May 29, 2018
b2d0966
documented behavior of column labels
BryanCutler May 29, 2018
5a7edb2
passing conf map to runner, tests pass
BryanCutler Jun 18, 2018
59972d6
added conf for assignment by position, passing general map of conf to…
BryanCutler Jun 18, 2018
27b4cad
fix some typos
BryanCutler Jun 18, 2018
c593650
address nits
BryanCutler Jun 22, 2018
2d2ced6
changed from catching KeyError to explicit check for string labels
BryanCutler Jun 22, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
address nits
  • Loading branch information
@BryanCutler
BryanCutler committed Jun 22, 2018
commit c593650aa527241da1ddd7e433c626db0fa26bb5
4 changes: 2 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/python/AggregateInPandasExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ case class AggregateInPandasExec(
val bufferSize = inputRDD.conf.getInt("spark.buffer.size", 65536)
val reuseWorker = inputRDD.conf.getBoolean("spark.python.worker.reuse", defaultValue = true)
val sessionLocalTimeZone = conf.sessionLocalTimeZone
val runnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)

val (pyFuncs, inputs) = udfExpressions.map(collectFunctions).unzip

Expand Down Expand Up @@ -143,7 +143,7 @@ case class AggregateInPandasExec(
argOffsets,
aggInputSchema,
sessionLocalTimeZone,
runnerConf).compute(projectedRowIter, context.partitionId(), context)
pythonRunnerConf).compute(projectedRowIter, context.partitionId(), context)

val joinedAttributes =
groupingExpressions.map(_.toAttribute) ++ udfExpressions.map(_.resultAttribute)
Expand Down
4 changes: 2 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowEvalPythonExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ case class ArrowEvalPythonExec(udfs: Seq[PythonUDF], output: Seq[Attribute], chi

private val batchSize = conf.arrowMaxRecordsPerBatch
private val sessionLocalTimeZone = conf.sessionLocalTimeZone
private val runnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
private val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)

protected override def evaluate(
funcs: Seq[ChainedPythonFunctions],
Expand All @@ -88,7 +88,7 @@ case class ArrowEvalPythonExec(udfs: Seq[PythonUDF], output: Seq[Attribute], chi
argOffsets,
schema,
sessionLocalTimeZone,
runnerConf).compute(batchIter, context.partitionId(), context)
pythonRunnerConf).compute(batchIter, context.partitionId(), context)

new Iterator[InternalRow] {

Expand Down
3 changes: 3 additions & 0 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowPythonRunner.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,11 +58,14 @@ class ArrowPythonRunner(
new WriterThread(env, worker, inputIterator, partitionIndex, context) {

protected override def writeCommand(dataOut: DataOutputStream): Unit = {

// Write config for the worker as a number of key -> value pairs of strings
dataOut.writeInt(conf.size)
Copy link
Contributor

@icexelloss icexelloss Jun 19, 2018 •
edited
Loading

Choose a reason for hiding this comment

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

maybe put this in a writeConf method to be more specific?

Copy link
Member Author

@BryanCutler BryanCutler Jun 22, 2018

Choose a reason for hiding this comment

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

I think it's fine, but I will add some comments

Copy link
Contributor

@icexelloss icexelloss Jun 22, 2018

Choose a reason for hiding this comment

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

Ok, SGTM.

for ((k, v) <- conf) {
PythonRDD.writeUTF(k, dataOut)
PythonRDD.writeUTF(v, dataOut)
}

PythonUDFRunner.writeUDFs(dataOut, funcs, argOffsets)
}

Expand Down
4 changes: 2 additions & 2 deletions ...core/src/main/scala/org/apache/spark/sql/execution/python/FlatMapGroupsInPandasExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,7 @@ case class FlatMapGroupsInPandasExec(
val reuseWorker = inputRDD.conf.getBoolean("spark.python.worker.reuse", defaultValue = true)
val chainedFunc = Seq(ChainedPythonFunctions(Seq(pandasFunction)))
val sessionLocalTimeZone = conf.sessionLocalTimeZone
val runnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)

// Deduplicate the grouping attributes.
// If a grouping attribute also appears in data attributes, then we don't need to send the
Expand Down Expand Up @@ -147,7 +147,7 @@ case class FlatMapGroupsInPandasExec(
argOffsets,
dedupSchema,
sessionLocalTimeZone,
runnerConf).compute(grouped, context.partitionId(), context)
pythonRunnerConf).compute(grouped, context.partitionId(), context)

columnarBatchIter.flatMap(_.rowIterator.asScala).map(UnsafeProjection.create(output, output))
}
Expand Down
4 changes: 2 additions & 2 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/python/WindowInPandasExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ case class WindowInPandasExec(
val bufferSize = inputRDD.conf.getInt("spark.buffer.size", 65536)
val reuseWorker = inputRDD.conf.getBoolean("spark.python.worker.reuse", defaultValue = true)
val sessionLocalTimeZone = conf.sessionLocalTimeZone
val runnerConf = ArrowUtils.getPythonRunnerConfMap(conf)
val pythonRunnerConf = ArrowUtils.getPythonRunnerConfMap(conf)

// Extract window expressions and window functions
val expressions = windowExpression.flatMap(_.collect { case e: WindowExpression => e })
Expand Down Expand Up @@ -162,7 +162,7 @@ case class WindowInPandasExec(
argOffsets,
windowInputSchema,
sessionLocalTimeZone,
runnerConf).compute(pythonInput, context.partitionId(), context)
pythonRunnerConf).compute(pythonInput, context.partitionId(), context)

val joined = new JoinedRow
val resultProj = createResultProjection(expressions)
Expand Down