absl::Status PromoteArrayDataToType()

in tfx_bsl/cc/arrow/table_util.cc [158:225]


absl::Status PromoteArrayDataToType(
    const std::shared_ptr<ArrayData>& array_data,
    const std::shared_ptr<DataType>& target_type,
    std::shared_ptr<ArrayData>* promoted) {
  const auto& current_type = array_data->type;
  if (current_type->Equals(target_type)) {
    *promoted = array_data;
    return absl::OkStatus();
  }
  if (current_type->id() == Type::NA) {
    std::shared_ptr<Array> array_of_null;
    TFX_BSL_ASSIGN_OR_RETURN_ARROW(
        array_of_null, arrow::MakeArrayOfNull(target_type, array_data->length));
    *promoted = array_of_null->data();
    return absl::OkStatus();
  }
  if (current_type->id() != target_type->id()) {
    return absl::FailedPreconditionError(absl::StrCat(
        "Unable to promote array to incompatible type: ",
        current_type->ToString(), " vs ", target_type->ToString()));
  }
  switch (target_type->id()) {
      case Type::LIST:
      case Type::LARGE_LIST: {
        std::shared_ptr<ArrayData> promoted_value_array_data;
        TFX_BSL_RETURN_IF_ERROR(PromoteArrayDataToType(
            array_data->child_data[0], GetListValueType(*target_type),
            &promoted_value_array_data));

        *promoted = array_data->Copy();
        (*promoted)->type = target_type;
        (*promoted)->child_data[0] = std::move(promoted_value_array_data);
        break;
      }
      case Type::STRUCT: {
        const auto& target_struct_type =
            static_cast<const arrow::StructType&>(*target_type);
        const auto& current_struct_type =
            static_cast<const arrow::StructType&>(*array_data->type);
        const int target_num_fields = target_struct_type.num_fields();
        std::vector<std::shared_ptr<ArrayData>> child_data(target_num_fields);
        for (int i = 0; i < target_num_fields; ++i) {
          const std::string& field_name = target_struct_type.field(i)->name();
          const int current_field_index =
              current_struct_type.GetFieldIndex(field_name);
          if (current_field_index < 0) {
            std::shared_ptr<Array> array_of_null;
            TFX_BSL_ASSIGN_OR_RETURN_ARROW(
                array_of_null,
                arrow::MakeArrayOfNull(target_struct_type.field(i)->type(),
                                       array_data->length));
            child_data[i] = array_of_null->data();
          } else {
            TFX_BSL_RETURN_IF_ERROR(PromoteArrayDataToType(
                array_data->child_data[current_field_index],
                target_struct_type.field(i)->type(), &child_data[i]));
          }
        }
        *promoted = array_data->Copy();
        (*promoted)->type = target_type;
        (*promoted)->child_data = std::move(child_data);
        break;
      }
      default:
        return absl::UnimplementedError("Not implemented");
  }
  return absl::OkStatus();
}