blob: bbb231356ec6db201878cd40e8c7fa1327803982 [file]
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
/* This file is based on source code from the Spark Project (http://spark.apache.org/), licensed
* by the Apache Software Foundation (ASF) under the Apache License, Version 2.0. See the NOTICE
* file distributed with this work for additional information regarding copyright ownership. */
#include "paimon/common/data/variant/variant_shredding_utils.h"
#include <string>
#include <utility>
#include <vector>
#include "arrow/api.h"
#include "arrow/util/checked_cast.h"
#include "fmt/format.h"
#include "paimon/common/data/variant/variant_defs.h"
#include "paimon/common/data/variant/variant_type_utils.h"
namespace paimon {
namespace {
Status InvalidVariantShreddingSchema(const std::shared_ptr<arrow::DataType>& type) {
return Status::Invalid(
fmt::format("Invalid variant shredding schema: {}", type ? type->ToString() : "null"));
}
// Mirrors the Java `PaimonShreddingUtils.variantShreddingSchema(dataType, isTopLevel,
// isObjectField)`.
Result<std::shared_ptr<arrow::DataType>> VariantShreddingSchemaImpl(
const std::shared_ptr<arrow::DataType>& data_type, bool is_top_level, bool is_object_field) {
arrow::FieldVector fields;
if (is_top_level) {
fields.push_back(arrow::field(VariantDefs::kMetadataFieldName, arrow::binary(),
/*nullable=*/false));
}
switch (data_type->id()) {
case arrow::Type::LIST: {
const auto& list_type = std::static_pointer_cast<arrow::ListType>(data_type);
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::DataType> element_type,
VariantShreddingSchemaImpl(list_type->value_type(),
/*is_top_level=*/false,
/*is_object_field=*/false));
fields.push_back(
arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true));
fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName,
arrow::list(element_type), /*nullable=*/true));
break;
}
case arrow::Type::STRUCT: {
// The field name level is always non-nullable: Variant null values are represented in
// the "value" column as "00", and missing values are represented by setting both
// "value" and "typed_value" to null.
const auto& struct_type = std::static_pointer_cast<arrow::StructType>(data_type);
arrow::FieldVector shredded_fields;
for (const auto& field : struct_type->fields()) {
PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::DataType> field_type,
VariantShreddingSchemaImpl(field->type(),
/*is_top_level=*/false,
/*is_object_field=*/true));
shredded_fields.push_back(
arrow::field(field->name(), field_type, /*nullable=*/false));
}
fields.push_back(
arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true));
fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName,
arrow::struct_(shredded_fields), /*nullable=*/true));
break;
}
case arrow::Type::NA: {
// `arrow::null()` denotes an untyped VARIANT leaf in shredding types. It doesn't
// need a typed column. If there is no typed column, value is required for array
// elements or top-level fields, but optional for objects (where a null represents a
// missing field).
fields.push_back(arrow::field(VariantDefs::kValueFieldName, arrow::binary(),
/*nullable=*/is_object_field));
break;
}
case arrow::Type::STRING:
case arrow::Type::BOOL:
case arrow::Type::BINARY:
case arrow::Type::DECIMAL128:
case arrow::Type::INT8:
case arrow::Type::INT16:
case arrow::Type::INT32:
case arrow::Type::INT64:
case arrow::Type::FLOAT:
case arrow::Type::DOUBLE: {
fields.push_back(
arrow::field(VariantDefs::kValueFieldName, arrow::binary(), /*nullable=*/true));
fields.push_back(arrow::field(VariantDefs::kTypedValueFieldName, data_type,
/*nullable=*/true));
break;
}
default:
return InvalidVariantShreddingSchema(data_type);
}
return arrow::struct_(fields);
}
Result<std::shared_ptr<VariantSchema>> BuildVariantSchemaImpl(
const std::shared_ptr<arrow::DataType>& type, bool top_level) {
if (type->id() != arrow::Type::STRUCT) {
return InvalidVariantShreddingSchema(type);
}
const auto& struct_type = std::static_pointer_cast<arrow::StructType>(type);
// The struct must not be empty or contain duplicate field names. The latter is enforced in
// the loop below.
if (struct_type->num_fields() == 0) {
return InvalidVariantShreddingSchema(type);
}
auto schema = std::make_shared<VariantSchema>();
schema->num_fields = struct_type->num_fields();
for (int32_t i = 0; i < struct_type->num_fields(); ++i) {
const auto& field = struct_type->field(i);
const auto& field_type = field->type();
if (field->name() == VariantDefs::kTypedValueFieldName) {
if (schema->typed_idx != -1) {
return InvalidVariantShreddingSchema(type);
}
schema->typed_idx = i;
switch (field_type->id()) {
case arrow::Type::STRUCT: {
const auto& object_type =
std::static_pointer_cast<arrow::StructType>(field_type);
schema->has_object_schema = true;
schema->object_schema.reserve(object_type->num_fields());
for (int32_t index = 0; index < object_type->num_fields(); ++index) {
const auto& object_field = object_type->field(index);
PAIMON_ASSIGN_OR_RAISE(
std::shared_ptr<VariantSchema> field_schema,
BuildVariantSchemaImpl(object_field->type(), /*top_level=*/false));
schema->object_schema.push_back(
VariantSchema::ObjectField{object_field->name(), field_schema});
auto [it, inserted] =
schema->object_schema_map.emplace(object_field->name(), index);
if (!inserted) {
return InvalidVariantShreddingSchema(type);
}
}
break;
}
case arrow::Type::LIST: {
const auto& list_type = std::static_pointer_cast<arrow::ListType>(field_type);
PAIMON_ASSIGN_OR_RAISE(
schema->array_schema,
BuildVariantSchemaImpl(list_type->value_type(), /*top_level=*/false));
break;
}
case arrow::Type::BOOL:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kBoolean};
break;
case arrow::Type::INT8:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kByte};
break;
case arrow::Type::INT16:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kShort};
break;
case arrow::Type::INT32:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kInt};
break;
case arrow::Type::INT64:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kLong};
break;
case arrow::Type::FLOAT:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kFloat};
break;
case arrow::Type::DOUBLE:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kDouble};
break;
case arrow::Type::STRING:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kString};
break;
case arrow::Type::BINARY:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kBinary};
break;
case arrow::Type::DATE32:
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kDate};
break;
case arrow::Type::DECIMAL128: {
const auto& decimal_type =
std::static_pointer_cast<arrow::Decimal128Type>(field_type);
schema->scalar_schema =
VariantSchema::ScalarType{VariantSchema::ScalarKind::kDecimal,
decimal_type->precision(), decimal_type->scale()};
break;
}
case arrow::Type::TIMESTAMP: {
const auto& timestamp_type =
std::static_pointer_cast<arrow::TimestampType>(field_type);
// The variant binary stores timestamps as microseconds since the epoch; a
// typed_value column of any other precision would misinterpret the values.
if (timestamp_type->unit() != arrow::TimeUnit::MICRO) {
return InvalidVariantShreddingSchema(type);
}
schema->scalar_schema =
VariantSchema::ScalarType{timestamp_type->timezone().empty()
? VariantSchema::ScalarKind::kTimestampNtz
: VariantSchema::ScalarKind::kTimestampLtz};
break;
}
default:
return InvalidVariantShreddingSchema(type);
}
} else if (field->name() == VariantDefs::kValueFieldName) {
if (schema->variant_idx != -1 || field_type->id() != arrow::Type::BINARY) {
return InvalidVariantShreddingSchema(type);
}
schema->variant_idx = i;
} else if (field->name() == VariantDefs::kMetadataFieldName) {
if (schema->top_level_metadata_idx != -1 || field_type->id() != arrow::Type::BINARY) {
return InvalidVariantShreddingSchema(type);
}
schema->top_level_metadata_idx = i;
} else {
return InvalidVariantShreddingSchema(type);
}
}
if (top_level != (schema->top_level_metadata_idx >= 0)) {
return InvalidVariantShreddingSchema(type);
}
return schema;
}
} // namespace
Result<std::shared_ptr<arrow::DataType>> VariantShreddingUtils::VariantShreddingSchema(
const std::shared_ptr<arrow::DataType>& shredding_type) {
return VariantShreddingSchemaImpl(shredding_type, /*is_top_level=*/true,
/*is_object_field=*/false);
}
Result<std::shared_ptr<VariantSchema>> VariantShreddingUtils::BuildVariantSchema(
const std::shared_ptr<arrow::DataType>& struct_type) {
return BuildVariantSchemaImpl(struct_type, /*top_level=*/true);
}
Result<std::shared_ptr<arrow::DataType>> VariantShreddingUtils::ScalarSchemaToArrowType(
const VariantSchema::ScalarType& scalar) {
switch (scalar.kind) {
case VariantSchema::ScalarKind::kBoolean:
return arrow::boolean();
case VariantSchema::ScalarKind::kByte:
return arrow::int8();
case VariantSchema::ScalarKind::kShort:
return arrow::int16();
case VariantSchema::ScalarKind::kInt:
return arrow::int32();
case VariantSchema::ScalarKind::kLong:
return arrow::int64();
case VariantSchema::ScalarKind::kFloat:
return arrow::float32();
case VariantSchema::ScalarKind::kDouble:
return arrow::float64();
case VariantSchema::ScalarKind::kString:
return arrow::utf8();
case VariantSchema::ScalarKind::kBinary:
return arrow::binary();
case VariantSchema::ScalarKind::kDecimal:
return arrow::decimal128(scalar.precision, scalar.scale);
case VariantSchema::ScalarKind::kDate:
return arrow::date32();
case VariantSchema::ScalarKind::kTimestampLtz:
return arrow::timestamp(arrow::TimeUnit::MICRO, "UTC");
case VariantSchema::ScalarKind::kTimestampNtz:
return arrow::timestamp(arrow::TimeUnit::MICRO);
default:
return Status::NotImplemented(fmt::format("Unsupported variant scalar kind: {}",
static_cast<int32_t>(scalar.kind)));
}
}
bool VariantShreddingUtils::IsShreddedFileType(
const std::shared_ptr<arrow::DataType>& file_variant_type) {
if (!file_variant_type || file_variant_type->id() != arrow::Type::STRUCT) {
return false;
}
const auto& struct_type = std::static_pointer_cast<arrow::StructType>(file_variant_type);
return struct_type->GetFieldByName(VariantDefs::kTypedValueFieldName) != nullptr;
}
bool VariantShreddingUtils::IsUntypedPhysicalVariantType(
const std::shared_ptr<arrow::DataType>& file_variant_type) {
if (!file_variant_type || file_variant_type->id() != arrow::Type::STRUCT) {
return false;
}
const auto& struct_type = std::static_pointer_cast<arrow::StructType>(file_variant_type);
if (struct_type->num_fields() != 2) {
return false;
}
const auto& metadata = struct_type->field(0);
const auto& value = struct_type->field(1);
return metadata->name() == VariantDefs::kMetadataFieldName &&
metadata->type()->id() == arrow::Type::BINARY &&
value->name() == VariantDefs::kValueFieldName &&
value->type()->id() == arrow::Type::BINARY;
}
} // namespace paimon