Come accedere all’elemento di una colonna VectorUDT in Spark DataFrame?

Ho un datafame df con una colonna VectorUDT denominata features . Come ottengo un elemento della colonna, ad esempio il primo elemento?

Ho provato a fare quanto segue

 from pyspark.sql.functions import udf first_elem_udf = udf(lambda row: row.values[0]) df.select(first_elem_udf(df.features)).show() 

ma ottengo una net.razorvine.pickle.PickleException: expected zero arguments for construction of ClassDict(for numpy.dtype) . Lo stesso errore se faccio first_elem_udf = first_elem_udf(lambda row: row.toArray()[0]) .

Ho anche provato explode() ma ottengo un errore perché richiede un array o un tipo di mappa.

Questa dovrebbe essere un’operazione comune, penso.

Converti l’output in float :

 from pyspark.sql.types import DoubleType from pyspark.sql.functions import lit, udf def ith_(v, i): try: return float(v[i]) except ValueError: return None ith = udf(ith_, DoubleType()) 

Esempio di utilizzo:

 from pyspark.ml.linalg import Vectors df = sc.parallelize([ (1, Vectors.dense([1, 2, 3])), (2, Vectors.sparse(3, [1], [9])) ]).toDF(["id", "features"]) df.select(ith("features", lit(1))).show() ## +-----------------+ ## |ith_(features, 1)| ## +-----------------+ ## | 2.0| ## | 9.0| ## +-----------------+ 

Spiegazione:

I valori di output devono essere riserializzati in oggetti Java equivalenti. Se vuoi accedere ai values ( SparseVectors attenzione a SparseVectors ) dovresti usare il metodo item :

 v.values.item(0) 

che restituiscono scalari standard di Python. Allo stesso modo se si desidera accedere a tutti i valori come una struttura densa:

 v.toArray().tolist()