PySpark/Aws Glue 中的性能問題

Question

我有一個數據框。 我需要將每條記錄轉換為 JSON，然后使用 JSON 有效負載調用 API 以將數據插入 postgress。 我在數據框中有 14000 條記錄，要調用 api 並得到響應，需要 5 小時。 有什么辦法可以提高性能。 下面是我的代碼片段。

df_insert = spark.read \
.format(SNOWFLAKE_SOURCE_NAME) \
.options(**sfOptions) \
.option("dbtable", "source_table_name") \
.load()

json_insert = df_insert.toJSON().collect()

for row in json_insert:
  line = json.loads(row)
    headers = {
    'Authorization': authorization,
    'content-type': "application/json",
    'cache-control': "no-cache",
    }
  response = requests.request("POST", url_insert, data=payload, headers=headers)
  print(response.text)
  res = response.text
  response_result = json.loads(res)
  #print(response_result["httpStatus"])
  if response_result["message"] == 'success':
      print ("INFO : Record inserted successfully")
  else:
      print ("ERROR : Error in the record")
      status_code = response_result["status"]
      error_message =  response_result["error"]
      my_list = [(status_code,error_message,row)]
      df = sc.createDataFrame(my_list, ['status', 'error', 'json data'])
      df.write.format(SNOWFLAKE_SOURCE_NAME) \
      .options(**sfOptions) \
      .option("dbtable", "error_table") \
      .option("header", "true") \
      .option("truncate_table", "on") \
      .mode("append") \
      .save()

注意：我知道通過執行“json_insert = df_insert.toJSON().collect()”我失去了數據框的優勢。 有沒有更好的方法來實現。

Answer 1

df_insert.toJSON()返回一個RDD ，您可以使用flatMap覆蓋它。 ¹

source_rdd = df_insert.toJSON()

對這個 RDD 執行flatMap並返回一個只包含錯誤的 RDD。

headers = {
    'Authorization': authorization,
    'content-type': "application/json",
    'cache-control': "no-cache"
}

def post_service_error(row):
    # requests package may not be available in the node
    # see about adding files to the spark context
    response = requests.request("POST", url_insert, data=row, headers=headers)
    response_result = response.json()
    if response_result['message'] == 'success':
        print ("INFO : Record inserted successfully")
        return []
    print ("ERROR : Error in the record")
    status_code = response_result["status"]
    error_message =  response_result["error"]
    return [(status_code, error_message, row)]

errors_rdd = source_rdd.flatMap(post_service_error)

將錯誤 RDD 轉換為 spark DataFrame 並將其保存到表中。

errors_df = sc.createDataFrame(errors_rdd, ['status', 'error', 'json data'])
(errors_df.write.format(SNOWFLAKE_SOURCE_NAME)
  .options(**sfOptions)
  .option("dbtable", "error_table")
  .option("header", "true")
  .option("truncate_table", "on")
  .mode("append")
  .save())

如果您擁有要向其發出請求的 API，我建議您探索一個接受一批這些對象/數組的實現。 通過這種方式，您可以在將每個分區映射到批處理請求之前對 RDD 進行分區，然后處理錯誤。