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

Question

我需要从下面的模式结构派生两个新字段--新列beaver_id和llama_id。有一些(如果是其他逻辑)需要应用于一个结构数组。所需的最终结果是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 dataframe函数从输入(拼板)中选择元素来创建csv输出。除了这个问题之外，我还用这个方法提取了许多其他的元素。

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

示例输入(在JSON中)，实际输入是parquet：

{
  "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

可以在filter数组列上使用Animal.Species函数，如下所示：

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。

注意，从Spark3开始，您还可以在filter API中使用DataFrame函数。

Answer 2

当您必须处理大量的检查或复杂的检查时，最好使用模式策略来简化此任务，并在需要添加更多检查的情况下增加可伸缩性。

从这个意义上说，对我来说，创建一个上下文似乎更方便一些，它决定了什么类型的“物种”是您正在验证的记录，为了应用您需要的逻辑，在这种情况下，它将是提取所需的值并随后将其赋值到由上下文确定的字段。

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