foreachPartition（）自定義函數中的PySpark Access DataFrame列

Question

我有一個名為“內部”的函數。 我想將此功能應用於pyspark數據框。 為此，我在創建的數據幀上調用“ foreachPartition（inside）”方法。 “內部”功能需要數據框的值。

數據框如下所示：

>>> small_df
DataFrame[lon: double, lat: double, t: bigint]

代碼如下：

def inside(iterator):
    row=iterator
    x=row.lon
    y=row.lat
    i=row.t 
    #do more stuff

small=pliades.iloc[0:20000,:] #take sample of rows from big dataset
small_df=sqlContext.createDataFrame(small) #create dataframe
test=small_df.foreachPartition(inside)

我的問題是：x，y，i如何分別獲取數據幀的第一列（lon），第二列（lat）和第三列（t）的值？

我也嘗試使用row.lon，row.select進行處理，將其視為列表，但無法獲得所需的結果。

Answer 1

foreach在RDD[Row] ，每個分區都是Iterator[Row] 。 如果您想獲得所有值的列表（由於可能的內存問題，不建議使用

def inside(iterator):
    x, y, i = zip(*iterator)
    ...
    yield ...

通常，最好只逐行遍歷所有行，而不將所有行保留在內存中：

def inside(iterator):
    for x, y, i in iterator:
        yield ...

您也可以考慮使用pandas_udf ：

• 如果函數返回相同數量的值並且只有一個列，則可以使用標量類型，該標量類型采用pandas.Series並返回pandas.Series

   from pyspark.sql.functions import pandas_udf, PandasUDFType @pandas_udf(schema, PandasUDFType.SCALAR) def f(*cols: pandas.Series) -> pandas.Series: ... df.select(f("col1", "col2", ...)) 

• 分組后的變體，采用pandas.DataFrame並返回具有相同或不同行數的pandas.DataFrame ：

   from pyspark.sql.functions import spark_partition_id @pandas_udf(schema, PandasUDFType.GROUPED_MAP) def g(df: pandas.DataFrame) -> pandas.DataFrame: ... df.groupby(spark_partition_id()).apply(g)