PySpark не может соответствовать массиву

Question

Я использую PySpark для простой фильтрации данных.Кадр данных Spark df_rules выглядит следующим образом:

Я получил это df_rules следующим образом:

from pyspark.ml.fpm import FPGrowth
from pyspark.sql import SparkSession

spark = SparkSession.builder.master("local")\
   .appName("Association Rules FP-Growth")\
   .config("spark.some.config.option", "some-value")\
   .getOrCreate()

df = spark.createDataFrame([
    (0, [1, 2, 5]),
    (1, [1, 2, 3, 5]),
    (2, [1, 2])
], ["id", "items"])

fpGrowth = FPGrowth(itemsCol="items", minSupport=0.5, minConfidence=0.6)
model = fpGrowth.fit(df)

# Display frequent itemsets.
model.freqItemsets.show()

# Display generated association rules.
df_rules = model.associationRules

Я просто хочу сделать df_rules.where(df_rules.consequent == [1]).Сначала он дал мне ошибку несоответствия типов данных, поскольку df_rules.consequent равно array<bigint>.Поэтому я преобразовал consequent тип данных столбца через:

from pyspark.sql.types import ArrayType, IntegerType
df_rules = df_rules.withColumn("consequent", df_rules.consequent.cast(ArrayType(IntegerType())))

Но все равно получил ошибку:

Знаете ли вы, какя могу успешно выполнить фильтрацию?

Answer 1

Вам не нужно конвертировать array<bigint> в array<int>, просто используйте long:

from pyspark.sql.functions import array, lit

df_rules.where(df_rules.consequent == array(lit(1L)))