DataFrame-ified zipWithIndex

Ответ 1

Так как Spark 1.6 существует функция, называемая monotonically_increasing_id()
Он генерирует новый столбец с уникальным 64-битным монотонным индексом для каждой строки
Но это не означает, что каждый раздел запускает новый диапазон, поэтому мы должны рассчитать каждое смещение раздела перед его использованием.
Пытаясь предоставить "rdd-free" решение, у меня получилось с помощью some 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")
}

Это решение не перепаковывает исходные строки и не переделывает оригинальный огромный фреймворк данных, поэтому он довольно быстро работает в реальном мире: 200 ГБ данных CSV (43 миллиона строк с 150 столбцами) читаются, индексируются и упаковываются в паркет через 2 минуты на 240 ядрах После тестирования моего решения я выполнил решение Kirk Broadhurst, и это было на 20 секунд медленнее
Вы можете или не хотите использовать dfWithPartitionId.cache(), зависит от задачи

Ответ 2

Следующее было опубликовано от имени Дэвида Гриффина (отредактировано).

Всеподобный, все-танцевальный метод 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)))
    )
  ) 
}

Ответ 3

Начиная с Spark 1.5, в Spark были добавлены выражения Window. Вместо преобразования DataFrame в RDD теперь вы можете использовать org.apache.spark.sql.expressions.row_number. Обратите внимание, что я нашел производительность для вышеперечисленного dfZipWithIndex значительно быстрее, чем приведенный ниже алгоритм. Но я отправляю его, потому что:

• У кого-то еще будет соблазн попробовать это.
• Возможно, кто-то может оптимизировать выражения ниже

Во всяком случае, вот что работает для меня:

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

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

Обратите внимание, что я использую lit(1) как для разбиения на разделы, так и для упорядочения - это делает все в одном разделе и, похоже, сохраняет исходный порядок DataFrame, но я полагаю, что это то, что замедляет его вниз.

Я протестировал его на 4-столбце DataFrame с 7 000 000 строк, и разница в скорости значительно велика между этим и выше dfZipWithIndex (как я уже сказал, функции RDD намного, намного быстрее).

Ответ 4

Версия 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() 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

Ответ 5

Версия Spark Java API:

Я реализовал решение @Evgeny для выполнения zipWithIndex на DataFrames в Java и хотел поделиться кодом.

Он также содержит улучшения, предложенные @fylb в его решении. Я могу подтвердить для Spark 2.4, что выполнение завершается неудачно, когда записи, возвращаемые spark_partition_id(), не начинаются с 0 или не увеличиваются последовательно. Поскольку эта функция задокументирована как недетерминированная, очень вероятно, что произойдет один из указанных выше случаев. Один пример вызван увеличением количества разделов.

Реализация 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");
    }

Ответ 6

@Evgeny, ваше решение интересно. Обратите внимание, что есть ошибка, когда у вас есть пустые разделы (в массиве отсутствуют эти индексы разделов, по крайней мере, это происходит со мной с помощью spark 1.6), поэтому я преобразовал массив в Map (partitionId → смещения).

В дополнение, я извлек источники 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 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/info/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")
  }
}