Filtrar por si el valor de la columna es igual a una lista en Spark

Estoy tratando de filtrar un dataframe de Spark en función de si los valores en una columna son iguales a una lista. Me gustaría hacer algo como esto:

filtered_df = df.where(df.a == ['list','of' , 'stuff']) 

Donde filtered_df solo contiene filas donde el valor de filtered_df.a es ['list','of' , 'stuff'] y el tipo de a es array (nullable = true) .

Actualización :

Con las versiones actuales puedes usar una array de literales:

 from pyspark.sql.functions import array, lit df.where(df.a == array(*[lit(x) for x in ['list','of' , 'stuff']])) 

Respuesta original :

Bueno, una forma un poco pirateada de hacerlo, que no requiere un trabajo por lotes de Python, es algo como esto:

 from pyspark.sql.functions import col, lit, size from functools import reduce from operator import and_ def array_equal(c, an_array): same_size = size(c) == len(an_array) # Check if the same size # Check if all items equal same_items = reduce( and_, (c.getItem(i) == an_array[i] for i in range(len(an_array))) ) return and_(same_size, same_items) 

Examen rápido:

 df = sc.parallelize([ (1, ['list','of' , 'stuff']), (2, ['foo', 'bar']), (3, ['foobar']), (4, ['list','of' , 'stuff', 'and', 'foo']), (5, ['a', 'list','of' , 'stuff']), ]).toDF(['id', 'a']) df.where(array_equal(col('a'), ['list','of' , 'stuff'])).show() ## +---+-----------------+ ## | id| a| ## +---+-----------------+ ## | 1|[list, of, stuff]| ## +---+-----------------+ 

Puedes crear un udf. Por ejemplo:

 def test_in(x): return x == ['list','of' , 'stuff'] from pyspark.sql.functions import udf f = udf(test_in, pyspark.sql.types.BooleanType()) filtered_df = df.where(f(df.a))