Question

我有一个PySpark数据框（例如df1），其中包含以下列

1.> category：一些字符串 2.> array1：元素数组 3.> array2：元素数组

以下是df1

的示例

+--------+--------------+--------------+
|category|        array1|        array2|
+--------+--------------+--------------+
|A       |  [x1, x2, x3]|  [y1, y2, y3]|
|B       |      [u1, u2]|      [v1, v2]|
+--------+--------------+--------------+

对于每一行，array1的长度等于array2的长度。在每一列中，我希望array1（和array2）的不同行具有不同大小的数组。

我想形成单独的列（例如element1和element2），以便在每一行中，列element1和element2包含来自相同位置的元素分别为array1和array2

。

以下是我想要的输出数据帧（例如df2）的示例：

+--------+--------------+--------------+----------+----------+
|category|        array1|        array2|  element1|  element2|
+--------+--------------+--------------+----------+----------+
|A       |  [x1, x2, x3]|  [y1, y2, y3]|        x1|        y1|
|A       |  [x1, x2, x3]|  [y1, y2, y3]|        x2|        y2|
|A       |  [x1, x2, x3]|  [y1, y2, y3]|        x3|        y3|
|B       |      [u1, u2]|      [v1, v2]|        u1|        v1|
|B       |      [u1, u2]|      [v1, v2]|        u2|        v2|
+--------+--------------+--------------+----------+----------+

以下是我到目前为止尝试过的内容（但除了我想要的内容之外，它还提供了来自其他位置的element1和element2中的值。）

df2 = df1.select( "*", F.explode("array1").alias("element1") ).select( "*", F.explode("array2").alias("element2") )

Answer 1

对于Spark> = 2.4，您可以使用Higher-Order Functions：

data = [('A', ['x1', 'x2', 'x3'], ['y1', 'y2', 'y3']),
        ('B', ['u1', 'u2'], ['v1', 'v2'])
      ]
df = spark.createDataFrame(data, ["category", "array1", "array2"])

# tranform array1, array2 => [struct(element1, element2)]
transform_expr = "transform(array1, (x, i) -> struct(x as element1, array2[i] as element2))"

# explode transformed arrays and extract values of element1 and element2
df.withColumn("merged_arrays", explode(expr(transform_expr))) \
    .withColumn("element1", col("merged_arrays.element1")) \
    .withColumn("element2", col("merged_arrays.element2")) \
    .drop("merged_arrays") \
    .show(truncate=False)

输出：

+--------+------------+------------+--------+--------+
|category|array1      |array2      |element1|element2|
+--------+------------+------------+--------+--------+
|A       |[x1, x2, x3]|[y1, y2, y3]|x1      |y1      |
|A       |[x1, x2, x3]|[y1, y2, y3]|x2      |y2      |
|A       |[x1, x2, x3]|[y1, y2, y3]|x3      |y3      |
|B       |[u1, u2]    |[v1, v2]    |u1      |v1      |
|B       |[u1, u2]    |[v1, v2]    |u2      |v2      |
+--------+------------+------------+--------+--------+

transform功能的说明：

该函数采用第一个数组array1并应用lambda函数(x, i) -> struct(string, string)，其中x是实际值，而i是其在数组中的索引。对于每个值，我们返回一个结构，该结构包含该值作为element1和array2中的对应值（使用索引i）作为element2。

其余的只是分解转换结果并访问我们创建的struct元素。

Answer 2

初始化

import pyspark.sql.functions as F

sc = pyspark.SparkContext()
sqlContext = SQLContext(sc)

columns = ['category','array1','array2']
vals = [
     ('A', ['x1', 'x2', 'x3'], ['y1','y2','y3']),
     ('B', ['u1', 'u2',], ['v1','v2'])
]
df = sqlContext.createDataFrame(vals, columns)

Based on arrays_zip [docs]火花> = 2.4

df.withColumn('new', F.arrays_zip('array1','array2')).withColumn('ex',explode('new'))\
    .select('category','array1','array2',
            col("ex.array1").alias('element1'),
            col("ex.array2").alias('element2')\
           ).drop('new','ex').show()

输出

+--------+------------+------------+--------+--------+
|category|      array1|      array2|element1|element2|
+--------+------------+------------+--------+--------+
|       A|[x1, x2, x3]|[y1, y2, y3]|      x1|      y1|
|       A|[x1, x2, x3]|[y1, y2, y3]|      x2|      y2|
|       A|[x1, x2, x3]|[y1, y2, y3]|      x3|      y3|
|       B|    [u1, u2]|    [v1, v2]|      u1|      v1|
|       B|    [u1, u2]|    [v1, v2]|      u2|      v2|
+--------+------------+------------+--------+--------+

说明看arrays_zip产生的内容基本上可以解释一切。我们将cols与其一起压缩/压缩，然后explode。然后只需在explode创建的新列中引用相应的结构即可。

>>> df.withColumn('new', F.arrays_zip('array1','array2')).show(truncate=False)
+--------+------------+------------+------------------------------+
|category|array1      |array2      |new                           |
+--------+------------+------------+------------------------------+
|A       |[x1, x2, x3]|[y1, y2, y3]|[[x1, y1], [x2, y2], [x3, y3]]|
|B       |[u1, u2]    |[v1, v2]    |[[u1, v1], [u2, v2]]          |
+--------+------------+------------+------------------------------+

分解两个PySpark阵列并使元素保持在相同位置

2 个答案: