PySpark добавить Id столбец и фильтр не работает

Question

У меня есть набор данных с 233 465 строками, который растет примерно на 10000 строк в день. Мне нужно случайным образом выбрать строки из полного набора данных для использования в обучении ML. Я добавил столбец "id" для "index".

from pyspark.sql.functions import monotonically_increasing_id
spark_df = n_data.withColumn("id", monotonically_increasing_id())

Я выполняю следующий код, ожидая увидеть 5 возвращенных строк, где идентификаторы совпадают со списком "индексы" с количеством 5.

indices = [1000, 999, 45, 1001, 1823, 123476]
result = spark_df.filter(col("id").isin(indices))
result.show()
print(result.count())

вместо этого я получаю 3 ряда. Я получаю идентификаторы для 45, 1000 и 1001.

Есть идеи о том, что здесь может быть не так? Это кажется довольно резким и сухим.

Спасибо!

Answer 1

Не существует прямой функции для вызова уникальных последовательных идентификаторов каждой строке. Но есть обходной путь, использующий windowфункция на основе.

df = spark.createDataFrame([(3,),(7,),(9,),(1,),(-3,),(5,)], ["values"])
df.show()

+------+
|values|
+------+
|     3|
|     7|
|     9|
|     1|
|    -3|
|     5|
+------+



df = (df.withColumn('dummy', F.monotonically_increasing_id())
       .withColumn('ID', F.row_number().over(Window.orderBy('dummy')))
       .drop('dummy'))
df.show()

+------+---+
|values| ID|
+------+---+
|     3|  1|
|     7|  2|
|     9|  3|
|     1|  4|
|    -3|  5|
|     5|  6|
+------+---+