DataFrame 化的 zipWithIndex答案

【问题标题】：DataFrame-ified zipWithIndexDataFrame 化的 zipWithIndex
【发布时间】：2015-07-30 01:15:39
【问题描述】：

我正在尝试解决向数据集添加序列号的古老问题。我正在使用 DataFrame，似乎没有与 RDD.zipWithIndex 等效的 DataFrame。另一方面，以下工作或多或少符合我的要求：

val origDF = sqlContext.load(...)    

val seqDF= sqlContext.createDataFrame(
    origDF.rdd.zipWithIndex.map(ln => Row.fromSeq(Seq(ln._2) ++ ln._1.toSeq)),
    StructType(Array(StructField("seq", LongType, false)) ++ origDF.schema.fields)
)

在我的实际应用程序中，origDF 不会直接从文件中加载——它将通过将 2-3 个其他 DataFrame 连接在一起来创建，并且将包含超过 1 亿行。

有没有更好的方法来做到这一点？我可以做些什么来优化它？

【问题讨论】：

标签： scala apache-spark apache-spark-sql

【解决方案1】：

这是我的建议，其优点是：

它不涉及我们DataFrame的InternalRows的任何序列化/反序列化^[1]。
其逻辑极简，仅依赖RDD.zipWithIndex。

它的主要缺点是：

无法直接从非 JVM API（pySpark、SparkR）中使用它。
必须在package org.apache.spark.sql; 下。

进口：

import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.execution.LogicalRDD
import org.apache.spark.sql.functions.lit

/**
  * Optimized Spark SQL equivalent of RDD.zipWithIndex.
  *
  * @param df
  * @param indexColName
  * @return `df` with a column named `indexColName` of consecutive unique ids.
  */
