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();
}