在 PySpark 中定义一个 UDF，其中返回类型基于列

Question

我必须从 REST API 中提取数据，然后将结果转换为表格。 作为参考，我使用这个链接。 我的问题来自链接中模式的定义。 这是一个例子。

def executeRestApi(verb, url, headers, body):
  #
  headers = {
      'content-type': "application/json"
  }

  res = None
  # Make API request, get response object back, create dataframe from above schema.
  try:
    if verb == "get":
      res = requests.get(url, data=body, headers=headers)
    else:
      res = requests.post(url, data=body, headers=headers)
  except Exception as e:
    return e

  if res != None and res.status_code == 200:
    return json.loads(res.text)

  return None

schema = StructType([
  StructField("Count", IntegerType(), True),
  StructField("Message", StringType(), True),
  StructField("SearchCriteria", StringType(), True),
  StructField("Results", ArrayType(
    StructType([
      StructField("Make_ID", IntegerType()),
      StructField("Make_Name", StringType())
    ])
  ))
])

这是UDF的定义：

udf_executeRestApi = udf(executeRestApi, schema)

就我而言，每个 REST API 的架构都不同。 我有一个存储表名和模式的表：

my_table = spark.createDataFrame([["A", "schema_A", "get"], ["B", "schema_B", "get"]], schema=["TableName", "Schema", "verb"])\
    .withColumn("url", F.concat(F.lit("my_url/"), F.col("TableName")))
+---------+--------+----+--------+
|TableName|  Schema|verb|     url|
+---------+--------+----+--------+
|        A|schema_A| get|my_url/A|
|        B|schema_B| get|my_url/B|
+---------+--------+----+--------+

我该如何申请：

my_table.withColumn("result", udf_executeRestApi(col("verb"), col("url"), col("headers"), col("body")))

架构基于列Schema在哪里？

Answer 1

所以这里的目标基本上是并行化 API 调用，但不同 API 的架构不匹配。 我们先不要从 Spark 开始。 我们可以在 Python 和 Pandas 中编写代码，然后用Fugue轻松将其带到 Spark 中。

首先进行一些设置。 我正在使用口袋妖怪 API 来获得端到端的示例。 假设 URL 是不同的。 我们首先不关心 Schema。

import pandas as pd
from typing import List, Iterable, Any, Dict
import requests as re
import pickle

df = pd.DataFrame({"table": ["ditto", "pikachu"], 
                   "verb": ["get", "get"], 
                   "url":["https://pokeapi.co/api/v2/pokemon/ditto",
                          "https://pokeapi.co/api/v2/pokemon/pikachu"]})

我们可以为 Pandas DataFrame 中的一行创建一个逻辑来执行 API 调用。 我们可以使用 pickle 来创建一个包含所有结果的二进制类型列，而不是提取模式。

def call_api(df: List[Dict[str,Any]]) -> Iterable[Dict[str,Any]]:
    # this is for one function call
    # if the dataframe coming in is a list of dict, this operation is easy
    # lets assume df has one row
    row = df[0]
    res = re.get(row["url"])
    row["result"] = pickle.dumps(res.text)
    yield row

Fugue 现在可以读取注释并应用转换。 我们所要做的就是使用transform function 来测试：

from fugue import transform
transform(df.iloc[0:1], call_api, schema="*,result:binary")

Schema 是 Spark 的要求，所以我们在这里也需要它。 这会给你这样的东西：

table   verb    url     result
ditto   get     https://pokeapi.co/api/v2/pokemon/ditto 
b'\x80\x04\x95\xc2X\x00\x00\x00\x00\x00\x00X\x...

它未对齐，但请注意结果是二进制的。 这适用于一行，那么我们如何将它应用于每一行？ 我们只是传入一个分区策略。

transform(df, call_api, schema="*,result:binary", partition={"by":"url"})

我不知道您将在哪个列上进行分区，但目标是每个分区 1 行，然后我们可以并行化 DataFrame 的行。 在我的示例中，url 是独一无二的。

然后现在我们可以把它带到 Spark 上。 我们只需要将 SparkSession 传递给同一个transform调用，它就会在 Spark 上运行。

from pyspark.sql import SparkSession
spark = SparkSession.builder.getOrCreate()
sdf = transform(df, call_api, schema="*,result:binary", partition={"by":"url"},engine=spark)
sdf.show()

这将给出 output。

几点注意事项：

1. 如果您想稍后使用 API 结果，只需将列解压缩为原始形式即可。 使用 pickle.loads
2. 如果您想在 pickle 之前对结果进行子集化，只需将逻辑添加到call_api并添加一些 if else
3. Fugue transform()的模式表达式将转换为 Spark 的模式表达式。

如果您需要更多帮助，请随时给我留言。 个人资料中的联系信息。