diff --git a/connector/connect/common/src/main/protobuf/spark/connect/relations.proto b/connector/connect/common/src/main/protobuf/spark/connect/relations.proto index c965a6c8d32f4..87454dd00244a 100644 --- a/connector/connect/common/src/main/protobuf/spark/connect/relations.proto +++ b/connector/connect/common/src/main/protobuf/spark/connect/relations.proto @@ -795,7 +795,7 @@ message MapPartitions { // (Required) Input user-defined function. CommonInlineUserDefinedFunction func = 2; - // (Optional) isBarrier. + // (Optional) Whether to use barrier mode execution or not. optional bool is_barrier = 3; } diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py index 65f270f21d48c..94daaa4d57d7f 100644 --- a/python/pyspark/sql/connect/dataframe.py +++ b/python/pyspark/sql/connect/dataframe.py @@ -1645,7 +1645,7 @@ def _map_partitions( func: "PandasMapIterFunction", schema: Union[StructType, str], evalType: int, - is_barrier: bool, + barrier: bool, ) -> "DataFrame": from pyspark.sql.connect.udf import UserDefinedFunction @@ -1660,7 +1660,7 @@ def _map_partitions( return DataFrame.withPlan( plan.MapPartitions( - child=self._plan, function=udf_obj, cols=self.columns, is_barrier=is_barrier + child=self._plan, function=udf_obj, cols=self.columns, is_barrier=barrier ), session=self._session, ) @@ -1669,11 +1669,9 @@ def mapInPandas( self, func: "PandasMapIterFunction", schema: Union[StructType, str], - is_barrier: bool = False, + barrier: bool = False, ) -> "DataFrame": - return self._map_partitions( - func, schema, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, is_barrier - ) + return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_PANDAS_ITER_UDF, barrier) mapInPandas.__doc__ = PySparkDataFrame.mapInPandas.__doc__ @@ -1681,9 +1679,9 @@ def mapInArrow( self, func: "ArrowMapIterFunction", schema: Union[StructType, str], - is_barrier: bool = False, + barrier: bool = False, ) -> "DataFrame": - return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_ARROW_ITER_UDF, is_barrier) + return self._map_partitions(func, schema, PythonEvalType.SQL_MAP_ARROW_ITER_UDF, barrier) mapInArrow.__doc__ = PySparkDataFrame.mapInArrow.__doc__ diff --git a/python/pyspark/sql/connect/proto/relations_pb2.pyi b/python/pyspark/sql/connect/proto/relations_pb2.pyi index c3cf733a995b1..cd2b0fb2d3ff5 100644 --- a/python/pyspark/sql/connect/proto/relations_pb2.pyi +++ b/python/pyspark/sql/connect/proto/relations_pb2.pyi @@ -2767,7 +2767,7 @@ class MapPartitions(google.protobuf.message.Message): def func(self) -> pyspark.sql.connect.proto.expressions_pb2.CommonInlineUserDefinedFunction: """(Required) Input user-defined function.""" is_barrier: builtins.bool - """(Optional) isBarrier.""" + """(Optional) Whether to use barrier mode execution or not.""" def __init__( self, *,