ARROW-14615: [C++] Refactor nested field refs and add union support

cpp/src/arrow/array/array_nested.cc

@@ -35,6 +35,7 @@
 #include "arrow/type_traits.h"
 #include "arrow/util/atomic_shared_ptr.h"
 #include "arrow/util/bit_util.h"
+#include "arrow/util/bitmap_generate.h"
 #include "arrow/util/bitmap_ops.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/logging.h"
@@ -650,6 +651,43 @@ SparseUnionArray::SparseUnionArray(std::shared_ptr<DataType> type, int64_t lengt
   SetData(std::move(internal_data));
 }
 
+Result<std::shared_ptr<Array>> SparseUnionArray::GetFlattenedField(
+    int index, MemoryPool* pool) const {
+  if (index < 0 || index >= num_fields()) {
+    return Status::Invalid("Index out of range: ", index);
+  }
+  auto child_data = data_->child_data[index]->Copy();
+  // Adjust the result offset/length to be absolute.
+  if (data_->offset != 0 || data_->length != child_data->length) {
+    child_data = child_data->Slice(data_->offset, data_->length);
+  }
+  std::shared_ptr<Buffer> child_null_bitmap = child_data->buffers[0];
+  const int64_t child_offset = child_data->offset;
+
+  // Synthesize a null bitmap based on the union discriminant.
+  // Make sure the bitmap has extra bits corresponding to the child offset.
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> flattened_null_bitmap,
+                        AllocateEmptyBitmap(child_data->length + child_offset, pool));
+  const int8_t type_code = union_type()->type_codes()[index];
+  const int8_t* type_codes = raw_type_codes();
+  int64_t offset = 0;
+  internal::GenerateBitsUnrolled(flattened_null_bitmap->mutable_data(), child_offset,
+                                 data_->length,
+                                 [&] { return type_codes[offset++] == type_code; });
+
+  // The validity of a flattened datum is the logical AND of the synthesized
+  // null bitmap buffer and the individual field element's validity.
+  if (child_null_bitmap) {
+    BitmapAnd(flattened_null_bitmap->data(), child_offset, child_null_bitmap->data(),
+              child_offset, child_data->length, child_offset,
+              flattened_null_bitmap->mutable_data());
+  }
+
+  child_data->buffers[0] = std::move(flattened_null_bitmap);
+  child_data->null_count = kUnknownNullCount;
+  return MakeArray(child_data);
+}
+
 DenseUnionArray::DenseUnionArray(const std::shared_ptr<ArrayData>& data) {
   SetData(data);
 }

cpp/src/arrow/array/array_nested.h

@@ -464,6 +464,14 @@ class ARROW_EXPORT SparseUnionArray : public UnionArray {
     return internal::checked_cast<const SparseUnionType*>(union_type_);
   }
 
+  /// \brief Get one of the child arrays, adjusting its null bitmap
+  /// where the union array type code does not match.
+  ///
+  /// \param[in] index Which child array to get (i.e. the physical index, not the type
+  /// code) \param[in] pool The pool to allocate null bitmaps from, if necessary
+  Result<std::shared_ptr<Array>> GetFlattenedField(
+      int index, MemoryPool* pool = default_memory_pool()) const;
+
  protected:
   void SetData(std::shared_ptr<ArrayData> data);
 };

cpp/src/arrow/array/array_union_test.cc

