使用 If Else 逻辑提取结构的嵌套数组

Question

我需要从以下模式结构中派生出两个新字段 - 新列 beaver_id 和 llama_id。有一些 if else 逻辑需要应用于结构数组。所需的最终结果是 csv 输出。最好的方法是什么？

架构：

root
 |-- Animal: struct (nullable = true)
 |    |-- Species: array (nullable = true)
 |    |    |-- element: struct (containsNull = true)
 |    |    |    |-- name: string (nullable = true)
 |    |    |    |-- color: string (nullable = true)
 |    |    |    |-- unique_id: string (nullable = true)

伪代码：

If name == "Beaver" 
   then get unique_id and put in dataframe column "beaver_id"
     else
   null in column "beaver_id"

If name == "Llama"
   then get unique_id and put in dataframe column "llama_id"
     else
   null in column "llama_id"

If array of names does not contain "Llama" or "Beaver"
   then null for both "beaver_id" and "llama_id"

目前：我正在使用 select 数据帧函数从输入（parquet）中提取 select 元素来创建 csv 输出。除了这个问题中的那些之外，我正在使用这种方法提取许多其他元素。

var select_df = raw_df.select(
  col(Animal.Species.name).getField("name")
  col(Animal.Species.color).getField("color")
)

示例输入（在 JSON 中），实际输入是镶木地板：

{
  "Animal": {
    "Species": [
      {
        "name": "Beaver",
        "color": "red",
        "unique_id": "1001"
      },
      {
        "name": "Llama",
        "color": "blue",
        "unique_id": "2222"
      }
    ]
  }
}

预期的 csv 输出：

beaver_id, llama_id
1001, 2222

Answer 1

当您必须处理大量检查或复杂检查时，最好始终使用可简化此任务并增加可伸缩性的模式策略，以防您想添加更多检查。

从这个意义上说，您创建一个上下文来确定您正在验证的记录是哪种“物种”，以便应用您需要的逻辑，在这种情况下，它似乎更方便将是所需值的提取及其随后分配给由上下文确定的字段。

如果我正确理解了你的问题，我认为这将是完成这项任务的最佳方法，如果我的观点适合你的情况以及你对此有任何疑问，请告诉我。

Answer 2

您可以像这样在 Animal.Species 数组列上使用 filter 函数：

val select_df = raw_df.select(
  element_at(expr("filter(Animal.Species, x -> x.name = 'Beaver')"), 1)
    .getField("unique_id")
    .as("beaver_id"),
  element_at(expr("filter(Animal.Species, x -> x.name = 'Llama')"), 1)
    .getField("unique_id")
    .as("llama_id")
)

select_df.show
//+---------+--------+
//|beaver_id|llama_id|
//+---------+--------+
//|     1001|    2222|
//+---------+--------+

逻辑很简单，我们过滤数组找到包含name = Beaver|Llama的内部结构并得到它的unique_id。如果找不到，则返回 null。

请注意，自 Spark 3 起，您还可以在 DataFrame API.

中使用 filter 函数

使用 If Else 逻辑提取结构的嵌套数组

Extract Nested Array of Struct with If Else Logic

scala

dataframe

apache-spark

apache-spark-sql

aws-glue