Как писать IN и NOT IN с использованием Apache Spark Dataframes

Question

У меня ниже 2 примеров SQL-запросов в SQL:

a) update DBTABLE1
    set col1 = 'Yes'
where ID IN ( '100' ) and City = any(select City from DBTable2 where Country = 'USA');

b) update DBTABLE1
    set col2 = 'No'
where ID NOT IN ( '100' ) and City = any(select City from DBTable2 where Country = 'USA');


How to write above 2 SQLs using Apache Spark Dataframes (Not Select subquery etc). A dataframe is already having these 2 columns - col1 and col2, I am changing their values using WITHCOLUMN and WHEN clause. 

CitiDF содержит набор данных о количестве городов.

I tried below but giving compile errors:

c) This is for (a) above:

withcolumn(col("col1"),when(col("id") === lit("100") 
and col("city").isin(CitiDF("city")), lit("yes")))

d) This is for (b) above:

withcolumn(col("col2"),when(col("id") === lit("100") 
and ! (col("city").isin(CitiDF("city"))), lit("yes")))

Answer 1

С примерами данных, использованными Джейсоном, и без смешивания списков снаружи, решение может быть достигнуто с помощью самого spark-sql. Проверьте это:

val df = Seq((100,"Frankfurt","filler","filler"),(200,"Berlin","filler","filler"),(100,"Vienna","filler","filler"),(500,"Victoria","filler","filler"),(600,"Shanghai","filler","filler"),(100,"Cologne","filler","filler")).toDF("id","city","col1","col2")
df.createOrReplaceTempView("city_details")
val city = Seq(("Frankfurt"),("Vienna")).toDF("cityName")
city.createOrReplaceTempView("city_list")
df.show(false)
spark.sql(
  """ select id,city, case when id=100 and array_contains((select collect_list(cityname) from city_list), city) then 'yes' else null end as col1,
      case when id=100 and not array_contains((select collect_list(cityname) from city_list), city) then 'yes' else null end as col2
      from city_details
  """).show(false)

Выход:

+---+---------+----+----+
|id |city     |col1|col2|
+---+---------+----+----+
|100|Frankfurt|yes |null|
|200|Berlin   |null|null|
|100|Vienna   |yes |null|
|500|Victoria |null|null|
|600|Shanghai |null|null|
|100|Cologne  |null|yes |
+---+---------+----+----+

Answer 2

Чтобы сделать вещи более конкретными, давайте рассмотрим некоторые игрушечные данные.У нас есть один DataFrame с именем df, который выглядит следующим образом:

+---+---------+------+------+
| id|     city|  col1|  col2|
+---+---------+------+------+
|100|Frankfurt|filler|filler|
|200|   Berlin|filler|filler|
|100|   Vienna|filler|filler|
|500| Victoria|filler|filler|
|600| Shanghai|filler|filler|
|100|  Cologne|filler|filler| 
+---+---------+------+------+

и другой, с именем cities, который выглядит следующим образом:

+---------+
| cityName|
+---------+
|Frankfurt|
|   Vienna|
+---------+

Мы можем выполнить ваши запросынапример:

val cityList = cities.collect.map(x => x(0))
val df1 = df.withColumn("col1", when($"id" === "100" and $"city".isin(cityList: _*), "yes"))

Результат, который мы получаем:

+---+---------+----+------+
| id|     city|col1|  col2|
+---+---------+----+------+
|100|Frankfurt| yes|filler|
|200|   Berlin|null|filler|
|100|   Vienna| yes|filler|
|500| Victoria|null|filler|
|600| Shanghai|null|filler|
|100|  Cologne|null|filler|
+---+---------+----+------+

Для вашего второго запроса мы используем такие же cityList:

val df2 = df.withColumn("col2", when($"id" === "100" and !$"city".isin(cityList: _*), "yes"))

us

+---+---------+------+----+
| id|     city|  col1|col2|
+---+---------+------+----+
|100|Frankfurt|filler|null|
|200|   Berlin|filler|null|
|100|   Vienna|filler|null|
|500| Victoria|filler|null|
|600| Shanghai|filler|null|
|100|  Cologne|filler| yes|
+---+---------+------+----+

Тем не менее, есть большой недостаток в этом подходе.Если количество городов велико, вам может не хватить памяти, собрав все названия.Вместо этого я хотел бы рассмотреть альтернативный подход, такой как внешнее соединение:

df.join(cities, df("city") === cities("cityName"), "outer").
  withColumn("col1", when($"cityName".isNotNull and $"id" === "100", "yes")).
  withColumn("col2", when($"cityName".isNull and $"id" === "100", "yes")).
  drop("cityName")

, дающий нам

+---+---------+----+----+
| id|     city|col1|col2|
+---+---------+----+----+
|100|Frankfurt| yes|null|
|500| Victoria|null|null|
|200|   Berlin|null|null|
|100|   Vienna| yes|null|
|100|  Cologne|null| yes|
|600| Shanghai|null|null|
+---+---------+----+----+

Да, он вводит дополнительный столбец, но только временно, и избегаетпотянув потенциально большой список городов в память водителя.