@@ -32,6 +32,7 @@
 namespace arrow {
 
 using internal::checked_cast;
+using internal::checked_pointer_cast;
 
 TEST(TestUnionArray, TestSliceEquals) {
   std::shared_ptr<RecordBatch> batch;
@@ -68,6 +69,70 @@ TEST(TestUnionArray, TestSliceEquals) {
   CheckUnion(batch->column(1));
 }
 
+TEST(TestSparseUnionArray, GetFlattenedField) {
+  auto ty = sparse_union({field("ints", int64()), field("strs", utf8())}, {2, 7});
+  auto ints = ArrayFromJSON(int64(), "[0, 1, 2, 3]");
+  auto strs = ArrayFromJSON(utf8(), R"(["a", null, "c", "d"])");
+  auto ids = ArrayFromJSON(int8(), "[2, 7, 2, 7]")->data()->buffers[1];
+  const int length = 4;
+
+  {
+    SparseUnionArray arr(ty, length, {ints, strs}, ids);
+    ASSERT_OK(arr.ValidateFull());
+
+    ASSERT_OK_AND_ASSIGN(auto flattened, arr.GetFlattenedField(0));
+    AssertArraysEqual(*ArrayFromJSON(int64(), "[0, null, 2, null]"), *flattened,
+                      /*verbose=*/true);
+
+    ASSERT_OK_AND_ASSIGN(flattened, arr.GetFlattenedField(1));
+    AssertArraysEqual(*ArrayFromJSON(utf8(), R"([null, null, null, "d"])"), *flattened,
+                      /*verbose=*/true);
+
+    const auto sliced = checked_pointer_cast<SparseUnionArray>(arr.Slice(1, 2));
+
+    ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(0));
+    AssertArraysEqual(*ArrayFromJSON(int64(), "[null, 2]"), *flattened, /*verbose=*/true);
+
+    ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(1));
+    AssertArraysEqual(*ArrayFromJSON(utf8(), R"([null, null])"), *flattened,
+                      /*verbose=*/true);
+
+    ASSERT_RAISES(Invalid, arr.GetFlattenedField(-1));
+    ASSERT_RAISES(Invalid, arr.GetFlattenedField(2));
+  }
+  {
+    SparseUnionArray arr(ty, length - 2, {ints->Slice(1, 2), strs->Slice(1, 2)}, ids);
+    ASSERT_OK(arr.ValidateFull());
+
+    ASSERT_OK_AND_ASSIGN(auto flattened, arr.GetFlattenedField(0));
+    AssertArraysEqual(*ArrayFromJSON(int64(), "[1, null]"), *flattened, /*verbose=*/true);
+
+    ASSERT_OK_AND_ASSIGN(flattened, arr.GetFlattenedField(1));
+    AssertArraysEqual(*ArrayFromJSON(utf8(), R"([null, "c"])"), *flattened,
+                      /*verbose=*/true);
+
+    const auto sliced = checked_pointer_cast<SparseUnionArray>(arr.Slice(1, 1));
+
+    ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(0));
+    AssertArraysEqual(*ArrayFromJSON(int64(), "[null]"), *flattened, /*verbose=*/true);
+
+    ASSERT_OK_AND_ASSIGN(flattened, sliced->GetFlattenedField(1));
+    AssertArraysEqual(*ArrayFromJSON(utf8(), R"(["c"])"), *flattened, /*verbose=*/true);
+  }
+  {
+    SparseUnionArray arr(ty, /*length=*/0, {ints->Slice(length), strs->Slice(length)},
+                         ids);
+    ASSERT_OK(arr.ValidateFull());
+
+    ASSERT_OK_AND_ASSIGN(auto flattened, arr.GetFlattenedField(0));
+    AssertArraysEqual(*ArrayFromJSON(int64(), "[]"), *flattened, /*verbose=*/true);
+
+    ASSERT_OK_AND_ASSIGN(flattened, arr.GetFlattenedField(1));
+    AssertArraysEqual(*ArrayFromJSON(utf8(), "[]"), *flattened,
+                      /*verbose=*/true);
+  }
+}
+
 TEST(TestSparseUnionArray, Validate) {
   auto a = ArrayFromJSON(int32(), "[4, 5]");
   auto type = sparse_union({field("a", int32())});

cpp/src/arrow/compute/api_scalar.cc

@@ -223,6 +223,8 @@ static auto kExtractRegexOptionsType = GetFunctionOptionsType<ExtractRegexOption
 static auto kSetLookupOptionsType = GetFunctionOptionsType<SetLookupOptions>(
     DataMember("value_set", &SetLookupOptions::value_set),
     DataMember("skip_nulls", &SetLookupOptions::skip_nulls));
+static auto kStructFieldOptionsType = GetFunctionOptionsType<StructFieldOptions>(
+    DataMember("indices", &StructFieldOptions::indices));
 static auto kStrptimeOptionsType = GetFunctionOptionsType<StrptimeOptions>(
     DataMember("format", &StrptimeOptions::format),
     DataMember("unit", &StrptimeOptions::unit));
@@ -351,6 +353,11 @@ SetLookupOptions::SetLookupOptions(Datum value_set, bool skip_nulls)
 SetLookupOptions::SetLookupOptions() : SetLookupOptions({}, false) {}
 constexpr char SetLookupOptions::kTypeName[];
 
+StructFieldOptions::StructFieldOptions(std::vector<int> indices)
+    : FunctionOptions(internal::kStructFieldOptionsType), indices(std::move(indices)) {}
+StructFieldOptions::StructFieldOptions() : StructFieldOptions(std::vector<int>()) {}
+constexpr char StructFieldOptions::kTypeName[];
+
 StrptimeOptions::StrptimeOptions(std::string format, TimeUnit::type unit)
     : FunctionOptions(internal::kStrptimeOptionsType),
       format(std::move(format)),
@@ -444,6 +451,7 @@ void RegisterScalarOptions(FunctionRegistry* registry) {
   DCHECK_OK(registry->AddFunctionOptionsType(kReplaceSubstringOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kExtractRegexOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kSetLookupOptionsType));
+  DCHECK_OK(registry->AddFunctionOptionsType(kStructFieldOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kStrptimeOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kStrftimeOptionsType));
   DCHECK_OK(registry->AddFunctionOptionsType(kAssumeTimezoneOptionsType));

cpp/src/arrow/compute/api_scalar.h

@@ -223,6 +223,18 @@ class ARROW_EXPORT SetLookupOptions : public FunctionOptions {
   bool skip_nulls;
 };
 
+/// Options for struct_field function
+class ARROW_EXPORT StructFieldOptions : public FunctionOptions {
+ public:
+  explicit StructFieldOptions(std::vector<int> indices);
+  StructFieldOptions();
+  constexpr static char const kTypeName[] = "StructFieldOptions";
+
+  /// The child indices to extract. For instance, to get the 2nd child
+  /// of the 1st child of a struct or union, this would be {0, 1}.
+  std::vector<int> indices;
+};
+
 class ARROW_EXPORT StrptimeOptions : public FunctionOptions {
  public:
   explicit StrptimeOptions(std::string format, TimeUnit::type unit);

cpp/src/arrow/compute/exec/expression.cc

@@ -510,29 +510,11 @@ Result<Datum> ExecuteScalarExpression(const Expression& expr, const ExecBatch& i
     }
 
     Datum field = input[param->indices[0]];
-    for (auto it = param->indices.begin() + 1; it != param->indices.end(); ++it) {
-      if (field.type()->id() != Type::STRUCT) {
-        return Status::Invalid("Nested field reference into a non-struct: ",
-                               *field.type());
-      }
-      const int index = *it;
-      if (index < 0 || index >= field.type()->num_fields()) {
-        return Status::Invalid("Out of bounds field reference: ", index, " but type has ",
-                               field.type()->num_fields(), " fields");
-      }
-      if (field.is_scalar()) {
-        const auto& struct_scalar = field.scalar_as<StructScalar>();
-        if (!struct_scalar.is_valid) {
-          return MakeNullScalar(param->descr.type);
-        }
-        field = struct_scalar.value[index];
-      } else if (field.is_array()) {
-        const auto& struct_array = field.array_as<StructArray>();
-        ARROW_ASSIGN_OR_RAISE(
-            field, struct_array->GetFlattenedField(index, exec_context->memory_pool()));
-      } else {
-        return Status::NotImplemented("Nested field reference into a ", field.ToString());
-      }
+    if (param->indices.size() > 1) {
+      std::vector<int> indices(param->indices.begin() + 1, param->indices.end());
+      compute::StructFieldOptions options(std::move(indices));
+      ARROW_ASSIGN_OR_RAISE(
+          field, compute::CallFunction("struct_field", {std::move(field)}, &options));
     }
     if (!field.type()->Equals(param->descr.type)) {
       return Status::Invalid("Referenced field ", expr.ToString(), " was ",