def zipWithIndex(df: DataFrame, indexColName: String = "index"): DataFrame = {
  import df.sparkSession.implicits._

  val dfWithIndexCol: DataFrame = df
    .drop(indexColName)
    .select(lit(0L).as(indexColName), $"*")

  val internalRows: RDD[InternalRow] = dfWithIndexCol
    .queryExecution
    .toRdd
    .zipWithIndex()
    .map {
      case (internalRow: InternalRow, index: Long) =>
        internalRow.setLong(0, index)
        internalRow
    }

  Dataset.ofRows(
    df.sparkSession,
    LogicalRDD(dfWithIndexCol.schema.toAttributes, internalRows)(df.sparkSession)
  )

^[1]：（从/到InternalRow 的底层字节数组 GenericRow 的底层JVM 对象集合Array[Any]）。

【讨论】：

【解决方案2】：

我已经修改了@Tagar 的版本以在 Python 3.7 上运行，想要分享：

def dfZipWithIndex (df, offset=1, colName="rowId"):
'''
    Enumerates dataframe rows is native order, like rdd.ZipWithIndex(), but on a dataframe
    and preserves a schema

    :param df: source dataframe
    :param offset: adjustment to zipWithIndex()'s index
    :param colName: name of the index column
'''

new_schema = StructType(
                [StructField(colName,LongType(),True)]        # new added field in front
                + df.schema.fields                            # previous schema
            )

zipped_rdd = df.rdd.zipWithIndex()

new_rdd = zipped_rdd.map(lambda args: ([args[1] + offset] + list(args[0])))      # use this for python 3+, tuple gets passed as single argument so using args and [] notation to read elements within args
return spark.createDataFrame(new_rdd, new_schema)

【讨论】：

【解决方案3】：

Spark Java API 版本：

我已经实现了@Evgeny 的solution，用于在Java 中对DataFrames 执行zipWithIndex，并希望分享代码。

它还包含@fylb 在他的solution 中提供的改进。我可以确认 Spark 2.4 当 spark_partition_id() 返回的条目不以 0 开头或不按顺序增加时执行失败。由于此函数documented 是非确定性的，因此很可能会发生上述情况之一。一个示例是通过增加分区计数来触发的。

java实现如下：

public static Dataset<Row> zipWithIndex(Dataset<Row> df, Long offset, String indexName) {
        Dataset<Row> dfWithPartitionId = df
                .withColumn("partition_id", spark_partition_id())
                .withColumn("inc_id", monotonically_increasing_id());

        Object partitionOffsetsObject = dfWithPartitionId
                .groupBy("partition_id")
                .agg(count(lit(1)).alias("cnt"), first("inc_id").alias("inc_id"))
                .orderBy("partition_id")
                .select(col("partition_id"), sum("cnt").over(Window.orderBy("partition_id")).minus(col("cnt")).minus(col("inc_id")).plus(lit(offset).alias("cnt")))
                .collect();
        Row[] partitionOffsetsArray = ((Row[]) partitionOffsetsObject);
        Map<Integer, Long> partitionOffsets = new HashMap<>();
        for (int i = 0; i < partitionOffsetsArray.length; i++) {
            partitionOffsets.put(partitionOffsetsArray[i].getInt(0), partitionOffsetsArray[i].getLong(1));
        }

        UserDefinedFunction getPartitionOffset = udf(
                (partitionId) -> partitionOffsets.get((Integer) partitionId), DataTypes.LongType
        );

        return dfWithPartitionId
                .withColumn("partition_offset", getPartitionOffset.apply(col("partition_id")))
                .withColumn(indexName, col("partition_offset").plus(col("inc_id")))
                .drop("partition_id", "partition_offset", "inc_id");
    }

【讨论】：

【解决方案4】：

@Evgeny ，your solution 很有趣。请注意，当您有空分区时存在一个错误（数组缺少这些分区索引，至少这在 spark 1.6 中发生在我身上），所以我将数组转换为 Map(partitionId -> offsets)。

另外，我把 monotonically_increasing_id 的来源取出来让每个分区的“inc_id”从 0 开始。

这是一个更新版本：

import org.apache.spark.sql.catalyst.expressions.LeafExpression
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.types.LongType
import org.apache.spark.sql.catalyst.expressions.Nondeterministic
import org.apache.spark.sql.catalyst.expressions.codegen.GeneratedExpressionCode
import org.apache.spark.sql.catalyst.expressions.codegen.CodeGenContext
import org.apache.spark.sql.types.DataType
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions._
import org.apache.spark.sql.Column
import org.apache.spark.sql.expressions.Window

case class PartitionMonotonicallyIncreasingID() extends LeafExpression with Nondeterministic {

  /**
   * From org.apache.spark.sql.catalyst.expressions.MonotonicallyIncreasingID
   *
   * Record ID within each partition. By being transient, count's value is reset to 0 every time
   * we serialize and deserialize and initialize it.
   */
  @transient private[this] var count: Long = _

  override protected def initInternal(): Unit = {
    count = 1L // notice this starts at 1, not 0 as in org.apache.spark.sql.catalyst.expressions.MonotonicallyIncreasingID
  }

  override def nullable: Boolean = false

  override def dataType: DataType = LongType

  override protected def evalInternal(input: InternalRow): Long = {
    val currentCount = count
    count += 1
    currentCount
  }

  override def genCode(ctx: CodeGenContext, ev: GeneratedExpressionCode): String = {
    val countTerm = ctx.freshName("count")
    ctx.addMutableState(ctx.JAVA_LONG, countTerm, s"$countTerm = 1L;")
    ev.isNull = "false"
    s"""
      final ${ctx.javaType(dataType)} ${ev.value} = $countTerm;
      $countTerm++;
    """
  }
}

object DataframeUtils {
  def zipWithIndex(df: DataFrame, offset: Long = 0, indexName: String = "index") = {
    // from https://stackoverflow.com/questions/30304810/dataframe-ified-zipwithindex)
    val dfWithPartitionId = df.withColumn("partition_id", spark_partition_id()).withColumn("inc_id", new Column(PartitionMonotonicallyIncreasingID()))

    // collect each partition size, create the offset pages
    val partitionOffsets: Map[Int, Long] = dfWithPartitionId
      .groupBy("partition_id")
      .agg(max("inc_id") as "cnt") // in each partition, count(inc_id) is equal to max(inc_id) (I don't know which one would be faster)
      .select(col("partition_id"), sum("cnt").over(Window.orderBy("partition_id")) - col("cnt") + lit(offset) as "cnt")
      .collect()
      .map(r => (r.getInt(0) -> r.getLong(1)))
      .toMap

    def partition_offset(partitionId: Int): Long = partitionOffsets(partitionId)
    val partition_offset_udf = udf(partition_offset _)
    // and re-number the index
    dfWithPartitionId
      .withColumn("partition_offset", partition_offset_udf(col("partition_id")))
      .withColumn(indexName, col("partition_offset") + col("inc_id"))
      .drop("partition_id")
      .drop("partition_offset")
      .drop("inc_id")
  }
}

【讨论】：

【解决方案5】：

从 Spark 1.6 开始，有一个名为 monotonically_increasing_id() 的函数
它为每一行生成一个具有唯一 64 位单调索引的新列
但这不是必然的，每个分区都开始一个新的范围，所以我们必须在使用它之前计算每个分区的偏移量。
试图提供一个“无rdd”的解决方案，我最终得到了一些collect()，但它只收集偏移量，每个分区一个值，所以不会导致OOM

def zipWithIndex(df: DataFrame, offset: Long = 1, indexName: String = "index") = {
    val dfWithPartitionId = df.withColumn("partition_id", spark_partition_id()).withColumn("inc_id", monotonically_increasing_id())

    val partitionOffsets = dfWithPartitionId
        .groupBy("partition_id")
        .agg(count(lit(1)) as "cnt", first("inc_id") as "inc_id")
        .orderBy("partition_id")
        .select(sum("cnt").over(Window.orderBy("partition_id")) - col("cnt") - col("inc_id") + lit(offset) as "cnt" )
        .collect()
        .map(_.getLong(0))
        .toArray
        
     dfWithPartitionId
        .withColumn("partition_offset", udf((partitionId: Int) => partitionOffsets(partitionId), LongType)(col("partition_id")))
        .withColumn(indexName, col("partition_offset") + col("inc_id"))
        .drop("partition_id", "partition_offset", "inc_id")
}

此解决方案不会重新打包原始行，也不会重新分区原始的巨大数据帧，因此在现实世界中它非常快： 200GB 的 CSV 数据（4300 万行，150 列）在 240 个内核上在 2 分钟内读取、索引并打包到 parquet
在测试我的解决方案后，我运行了Kirk Broadhurst's solution，它慢了 20 秒
您可能想要或不想使用dfWithPartitionId.cache()，取决于任务

【讨论】：

首先，如果某些分区没有df 的行，可能会出现错误。请检查@fylb 的答案，并解释如何解决它。其次，最后一个as "cnt" 在那里非常令人困惑，并且由于之后不再使用，因此可以并且值得将其删除。第三，感谢非常有用、有趣且相当优雅的回答！
假设我正在使用monotonically_increasing_id() 为“product_id”列创建索引，我的问题是这个函数是否保证相同的 product_id 在分区/节点之间具有相同的索引？
@SarahData 不，monotonically_increasing_id() 在所有行中都是绝对唯一的，它不依赖于任何其他列，因此具有相同product_id 的不同行将始终具有不同的monotonically_increasing_id()' s

【解决方案6】：

PySpark 版本：

from pyspark.sql.types import LongType, StructField, StructType

def dfZipWithIndex (df, offset=1, colName="rowId"):
    '''
        Enumerates dataframe rows is native order, like rdd.ZipWithIndex(), but on a dataframe 
        and preserves a schema

        :param df: source dataframe
        :param offset: adjustment to zipWithIndex()'s index
        :param colName: name of the index column
    '''

    new_schema = StructType(
                    [StructField(colName,LongType(),True)]        # new added field in front
                    + df.schema.fields                            # previous schema
                )

    zipped_rdd = df.rdd.zipWithIndex()

    new_rdd = zipped_rdd.map(lambda (row,rowId): ([rowId +offset] + list(row)))

    return spark.createDataFrame(new_rdd, new_schema)

还创建了一个 jira 以在 Spark 中本地添加此功能：https://issues.apache.org/jira/browse/SPARK-23074

【讨论】：

【解决方案7】：

以下内容是代表 David Griffin 发布的（编辑无问题）。

全唱全能的 dfZipWithIndex 方法。您可以设置起始偏移量（默认为 1）、索引列名（默认为“id”），并将列放在前面或后面：

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.types.{LongType, StructField, StructType}
import org.apache.spark.sql.Row


def dfZipWithIndex(
  df: DataFrame,
  offset: Int = 1,
  colName: String = "id",
  inFront: Boolean = true
) : DataFrame = {
  df.sqlContext.createDataFrame(
    df.rdd.zipWithIndex.map(ln =>
      Row.fromSeq(
        (if (inFront) Seq(ln._2 + offset) else Seq())
          ++ ln._1.toSeq ++
        (if (inFront) Seq() else Seq(ln._2 + offset))
      )
    ),
    StructType(
      (if (inFront) Array(StructField(colName,LongType,false)) else Array[StructField]()) 
        ++ df.schema.fields ++ 
      (if (inFront) Array[StructField]() else Array(StructField(colName,LongType,false)))
    )
  ) 
}

【讨论】：

@eliasah -- 我找到了一种Window 表达方式来做到这一点。但是，它要慢得多，但认为您可能想看一下。请参阅下面的答案。
太棒了。对 PySpark 版本的任何引用？感谢分享。
这看起来很流畅。谁能帮我用JAVA写这个

【解决方案8】：

从 Spark 1.5 开始，Window 表达式被添加到 Spark。您现在可以使用org.apache.spark.sql.expressions.row_number，而不必将DataFrame 转换为RDD。请注意，我发现上述dfZipWithIndex 的性能明显快于以下算法。但我发布它是因为：

其他人会很想试试这个
也许有人可以优化下面的表达式

无论如何，这对我有用：

import org.apache.spark.sql.expressions._

df.withColumn("row_num", row_number.over(Window.partitionBy(lit(1)).orderBy(lit(1))))

请注意，我使用lit(1) 进行分区和排序——这使得所有内容都在同一个分区中，并且似乎保留了DataFrame 的原始排序，但我想这是减慢速度的原因下来。

我在包含 7,000,000 行的 4 列 DataFrame 上对其进行了测试，这与上面的 dfZipWithIndex 之间的速度差异很大（就像我说的，RDD 函数要快得多）。

【讨论】：

如果数据集不适合单个工作人员的内存，那会不会导致 OOM 错误？
我不知道——我只知道它比基于RDD 的zipWithIndex 慢得多，这足以让我停止思考它。我发布了上面的内容，以便其他人不会想在这条路上走得太远；原来的dfZipWithIndex 似乎仍然是最好的方法。
感谢分享，我以为不把DF转RDD的方式一开始会更快，现在也不会走得太远。
这种方法有两个问题：1）窗口函数需要 order by 子句 - 所以数据集将被排序 - 这就是为什么我猜你发现 dfZipWithIndex 更快； 2) 有时您没有要排序的列；例如，您需要文件中的本机行顺序，因此 dfZipWithIndex 可能是唯一的选择..
使用spark 2.3.0，如果数据超过它的内存容量，worker不会溢出到磁盘吗？