получить первые N элементов из столбца данных ArrayType в pyspark

Question

У меня есть искровой фрейм данных со строками в виде -

1   |   [a, b, c]
2   |   [d, e, f]
3   |   [g, h, i]

Теперь я хочу сохранить только первые 2 элемента из столбца массива.

1   |   [a, b]
2   |   [d, e]
3   |   [g, h]

Как этого достичь?

Примечание. Помните, что здесь я извлекаю не один элемент массива, а часть массива, которая может содержать несколько элементов.

Answer 1

Вот как это сделать с помощью функций API.

Предположим, что ваш DataFrame был следующим:

df.show()
#+---+---------+
#| id|  letters|
#+---+---------+
#|  1|[a, b, c]|
#|  2|[d, e, f]|
#|  3|[g, h, i]|
#+---+---------+

df.printSchema()
#root
# |-- id: long (nullable = true)
# |-- letters: array (nullable = true)
# |    |-- element: string (containsNull = true)

Вы можете использовать квадратные скобки для доступа к элементам в столбце letters с помощьюиндекс и оберните его при вызове pyspark.sql.functions.array() для создания нового столбца ArrayType.

import pyspark.sql.functions as f

df.withColumn("first_two", f.array([f.col("letters")[0], f.col("letters")[1]])).show()
#+---+---------+---------+
#| id|  letters|first_two|
#+---+---------+---------+
#|  1|[a, b, c]|   [a, b]|
#|  2|[d, e, f]|   [d, e]|
#|  3|[g, h, i]|   [g, h]|
#+---+---------+---------+

Или, если у вас слишком много индексов для списка, вы можете использовать понимание списка:

df.withColumn("first_two", f.array([f.col("letters")[i] for i in range(2)])).show()
#+---+---------+---------+
#| id|  letters|first_two|
#+---+---------+---------+
#|  1|[a, b, c]|   [a, b]|
#|  2|[d, e, f]|   [d, e]|
#|  3|[g, h, i]|   [g, h]|
#+---+---------+---------+

Answer 2

Либо мои навыки работы с pyspark стали ржавыми (признаюсь, что в настоящее время я их больше не оттачиваю), либо это действительно крепкий орешек ... Единственный способ, которым мне удалось это сделать, - использовать операторы SQL:

spark.version
#  u'2.3.1'

# dummy data:

from pyspark.sql import Row
x = [Row(col1="xx", col2="yy", col3="zz", col4=[123,234, 456])]
rdd = sc.parallelize(x)
df = spark.createDataFrame(rdd)
df.show()
# result:
+----+----+----+---------------+
|col1|col2|col3|           col4|
+----+----+----+---------------+
|  xx|  yy|  zz|[123, 234, 456]|
+----+----+----+---------------+

df.createOrReplaceTempView("df")
df2 = spark.sql("SELECT col1, col2, col3, (col4[0], col4[1]) as col5 FROM df")
df2.show()
# result:
+----+----+----+----------+ 
|col1|col2|col3|      col5|
+----+----+----+----------+ 
|  xx|  yy|  zz|[123, 234]|
+----+----+----+----------+

Для будущих вопросов было бы хорошо следовать предложенным рекомендациям по Как создать хорошие воспроизводимые примеры Apache Spark Dataframe .