PySpark - передать значение из другого столбца в качестве параметра искровой функции - PullRequest
Новая
       6

PySpark - передать значение из другого столбца в качестве параметра искровой функции

5 голосов
/ 20 июня 2020

У меня есть искровый фрейм данных, который выглядит так, где expr - это выражение фильтра SQL / Hive. 

+-----------------------------------------+
|expr                     |var1     |var2 |
+-------------------------+---------+-----+
|var1 > 7                 |9        |0    |
|var1 > 7                 |9        |0    |
|var1 > 7                 |9        |0    |
|var1 > 7                 |9        |0    |
|var1 = 3 AND var2 >= 0   |9        |0    |
|var1 = 3 AND var2 >= 0   |9        |0    |
|var1 = 3 AND var2 >= 0   |9        |0    |
|var1 = 3 AND var2 >= 0   |9        |0    |
|var1 = 2 AND var2 >= 0   |9        |0    |
+-------------------------+---------+-----+

Я хочу преобразовать этот фрейм данных в фрейм данных ниже, где flag - это логическое значение, найденное после оценки выражение в столбце 'expr' 

+---------------------------------------------------+
|expr                     |var1     |var2 |flag     |
+-------------------------+---------+-----+---------+
|var1 > 7                 |9        |0    |  True   |
|var1 > 7                 |9        |0    |  True   |
|var1 > 7                 |9        |0    |  True   |
|var1 > 7                 |9        |0    |  True   |
|var1 = 3 AND var2 >= 0   |9        |0    |     .   |
|var1 = 3 AND var2 >= 0   |9        |0    |     .   |
|var1 = 3 AND var2 >= 0   |9        |0    |     .   |
|var1 = 3 AND var2 >= 0   |9        |0    |     .   |
|var1 = 2 AND var2 >= 0   |9        |0    |     .   |
+-------------------------+---------+-----+---------+

Я пробовал использовать такую ​​функцию expr: 

df.withColumn('flag', expr(col('expr')))

Это не сработает, как и ожидалось, потому что функция expr ожидает строку в качестве параметра.

Еще одна идея, которую я подумывал использовать, - это создание UDF и передача в нее значения столбца expr, но это не позволит мне использовать функцию expr pyspark, потому что все UDF - это не-искровой код.

Каким должен быть мой подход? Есть предложения?

Ответы [ 2 ]

5 голосов
/ 20 июня 2020

Итак, вот решение PySpark без UDF. В Scala я считаю, что вы могли бы использовать map или foldleft с тем же logi c. 

exprs = [row.expr for row in df.select('expr').distinct().collect()]

for ex in exprs:
    df = df.withColumn('test', when(col('expr') == lit(ex), expr(ex)))
    
df.show()

+--------------------+----+----+----+
|                expr|var1|var2|test|
+--------------------+----+----+----+
|            var1 > 7|   9|   0|true|
|            var1 > 7|   9|   0|true|
|            var1 > 7|   9|   0|true|
|            var1 > 7|   9|   0|true|
|var1 = 3 AND var2...|   9|   0|null|
|var1 = 3 AND var2...|   9|   0|null|
|var1 = 3 AND var2...|   9|   0|null|
|var1 = 3 AND var2...|   9|   0|null|
|var1 = 2 AND var2...|   9|   0|null|
+--------------------+----+----+----+

Я должен указать, что я не понимаю, почему OP хочет это сделать, если они предоставят лучший контекст для проблемы, я уверен, что есть способ получше.

Итерация по DF не самая эффективная вещь, но в этом случае она будет работать очень быстро, поскольку не не перебирать данные, поэтому Spark фактически выполнит их в рамках одного плана. Кроме того, один метод collect () добавляет только 2 секунды ко времени выполнения для более чем 20 миллионов DF.

ОБНОВЛЕНИЕ:

Теперь я немного лучше понимаю проблему, это будет быстрее, так как Spark вычислит все фильтры сразу, прежде чем объединить их в один столбец. 

# Tip: perform the collect statement on the smaller DF that contains the filter expressions

exprs = [row.expr for row in df.select('expr').distinct().collect()]

df = df.withColumn('filter',
              coalesce(*[when(col('expr') == lit(ex), expr(ex)) for ex in exprs])
             )
df.show()

+--------------------+----+----+------+
|                expr|var1|var2|filter|
+--------------------+----+----+------+
|            var1 > 7|   9|   0|true  |
|            var1 > 7|   9|   0|true  |
|            var1 > 7|   9|   0|true  |
|            var1 > 7|   9|   0|true  |
|var1 = 3 AND var2...|   9|   0|null  |
|var1 = 3 AND var2...|   9|   0|null  |
|var1 = 3 AND var2...|   9|   0|null  |
|var1 = 3 AND var2...|   9|   0|null  |
|var1 = 2 AND var2...|   9|   0|null  |
+--------------------+----+----+------+
0 голосов
/ 20 июня 2020

не udf 

  val exprs5 =   sourceDF.select('expr).distinct().as[String].collect()
  val d1 = exprs5.map(i => {
    val df = sourceDF.filter('expr.equalTo(i))
    df.withColumn("flag", expr(i))
  })
  val d2 = d1.reduce(_ union _)

udf 

package spark

import org.apache.spark.sql.{DataFrame, SparkSession}

object Filter extends App {

  val spark = SparkSession.builder()
    .master("local")
    .appName("DataFrame-example")
    .getOrCreate()

  import spark.implicits._

  val sourceDF = Seq(("var1 > 7", 9, 0),
  ("var1 > 7", 9, 0),
  ("var1 > 7", 9, 0),
  ("var1 > 7", 9, 0),
  ("var1 = 3 AND var2 >= 0", 9, 0),
  ("var1 = 3 AND var2 >= 0", 9, 0),
  ("var1 = 3 AND var2 >= 0", 9, 0),
  ("var1 = 3 AND var2 >= 0", 9, 0),
  ("var1 = 2 AND var2 >= 0", 9, 0)).toDF("expr", "var1","var2")

  import org.apache.spark.sql.functions._

  val fCheck = udf((expr: String, val1: Int, val2: Int) => {
    expr.split(" ") match{
      case Array(vr, z, vl) if (vr == "var1" && z == ">") => Some(val1 > vl.toInt)
      case Array(vr1, z1,  vl1, logic1, vr2, z2, vl2)
        if (vr1 == "var1") && (z1 == "=") && (logic1 == "AND") && (vr2 == "var2") && (z2 == ">=")
      => Some((val1 == vl1.toInt ) && (val2 >= vl2.toInt))
      case _ => None
    }
  })

  val resultDF = sourceDF.withColumn("flag", lit(fCheck('expr, 'var1, 'var2)))

  resultDF.show(false)
//  +----------------------+----+----+-----+
//  |expr                  |var1|var2|flag |
//  +----------------------+----+----+-----+
//  |var1 > 7              |9   |0   |true |
//  |var1 > 7              |9   |0   |true |
//  |var1 > 7              |9   |0   |true |
//  |var1 > 7              |9   |0   |true |
//  |var1 = 3 AND var2 >= 0|9   |0   |false|
//  |var1 = 3 AND var2 >= 0|9   |0   |false|
//  |var1 = 3 AND var2 >= 0|9   |0   |false|
//  |var1 = 3 AND var2 >= 0|9   |0   |false|
//  |var1 = 2 AND var2 >= 0|9   |0   |false|
//  +----------------------+----+----+-----+

}
...