Skip to content

[SPARK-42896][SQL][PYTHON] Make mapInPandas / mapInArrow support barrier mode execution #40520

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
WeichenXu123 wants to merge 8 commits into from
Closed

[SPARK-42896][SQL][PYTHON] Make mapInPandas / mapInArrow support barrier mode execution #40520

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
4de9d07
init
WeichenXu123 Mar 22, 2023
f612bcb
update
WeichenXu123 Mar 22, 2023
e95fd95
update
WeichenXu123 Mar 22, 2023
aa7862f
update
WeichenXu123 Mar 23, 2023
9be4cad
update
WeichenXu123 Mar 24, 2023
9ebb3f2
Merge branch 'master' into barrier-udf
WeichenXu123 Mar 24, 2023
85b840b
fix test
WeichenXu123 Mar 24, 2023
66f13cb
update format
WeichenXu123 Mar 24, 2023
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
Next Next commit
init 
Signed-off-by: Weichen Xu <[email protected]>
  • Loading branch information
@WeichenXu123
WeichenXu123 committed Mar 22, 2023
commit 4de9d0791d7de4cdc042fd9f42c510f954c628da
8 changes: 4 additions & 4 deletions python/pyspark/sql/pandas/map_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ class PandasMapOpsMixin:
"""

def mapInPandas(
self, func: "PandasMapIterFunction", schema: Union[StructType, str]
self, func: "PandasMapIterFunction", schema: Union[StructType, str], is_barrier: bool
) -> "DataFrame":
"""
Maps an iterator of batches in the current :class:`DataFrame` using a Python native
Expand Down Expand Up @@ -93,11 +93,11 @@ def mapInPandas(
func, returnType=schema, functionType=PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
) # type: ignore[call-overload]
udf_column = udf(*[self[col] for col in self.columns])
jdf = self._jdf.mapInPandas(udf_column._jc.expr())
jdf = self._jdf.mapInPandas(udf_column._jc.expr(), is_barrier)
return DataFrame(jdf, self.sparkSession)

