как разделить фрейм данных на логический столбец

#apache-spark #spark-dataframe

#apache-spark #apache-spark-sql

Вопрос:

Я хочу разделить a DataFrame относительно логического столбца.

Я придумал :

 def partition(df: DataFrame, c: Column): (DataFrame, DataFrame) = 
  (df.filter(c === true), df.filter(c === false))
  

Примечание: в моем случае c это UDF.

Есть ли лучший способ?

Я бы хотел :

  • чтобы избежать сканирования фрейма данных 2 раза
  • чтобы избежать уродливых логических тестов

Вот пример :

 @ val df = sc.parallelize(Seq(1,2,3,4)).toDF("i")
df: org.apache.spark.sql.DataFrame = [i: int]
@ val u = udf((i: Int) => i % 2 == 0)
u: org.apache.spark.sql.UserDefinedFunction = UserDefinedFunction(<function1>, BooleanType, List(IntegerType))
@ partition(df, u($"i"))
res25: (org.apache.spark.sql.DataFrame, org.apache.spark.sql.DataFrame) = ([i: int], [i: int])
  

Комментарии:

1. можете ли вы поделиться воспроизводимым примером?

Ответ №1:

используйте combineByKey для логического столбца

 data.combineByKey(lambda value: (value, 1),
                             lambda x, value: (x[0]   value, x[1]   1),
                             lambda x, y: (x[0]   y[0], x[1]   y[1]))