python - 有没有办法在 PySpark 中为 pandas_udf 设置最小批量大小？

Question

我正在使用 pandas_udf 在我的 Spark 集群上应用机器学习模型，并且有兴趣预定义通过箭头发送到 UDF 的最小记录数。

我遵循了大部分 UDF 的 databricks 教程... https://docs.databricks.com/applications/deep-learning/inference/resnet-model-inference-tensorflow.html

在教程中，我将 spark session 设置为具有最大批量大小和启用箭头。我可以轻松设置最大批量大小，但是我想知道是否有类似的方法来设置 UDF 将处理的最小批量大小？

spark = SparkSession.builder.appName('App').getOrCreate()

spark.conf.set("spark.sql.execution.arrow.enabled", "true")

spark.conf.set('spark.sql.execution.arrow.maxRecordsPerBatch', PyArrowBatchSize)

我正在运行 spark 版本 2.4.3 和 python 3.6.0。

Answer 1

在 Spark docs 中没有办法设置最小批量大小，但在这种情况下 max有点误导。这应该类似于“余数前的批量大小”。

例如:如果您的数据集中有 100132 行，并且您的 maxRecordsPerBatch是 10000，那么您将得到 10 批大小为 10000 的批处理，以及一批大小为 132 的作为余数。 (如果你有多个执行者，你可能会有额外的批处理和剩余部分，这取决于事情是如何分割的。)

你可以知道你的大概min批量大小取决于您的余数，否则所有批量大小都将恰好是最小批量大小。