def mapInArrow(
self, func: "ArrowMapIterFunction", schema: Union[StructType, str]
self, func: "ArrowMapIterFunction", schema: Union[StructType, str], is_barrier: bool
) -> "DataFrame":
"""
Maps an iterator of batches in the current :class:`DataFrame` using a Python native
Expand Down Expand Up @@ -157,7 +157,7 @@ def mapInArrow(
func, returnType=schema, functionType=PythonEvalType.SQL_MAP_ARROW_ITER_UDF
) # type: ignore[call-overload]
udf_column = udf(*[self[col] for col in self.columns])
jdf = self._jdf.pythonMapInArrow(udf_column._jc.expr())
jdf = self._jdf.pythonMapInArrow(udf_column._jc.expr(), is_barrier)
return DataFrame(jdf, self.sparkSession)


Expand Down
6 changes: 4 additions & 2 deletions ...t/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/pythonLogicalOperators.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,8 @@ case class FlatMapGroupsInPandas(
case class MapInPandas(
functionExpr: Expression,
output: Seq[Attribute],
child: LogicalPlan) extends UnaryNode {
child: LogicalPlan,
isBarrier: Boolean) extends UnaryNode {

override val producedAttributes = AttributeSet(output)

Expand All @@ -68,7 +69,8 @@ case class MapInPandas(
case class PythonMapInArrow(
functionExpr: Expression,
output: Seq[Attribute],
child: LogicalPlan) extends UnaryNode {
child: LogicalPlan,
isBarrier: Boolean) extends UnaryNode {

override val producedAttributes = AttributeSet(output)

Expand Down
10 changes: 6 additions & 4 deletions sql/core/src/main/scala/org/apache/spark/sql/Dataset.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3283,27 +3283,29 @@ class Dataset[T] private[sql](
* This function uses Apache Arrow as serialization format between Java executors and Python
* workers.
*/
private[sql] def mapInPandas(func: PythonUDF): DataFrame = {
private[sql] def mapInPandas(func: PythonUDF, isBarrier: Boolean = false): DataFrame = {
Dataset.ofRows(
sparkSession,
MapInPandas(
func,
func.dataType.asInstanceOf[StructType].toAttributes,
logicalPlan))
logicalPlan,
isBarrier))
}

/**
* Applies a function to each partition in Arrow format. The user-defined function
* defines a transformation: `iter(pyarrow.RecordBatch)` -> `iter(pyarrow.RecordBatch)`.
* Each partition is each iterator consisting of `pyarrow.RecordBatch`s as batches.
*/
private[sql] def pythonMapInArrow(func: PythonUDF): DataFrame = {
private[sql] def pythonMapInArrow(func: PythonUDF, isBarrier: Boolean = false): DataFrame = {
Dataset.ofRows(
sparkSession,
PythonMapInArrow(
func,
func.dataType.asInstanceOf[StructType].toAttributes,
logicalPlan))
logicalPlan,
isBarrier))
}

/**
Expand Down
8 changes: 4 additions & 4 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/SparkStrategies.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -806,10 +806,10 @@ abstract class SparkStrategies extends QueryPlanner[SparkPlan] {
execution.python.FlatMapCoGroupsInPandasExec(
f.leftAttributes, f.rightAttributes,
func, output, planLater(left), planLater(right)) :: Nil
case logical.MapInPandas(func, output, child) =>
execution.python.MapInPandasExec(func, output, planLater(child)) :: Nil
case logical.PythonMapInArrow(func, output, child) =>
execution.python.PythonMapInArrowExec(func, output, planLater(child)) :: Nil
case logical.MapInPandas(func, output, child, isBarrier) =>
execution.python.MapInPandasExec(func, output, planLater(child), isBarrier) :: Nil
case logical.PythonMapInArrow(func, output, child, isBarrier) =>
execution.python.PythonMapInArrowExec(func, output, planLater(child), isBarrier) :: Nil
case logical.AttachDistributedSequence(attr, child) =>
execution.python.AttachDistributedSequenceExec(attr, planLater(child)) :: Nil
case logical.MapElements(f, _, _, objAttr, child) =>
Expand Down
9 changes: 8 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/python/MapInBatchExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ trait MapInBatchExec extends UnaryExecNode with PythonSQLMetrics {
protected val func: Expression
protected val pythonEvalType: Int

protected val isBarrier: Boolean

private val pythonFunction = func.asInstanceOf[PythonUDF].func

override def producedAttributes: AttributeSet = AttributeSet(output)
Expand All @@ -50,7 +52,7 @@ trait MapInBatchExec extends UnaryExecNode with PythonSQLMetrics {
override def outputPartitioning: Partitioning = child.outputPartitioning

override protected def doExecute(): RDD[InternalRow] = {
child.execute().mapPartitionsInternal { inputIter =>
val resultRDD = child.execute().mapPartitionsInternal { inputIter =>
// Single function with one struct.
val argOffsets = Array(Array(0))
val chainedFunc = Seq(ChainedPythonFunctions(Seq(pythonFunction)))
Expand Down Expand Up @@ -90,5 +92,10 @@ trait MapInBatchExec extends UnaryExecNode with PythonSQLMetrics {
flattenedBatch.rowIterator.asScala
}.map(unsafeProj)
}
if (isBarrier) {
resultRDD.barrier().mapPartitions(iter => iter)
} else {
resultRDD
}
}
}
3 changes: 2 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/python/MapInPandasExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ import org.apache.spark.sql.execution.SparkPlan
case class MapInPandasExec(
func: Expression,
output: Seq[Attribute],
child: SparkPlan)
child: SparkPlan,
override val isBarrier: Boolean)
Copy link
Member

@HyukjinKwon HyukjinKwon Mar 22, 2023

Choose a reason for hiding this comment

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

It would require an implementation much complicated than this actually. The current implementation might work for now because these specific physical plans don't move around much for now but the implementation is flaky because the physical plans can change via Catalyst Optimizer (e.g., predicate pushdown) but the barrier execution mode requires that the barrier is created exactly the location where it's invoked.

Copy link
Contributor Author

@WeichenXu123 WeichenXu123 Mar 22, 2023 •
edited
Loading

Choose a reason for hiding this comment

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

change via Catalyst Optimizer (e.g., predicate pushdown)

Do you mean pushdown like this ?

Filter
  MapInPandas
      XXOp
          ....

to

MapInPandas
  Filter
      XXOp
          ....

?

I don't think we should allow pushdown in this case. MapInPandas executes arbitrary user code. Can we modify optimizer code to prevent it change the plan when it finds "is_barrier" set ?

Copy link
Contributor Author

@WeichenXu123 WeichenXu123 Mar 22, 2023

Choose a reason for hiding this comment

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

The feature mapInPandas / mapInArrow supporting barrier mode is only for ML use case, we don't need to complicate it. We only need to support such scenario:

input_dataset.mapInPandas(..., is_barrier=True).collect()

Copy link
Member

@HyukjinKwon HyukjinKwon Mar 22, 2023

Choose a reason for hiding this comment

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

I meant that the optimizers can might switch the physical plan order in the future, and that's why the implementation is flaky. Now this works because we don't switch the order anywhere else.

Can we modify optimizer code to prevent it change the plan when it finds "is_barrier" set ?

To have a complete implementation, we should have something like #19873 which brought a lot of side effect, and it got reverted in the end.

Copy link
Member

@HyukjinKwon HyukjinKwon Mar 22, 2023

Choose a reason for hiding this comment

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

If this is only for internal purpose, let's probably don't expose is_barrier: bool = False to the API surface at least.

Copy link
Contributor Author

@WeichenXu123 WeichenXu123 Mar 22, 2023

Choose a reason for hiding this comment

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

To simply the implementation, we can implement a barrierMapInPandasAndCollect instead, and define a execution plan stage like BarrierMapInPandasAndCollectExec

Copy link
Contributor Author

@WeichenXu123 WeichenXu123 Mar 23, 2023

Choose a reason for hiding this comment

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

If this is only for internal purpose, let's probably don't expose is_barrier: bool = False to the API surface at least.

No. It should be public API. Third-party lib such as xgboost also need to use it.

extends MapInBatchExec {

override protected val pythonEvalType: Int = PythonEvalType.SQL_MAP_PANDAS_ITER_UDF
Expand Down
3 changes: 2 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/execution/python/PythonMapInArrowExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,8 @@ import org.apache.spark.sql.execution.SparkPlan
case class PythonMapInArrowExec(
func: Expression,
output: Seq[Attribute],
child: SparkPlan)
child: SparkPlan,
override val isBarrier: Boolean)
extends MapInBatchExec {

override protected val pythonEvalType: Int = PythonEvalType.SQL_MAP_ARROW_ITER_UDF
Expand Down