Spark – Creando DataFrame nested

Estoy empezando con PySpark y tengo problemas para crear DataFrames con objetos nesteds.

Este es mi ejemplo.

Tengo usuarios.

$ cat user.json {"id":1,"name":"UserA"} {"id":2,"name":"UserB"} 

Los usuarios tienen pedidos.

 $ cat order.json {"id":1,"price":202.30,"userid":1} {"id":2,"price":343.99,"userid":1} {"id":3,"price":399.99,"userid":2} 

Y me gusta unirme a él para obtener una estructura en la que los pedidos se anidan en los usuarios.

 $ cat join.json {"id":1, "name":"UserA", "orders":[{"id":1,"price":202.30,"userid":1},{"id":2,"price":343.99,"userid":1}]} {"id":2,"name":"UserB","orders":[{"id":3,"price":399.99,"userid":2}]} 

Cómo puedo hacer eso ? ¿Hay algún tipo de unión anidada o algo similar?

 >>> user = sqlContext.read.json("user.json") >>> user.printSchema(); root |-- id: long (nullable = true) |-- name: string (nullable = true) >>> order = sqlContext.read.json("order.json") >>> order.printSchema(); root |-- id: long (nullable = true) |-- price: double (nullable = true) |-- userid: long (nullable = true) >>> joined = sqlContext.read.json("join.json") >>> joined.printSchema(); root |-- id: long (nullable = true) |-- name: string (nullable = true) |-- orders: array (nullable = true) | |-- element: struct (containsNull = true) | | |-- id: long (nullable = true) | | |-- price: double (nullable = true) | | |-- userid: long (nullable = true) 

EDIT: Sé que existe la posibilidad de hacer esto usando join y foldByKey, pero ¿hay alguna forma más simple?

EDIT2: Estoy usando la solución por @ zero323

 def joinTable(tableLeft, tableRight, columnLeft, columnRight, columnNested, joinType = "left_outer"): tmpTable = sqlCtx.createDataFrame(tableRight.rdd.groupBy(lambda r: r.asDict()[columnRight])) tmpTable = tmpTable.select(tmpTable._1.alias("joinColumn"), tmpTable._2.data.alias(columnNested)) return tableLeft.join(tmpTable, tableLeft[columnLeft] == tmpTable["joinColumn"], joinType).drop("joinColumn") 

Añado la segunda estructura anidada ‘líneas’

 >>> lines = sqlContext.read.json(path + "lines.json") >>> lines.printSchema(); root |-- id: long (nullable = true) |-- orderid: long (nullable = true) |-- product: string (nullable = true) orders = joinTable(order, lines, "id", "orderid", "lines") joined = joinTable(user, orders, "id", "userid", "orders") joined.printSchema() root |-- id: long (nullable = true) |-- name: string (nullable = true) |-- orders: array (nullable = true) | |-- element: struct (containsNull = true) | | |-- id: long (nullable = true) | | |-- price: double (nullable = true) | | |-- userid: long (nullable = true) | | |-- lines: array (nullable = true) | | | |-- element: struct (containsNull = true) | | | | |-- _1: long (nullable = true) | | | | |-- _2: long (nullable = true) | | | | |-- _3: string (nullable = true) 

Después de esta columna se pierden los nombres de las líneas. Algunas ideas ?

EDIT 3: traté de especificar manualmente el esquema.

 from pyspark.sql.types import * fields = [] fields.append(StructField("_1", LongType(), True)) inner = ArrayType(lines.schema) fields.append(StructField("_2", inner)) new_schema = StructType(fields) print new_schema grouped = lines.rdd.groupBy(lambda r: r.orderid) grouped = grouped.map(lambda x: (x[0], list(x[1]))) g = sqlCtx.createDataFrame(grouped, new_schema) 

Error:

 TypeError: StructType(List(StructField(id,LongType,true),StructField(orderid,LongType,true),StructField(product,StringType,true))) can not accept object in type  

Esto funcionará solo en Spark 2.0 o posterior

Primero necesitaremos un par de importaciones:

 from pyspark.sql.functions import struct, collect_list 

El rest es una simple agregación y unión.

 orders = spark.read.json("/path/to/order.json") users = spark.read.json("/path/to/user.json") combined = users.join( orders .groupBy("userId") .agg(collect_list(struct(*orders.columns)).alias("orders")) .withColumnRenamed("userId", "id"), ["id"]) 

Para los datos de ejemplo el resultado es:

 combined.show(2, False) 
 +---+-----+---------------------------+ |id |name |orders | +---+-----+---------------------------+ |1 |UserA|[[1,202.3,1], [2,343.99,1]]| |2 |UserB|[[3,399.99,2]] | +---+-----+---------------------------+ 

con esquema:

 combined.printSchema() 
 root |-- id: long (nullable = true) |-- name: string (nullable = true) |-- orders: array (nullable = true) | |-- element: struct (containsNull = true) | | |-- id: long (nullable = true) | | |-- price: double (nullable = true) | | |-- userid: long (nullable = true) 

y representación JSON:

 for x in combined.toJSON().collect(): print(x) 
 {"id":1,"name":"UserA","orders":[{"id":1,"price":202.3,"userid":1},{"id":2,"price":343.99,"userid":1}]} {"id":2,"name":"UserB","orders":[{"id":3,"price":399.99,"userid":2}]} 

Primero, debe usar el userid como la clave de unión para el segundo DataFrame :

 user.join(order, user.id == order.userid) 

Luego, puede usar un paso del map para transformar los registros resultantes al formato deseado.