在pyspark对象框架中的数组上迭代，并基于与数组中的值同名的列创建新列

bybem2ql 于 2024-01-06 发布在 Spark

关注(0)|答案(1)|浏览(209)

我有一个这样格式的表格：
| 名称|水果|苹果|香蕉|橙子|
| --|--|--|--|--|
| 爱丽丝|[“苹果”、“香蕉”、“橙子”]| 5 | 8 | 3 |
| 鲍勃|[“苹果”]| 2 | 9 | 1 |
我想创建一个新列，其中包含以下格式的JSON包，其中键是数组的元素，值是列名称的结果值：
| 名称|水果|苹果|香蕉|橙子|新科尔|
| --|--|--|--|--|--|
| 爱丽丝|[“苹果”、“香蕉”、“橙子”]| 5 | 8 | 3 |{“苹果”：5，“香蕉”：8，“橙子”：3}|
| 鲍勃|[“苹果”]| 2 | 9 | 1 |{“apple”：2}|
我假设有一个UDF，但是我不能得到正确的语法。
这是我对代码的理解：

from pyspark.sql.functions import udf, col
from pyspark.sql.types import MapType, StringType
# Create a Spark session
spark = SparkSession.builder.appName("example").getOrCreate()
# Sample data
data = [("Alice", ["apple", "banana", "orange"], 5, 8, 3),
        ("Bob", ["apple"], 2, 9, 1)]
# Define the schema
schema = ["name", "fruits", "apple", "banana", "orange"]
# Create a DataFrame
df = spark.createDataFrame(data, schema=schema)
# Show the initial DataFrame
print("Initial DataFrame:")
display(df)
# Define a UDF to create a dictionary
@udf(MapType(StringType(), StringType()))
def json_map(fruits):
    result = {}
    for i in fruits:
        result[i] = col(i)
    return result
# Apply the UDF to the 'fruits' column
new_df = df.withColumn('test', json_map(col('fruits')))
# Display the updated DataFrame
display(new_df)

字符串

pyspark

来源：https://stackoverflow.com/questions/77628601/iterate-over-an-array-in-a-pyspark-dataframe-and-create-a-new-column-based-on-c

1条答案

按热度按时间

qyuhtwio1#

你可以使用Abdennacer在他的回答中分享的arrays_zip方法，但前提是数组元素应该与你的列对齐，这可能并不总是如此。
另一种方法是为列创建Map数组，并过滤该数组以仅保留fruits数组中可用的键的键值对。
这里有一个例子

# i've changed the input slightly to rearrange the fruits array
# +-----+-----------------------+-----+------+------+
# |name |fruits                 |apple|banana|orange|
# +-----+-----------------------+-----+------+------+
# |Alice|[orange, banana, apple]|5    |8     |3     |
# |Bob  |[apple]                |2    |9     |1     |
# +-----+-----------------------+-----+------+------+
data_sdf. \
    withColumn('fruitcols_arr', 
               func.array(*[func.create_map([func.lit(c), func.col(c)]) for c in data_sdf.drop('name', 'fruits').columns])
               ). \
    withColumn('fruitcols_arr', 
               func.expr('filter(fruitcols_arr, x -> array_contains(fruits, map_keys(x)[0]))')
               ). \
    withColumn('new_col',
               func.aggregate(func.expr('slice(fruitcols_arr, 2, size(fruitcols_arr))'),
                              func.col('fruitcols_arr')[0],
                              lambda x, y: func.map_concat(x, y)
                              )
               ). \
    drop('fruitcols_arr'). \
    show(truncate=False)
# +-----+-----------------------+-----+------+------+--------------------------------------+
# |name |fruits                 |apple|banana|orange|new_col                               |
# +-----+-----------------------+-----+------+------+--------------------------------------+
# |Alice|[orange, banana, apple]|5    |8     |3     |{apple -> 5, banana -> 8, orange -> 3}|
# |Bob  |[apple]                |2    |9     |1     |{apple -> 2}                          |
# +-----+-----------------------+-----+------+------+--------------------------------------+

字符串
第一个fruitcols_arr创建一个Map数组（column_name -> column_value）使用每个单独的fruit列.第二个过滤器基于fruits列数组元素的数组.这是基于每个fruits元素的Map的最终数组. new_col是通过使用aggregate高阶函数与map_concat创建的，连接最终Map数组中的所有单个Map。

展开查看全部

赞(0）回复(0）举报 2024-01-06

我来回答

在pyspark对象框架中的数组上迭代，并基于与数组中的值同名的列创建新列

1条答案

相关问题

热门标签

最新问答