Spark / Scala: использовать строковую переменную в условных выражениях в операциях DataFrame

Question

Позвольте мне объяснить это на примере. Начиная со следующего кадра данных

val df = Seq((1, "CS", 0, Array(0.1, 0.2, 0.4, 0.5)),
             (4, "Ed", 0, Array(0.4, 0.8, 0.3, 0.6)),
             (7, "CS", 0, Array(0.2, 0.5, 0.4, 0.7)),
             (101, "CS", 1, Array(0.5, 0.7, 0.3, 0.8)),
             (5, "CS", 1, Array(0.4, 0.2, 0.6, 0.9))).toDF("id", "dept", "test", "array")
df.show()
+---+----+----+--------------------+
| id|dept|test|               array|
+---+----+----+--------------------+
|  1|  CS|   0|[0.1, 0.2, 0.4, 0.5]|
|  4|  Ed|   0|[0.4, 0.8, 0.3, 0.6]|
|  7|  CS|   0|[0.2, 0.5, 0.4, 0.7]|
|101|  CS|   1|[0.5, 0.7, 0.3, 0.8]|
|  5|  CS|   1|[0.4, 0.2, 0.6, 0.9]|
+---+----+----+--------------------+

Рассматривая в качестве примера следующие две общие операции (но не ограничиваясь ими):

import org.apache.spark.sql.functions._ // for `when`
val dfFilter1 = df.where($"dept" === "CS")
val dfFilter3 = df.withColumn("category", when($"dept" === "CS" && $"id" === 101, 10).otherwise(0))

Теперь у меня есть строковая переменная colName = "dept". И требуется, чтобы $ "dept" в предыдущей операции был заменен на colName в некоторой форме для достижения той же функциональности. Мне удалось добиться первого, как показано ниже:

val dfFilter2 = df.where(s"${colName} = 'CS'")

Но подобная операция завершается неудачно во втором случае:

val dfFilter4 = df.withColumn("category", when(s"${colName} = 'CS'" && $"id" === 101, 10).otherwise(0))

В частности, это дает следующую ошибку:

Name: Unknown Error
Message: <console>:35: error: value && is not a member of String
       val dfFilter4 = df.withColumn("category", when(s"${colName} = 'CS'" && $"id" === 101, 10).otherwise(0))

Насколько я понимаю, после того, как я использую s "$ {variable}" для работы с переменной, все становится чистой строкой, и трудно задействовать логическую операцию.

Итак, мой вопрос: 1. Как лучше всего использовать такую ​​строковую переменную, как colName, для операций, подобных тем, что я перечислил выше (мне также не нравится решение, которое у меня есть для .where ())? 2. Существуют ли общие рекомендации по использованию такой строковой переменной в более общих операциях, кроме двух приведенных здесь примеров (я всегда чувствовал, что она очень специфична для конкретного случая, когда я имею дело со связанными со строками операциями).

Answer 1

1. Как лучше всего использовать такую ​​строковую переменную, как colName, для операций, аналогичных двум, перечисленным выше

Вы можете использовать col функцию от org.apache.spark.sql.functions

import org.apache.spark.sql.functions._
val colName = "dept"

Для dfFilter2

val dfFilter2 = df.where(col(colName) === "CS")

Для dfFilter4

val dfFilter4 = df.withColumn("category", when(col(colName) === "CS" && $"id" === 101, 10).otherwise(0))

Answer 2

Вы можете использовать expr функцию как

val dfFilter4 = df.withColumn("category", when(expr(s"${colName} = 'CS' and id = 101"), 10).otherwise(0))

Причина ошибки

where функция, если она определена с помощью строкового запроса, как показано ниже

val dfFilter2 = df.where(s"${colName} = 'CS'")

потому что есть поддержка apis для строки и столбца

/** * Filters rows using the given condition. This is an alias for фильтр . * {{{ * // The following are equivalent: * peopleDs.filter($"age" > 15) * peopleDs.where($"age" > 15) * }}} * * @group typedrel * @since 1.6.0 */ def where(condition: Column): Dataset[T] = filter(condition)

и

/** * Filters rows using the given SQL expression. * {{{ * peopleDs.where("age > 15") * }}} * * @group typedrel * @since 1.6.0 */ def where(conditionExpr: String): Dataset[T] = { filter(Column(sparkSession.sessionState.sqlParser.parseExpression(conditionExpr))) }

Но есть только один API для функции when, поддерживающей только тип столбца

/** * Evaluates a list of conditions and returns one of multiple possible result expressions. * If otherwise is not defined at the end, null is returned for unmatched conditions. * * {{{ * // Example: encoding gender string column into integer. * * // Scala: * people.select(when(people("gender") === "male", 0) * .when(people("gender") === "female", 1) * .otherwise(2)) * * // Java: * people.select(when(col("gender").equalTo("male"), 0) * .when(col("gender").equalTo("female"), 1) * .otherwise(2)) * }}} * * @group normal_funcs * @since 1.4.0 */ def when(condition: Column, value: Any): Column = withExpr { CaseWhen(Seq((condition.expr, lit(value).expr))) }

То есть вы не можете использовать строку sql запроса для when функции

Итак, правильный способ сделать это следующим образом

val dfFilter4 = df.withColumn("category", when(col(s"${colName}") === "CS" && $"id" === 101, 10).otherwise(0))

или коротко как

val dfFilter4 = df.withColumn("category", when(col(colName) === "CS" && col("id") === 101, 10).otherwise(0))