qzyu999 commented on code in PR #50252: URL: https://github.com/apache/arrow/pull/50252#discussion_r3858546338
########## cpp/src/parquet/variant/shred.cc: ########## @@ -0,0 +1,857 @@ +// 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. + +#include "parquet/variant/shred.h" + +#include <concepts> +#include <cstdint> +#include <limits> +#include <memory> +#include <optional> +#include <string> +#include <string_view> +#include <unordered_set> +#include <utility> +#include <vector> + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/array/builder_base.h" +#include "arrow/array/builder_binary.h" +#include "arrow/buffer.h" +#include "arrow/buffer_builder.h" +#include "arrow/compute/cast.h" +#include "arrow/compute/exec.h" +#include "arrow/extension/parquet_variant.h" +#include "arrow/extension/uuid.h" +#include "arrow/extension_type.h" +#include "arrow/scalar.h" +#include "arrow/type.h" +#include "arrow/type_traits.h" +#include "arrow/util/bitmap_ops.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" +#include "arrow/util/ubsan.h" +#include "arrow/util/unreachable.h" +#include "parquet/exception.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/builder.h" +#include "parquet/variant/decoding.h" + +namespace parquet::variant { + +using ::arrow::Array; +using ::arrow::ArrayBuilder; +using ::arrow::DataType; +using ::arrow::Field; +using ::arrow::MemoryPool; +using ::arrow::Scalar; +using ::arrow::StructArray; +using ::arrow::TimeUnit; +using ::arrow::extension::VariantArray; +using ::arrow::internal::checked_cast; + +namespace { + +struct ShreddedArrayParts { + std::shared_ptr<Array> value; + std::shared_ptr<Array> typed_value; + std::shared_ptr<::arrow::Buffer> null_bitmap; + int64_t null_count = 0; +}; + +std::shared_ptr<DataType> FieldGroupType(const std::shared_ptr<DataType>& typed_type) { + return ::arrow::struct_({::arrow::field("value", ::arrow::binary_view()), + ::arrow::field("typed_value", typed_type)}); +} + +std::shared_ptr<StructArray> MakeFieldGroup(ShreddedArrayParts parts) { + auto type = FieldGroupType(parts.typed_value->type()); + PARQUET_ASSIGN_OR_THROW( + auto out, + StructArray::Make({std::move(parts.value), std::move(parts.typed_value)}, + type->fields(), std::move(parts.null_bitmap), parts.null_count)); + return out; +} + +class ShredNode { + public: + ShredNode(std::shared_ptr<DataType> typed_type, MemoryPool* pool) + : typed_type_(std::move(typed_type)), residual_(pool), group_validity_(pool) {} + virtual ~ShredNode() = default; + + const std::shared_ptr<DataType>& typed_type() const { return typed_type_; } + std::shared_ptr<DataType> field_group_type() const { + return FieldGroupType(typed_type_); + } + + virtual void AppendParentNull() = 0; + virtual void AppendMissing(const VariantMetadataView& metadata) = 0; + virtual void AppendValue(const VariantMetadataView& metadata, + const VariantValueView& value) = 0; + virtual ShreddedArrayParts Finish() = 0; + + protected: + void AppendGroupValidity(bool valid) { + PARQUET_THROW_NOT_OK(group_validity_.Append(valid)); + } + void AppendResidualNull() { residual_.AppendNull(); } + void AppendResidual(const VariantValueView& value) { + residual_.AppendEncodedValue(value.value()); + } + + ShreddedArrayParts FinishParts(std::shared_ptr<Array> typed_value) { + const auto null_count = group_validity_.false_count(); + return ShreddedArrayParts{.value = residual_.Finish(), + .typed_value = std::move(typed_value), + .null_bitmap = internal::FinishNullBitmap(group_validity_), + .null_count = null_count}; + } + + std::shared_ptr<DataType> typed_type_; + VariantValueArrayBuilder residual_; + ::arrow::TypedBufferBuilder<bool> group_validity_; +}; + +template <typename T> +T LoadLittleEndian(std::string_view bytes) { + return ::arrow::bit_util::FromLittleEndian( + ::arrow::util::SafeLoadAs<T>(reinterpret_cast<const uint8_t*>(bytes.data()))); +} + +bool TryAppendString(const DataType& storage_type, ArrayBuilder& builder, + std::string_view value) { + switch (storage_type.id()) { + case ::arrow::Type::STRING: + PARQUET_THROW_NOT_OK(checked_cast<::arrow::StringBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::LARGE_STRING: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::LargeStringBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::STRING_VIEW: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::StringViewBuilder&>(builder).Append(value)); + return true; + default: + return false; + } +} + +bool TryAppendBinary(const DataType& storage_type, ArrayBuilder& builder, + std::string_view value) { + switch (storage_type.id()) { + case ::arrow::Type::BINARY: + PARQUET_THROW_NOT_OK(checked_cast<::arrow::BinaryBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::LARGE_BINARY: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::LargeBinaryBuilder&>(builder).Append(value)); + return true; + case ::arrow::Type::BINARY_VIEW: + PARQUET_THROW_NOT_OK( + checked_cast<::arrow::BinaryViewBuilder&>(builder).Append(value)); + return true; + default: + return false; + } +} + +bool TryAppendPrimitiveDirect(const VariantValueView& value, const DataType& storage_type, + ArrayBuilder& builder) { + if (const auto* short_string = std::get_if<VariantShortStringView>(&value.data())) { + return TryAppendString(storage_type, builder, short_string->string()); + } + const auto* primitive = std::get_if<VariantPrimitiveView>(&value.data()); + if (primitive == nullptr) { + return false; + } + + const auto payload = primitive->payload(); + switch (primitive->type()) { + case VariantPrimitiveType::kBinary: + return TryAppendBinary(storage_type, builder, payload.substr(4)); + case VariantPrimitiveType::kString: + return TryAppendString(storage_type, builder, payload.substr(4)); + default: + return false; + } +} + +std::shared_ptr<Scalar> DecodeSourceScalar(const VariantValueView& value) { + if (const auto* short_string = std::get_if<VariantShortStringView>(&value.data())) { + return std::make_shared<::arrow::StringScalar>(std::string(short_string->string())); + } + const auto* primitive = std::get_if<VariantPrimitiveView>(&value.data()); + if (primitive == nullptr) { + return nullptr; + } + + const auto payload = primitive->payload(); + switch (primitive->type()) { + case VariantPrimitiveType::kNull: + return nullptr; + case VariantPrimitiveType::kBooleanTrue: + return std::make_shared<::arrow::BooleanScalar>(true); + case VariantPrimitiveType::kBooleanFalse: + return std::make_shared<::arrow::BooleanScalar>(false); + case VariantPrimitiveType::kInt8: + return std::make_shared<::arrow::Int8Scalar>(LoadLittleEndian<int8_t>(payload)); + case VariantPrimitiveType::kInt16: + return std::make_shared<::arrow::Int16Scalar>(LoadLittleEndian<int16_t>(payload)); + case VariantPrimitiveType::kInt32: + return std::make_shared<::arrow::Int32Scalar>(LoadLittleEndian<int32_t>(payload)); + case VariantPrimitiveType::kInt64: + return std::make_shared<::arrow::Int64Scalar>(LoadLittleEndian<int64_t>(payload)); + case VariantPrimitiveType::kFloat: + return std::make_shared<::arrow::FloatScalar>(LoadLittleEndian<float>(payload)); + case VariantPrimitiveType::kDouble: + return std::make_shared<::arrow::DoubleScalar>(LoadLittleEndian<double>(payload)); + case VariantPrimitiveType::kDecimal4: { + const auto scale = static_cast<uint8_t>(payload[0]); + auto type = ::arrow::decimal32(/*precision=*/9, scale); + return std::make_shared<::arrow::Decimal32Scalar>( + ::arrow::Decimal32(LoadLittleEndian<int32_t>(payload.substr(1))), type); + } + case VariantPrimitiveType::kDecimal8: { + const auto scale = static_cast<uint8_t>(payload[0]); + auto type = ::arrow::decimal64(/*precision=*/18, scale); + return std::make_shared<::arrow::Decimal64Scalar>( + ::arrow::Decimal64(LoadLittleEndian<int64_t>(payload.substr(1))), type); + } + case VariantPrimitiveType::kDecimal16: { + const auto scale = static_cast<uint8_t>(payload[0]); + const auto low = LoadLittleEndian<uint64_t>(payload.substr(1)); + const auto high = LoadLittleEndian<int64_t>(payload.substr(9)); + auto type = ::arrow::decimal128(/*precision=*/38, scale); + return std::make_shared<::arrow::Decimal128Scalar>(::arrow::Decimal128(high, low), + type); + } + case VariantPrimitiveType::kDate: + return std::make_shared<::arrow::Date32Scalar>(LoadLittleEndian<int32_t>(payload)); + case VariantPrimitiveType::kTimeNTZMicros: + return std::make_shared<::arrow::Time64Scalar>(LoadLittleEndian<int64_t>(payload), + TimeUnit::MICRO); + case VariantPrimitiveType::kTimestampMicros: + case VariantPrimitiveType::kTimestampNTZMicros: + case VariantPrimitiveType::kTimestampNanos: + case VariantPrimitiveType::kTimestampNTZNanos: { + const bool nanos = primitive->type() == VariantPrimitiveType::kTimestampNanos || + primitive->type() == VariantPrimitiveType::kTimestampNTZNanos; + const bool adjusted = primitive->type() == VariantPrimitiveType::kTimestampMicros || + primitive->type() == VariantPrimitiveType::kTimestampNanos; + const auto unit = nanos ? TimeUnit::NANO : TimeUnit::MICRO; + auto type = adjusted ? ::arrow::timestamp(unit, "UTC") : ::arrow::timestamp(unit); + return std::make_shared<::arrow::TimestampScalar>( + LoadLittleEndian<int64_t>(payload), type); + } + case VariantPrimitiveType::kBinary: + return std::make_shared<::arrow::BinaryScalar>(std::string(payload.substr(4))); + case VariantPrimitiveType::kString: + return std::make_shared<::arrow::StringScalar>(std::string(payload.substr(4))); + case VariantPrimitiveType::kUuid: + return std::make_shared<::arrow::FixedSizeBinaryScalar>( + ::arrow::Buffer::FromString(std::string(payload)), + ::arrow::fixed_size_binary(16)); + } + ::arrow::Unreachable("Unexpected Variant primitive type"); +} + +bool CanAttemptCast(const Scalar& source, const DataType& target) { + const auto source_id = source.type->id(); + switch (target.id()) { + case ::arrow::Type::BOOL: + return source_id == ::arrow::Type::BOOL || ::arrow::is_signed_integer(source_id) || + ::arrow::is_physical_floating(source_id) || + source_id == ::arrow::Type::STRING; + case ::arrow::Type::INT8: + case ::arrow::Type::INT16: + case ::arrow::Type::INT32: + case ::arrow::Type::INT64: + case ::arrow::Type::FLOAT: + case ::arrow::Type::DOUBLE: + return source_id == ::arrow::Type::BOOL || ::arrow::is_signed_integer(source_id) || + ::arrow::is_physical_floating(source_id) || ::arrow::is_decimal(source_id); + case ::arrow::Type::DECIMAL32: + case ::arrow::Type::DECIMAL64: + case ::arrow::Type::DECIMAL128: + return ::arrow::is_signed_integer(source_id) || + ::arrow::is_physical_floating(source_id) || ::arrow::is_decimal(source_id) || + source_id == ::arrow::Type::STRING; + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_STRING: + case ::arrow::Type::STRING_VIEW: + return source_id == ::arrow::Type::STRING; + case ::arrow::Type::BINARY: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::BINARY_VIEW: + return source_id == ::arrow::Type::BINARY; + case ::arrow::Type::DATE32: + return source_id == ::arrow::Type::DATE32; + case ::arrow::Type::TIME64: + return source.type->Equals(target); + case ::arrow::Type::FIXED_SIZE_BINARY: + return source.type->Equals(target); + case ::arrow::Type::TIMESTAMP: { + if (source_id != ::arrow::Type::TIMESTAMP) { + return false; + } + const auto& source_type = checked_cast<const ::arrow::TimestampType&>(*source.type); + const auto& target_type = checked_cast<const ::arrow::TimestampType&>(target); + return source_type.timezone().empty() == target_type.timezone().empty() && + (source_type.unit() == target_type.unit() || + (source_type.unit() == TimeUnit::MICRO && + target_type.unit() == TimeUnit::NANO)); + } + default: + return false; + } +} + +class PrimitiveShredNode : public ShredNode { + public: + PrimitiveShredNode(std::shared_ptr<DataType> target, ::arrow::compute::ExecContext* ctx, + MemoryPool* pool) + : ShredNode(std::move(target), pool), ctx_(ctx) { + if (typed_type_->id() == ::arrow::Type::EXTENSION) { + storage_type_ = + checked_cast<const ::arrow::ExtensionType&>(*typed_type_).storage_type(); + } else { + storage_type_ = typed_type_; + } + PARQUET_ASSIGN_OR_THROW(builder_, ::arrow::MakeBuilder(storage_type_, pool)); + } + + void AppendParentNull() override { + AppendGroupValidity(false); + AppendResidualNull(); + PARQUET_THROW_NOT_OK(builder_->AppendNull()); + } + + void AppendMissing(const VariantMetadataView&) override { + AppendGroupValidity(true); + AppendResidualNull(); + PARQUET_THROW_NOT_OK(builder_->AppendNull()); + } + + void AppendValue(const VariantMetadataView&, const VariantValueView& value) override { + AppendGroupValidity(true); + if (TryAppendPrimitiveDirect(value, *storage_type_, *builder_)) { + AppendResidualNull(); + return; + } + const auto source = DecodeSourceScalar(value); + if (source != nullptr && CanAttemptCast(*source, *storage_type_)) { + auto casted = ::arrow::compute::Cast(::arrow::Datum(source), storage_type_, + ::arrow::compute::CastOptions::Safe(), ctx_); Review Comment: This decodes each row into a scalar and calls `compute::Cast` individually. For large arrays this is a bottleneck since scalar Cast has non-trivial per-call overhead (kernel lookup, ExecContext dispatch, output scalar allocation). A follow-up optimization could accumulate decoded source scalars into a temporary array grouped by source type, then batch-cast in one shot and scatter results back. This would amortize kernel dispatch and let the cast kernels use their SIMD/batch paths. ########## cpp/src/parquet/variant/slot_internal.cc: ########## @@ -0,0 +1,637 @@ +// 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. + +#include "parquet/variant/slot_internal.h" + +#include <optional> +#include <string_view> +#include <type_traits> +#include <unordered_set> +#include <utility> +#include <vector> + +#include "arrow/array.h" // IWYU pragma: keep +#include "arrow/extension_type.h" +#include "arrow/type.h" +#include "arrow/util/checked_cast.h" +#include "arrow/util/decimal.h" +#include "arrow/util/logging_internal.h" +#include "parquet/exception.h" +#include "parquet/variant/append_table_internal.h" +#include "parquet/variant/array_internal.h" +#include "parquet/variant/format.h" + +namespace parquet::variant::internal { + +namespace { + +using ::arrow::Array; +using ::arrow::ExtensionArray; +using ::arrow::StructArray; +using ::arrow::internal::checked_cast; + +#define VARIANT_TARGET_DECL(...) \ + (BuildTarget * target __VA_OPT__(, ) __VA_ARGS__) // NOLINT +#define VARIANT_OBJECT_CALL_ARGS(...) \ + destination.field_name __VA_OPT__(, ) __VA_ARGS__ // NOLINT + +#define DEFINE_TARGET_APPEND(NAME, decl_args, ...) \ + void Append##NAME VARIANT_TARGET_DECL decl_args { \ + if (target == nullptr) { \ + return; \ + } \ + std::visit( \ + [&](auto& destination) { \ + using Target = std::decay_t<decltype(destination)>; \ + auto& builder = destination.builder.get(); \ + if constexpr (std::is_same_v<Target, BuildTarget::ObjectField>) { \ + builder.Append##NAME(VARIANT_OBJECT_CALL_ARGS(__VA_ARGS__)); \ + } else { \ + builder.Append##NAME(__VA_ARGS__); \ + } \ + }, \ + target->destination); \ + } +PARQUET_VARIANT_DIRECT_APPEND_LIST(DEFINE_TARGET_APPEND) +PARQUET_VARIANT_SPECIAL_APPEND_LIST(DEFINE_TARGET_APPEND) +DEFINE_TARGET_APPEND(EncodedValue, (std::string_view value), value) +#undef DEFINE_TARGET_APPEND + +#undef VARIANT_OBJECT_CALL_ARGS +#undef VARIANT_TARGET_DECL + +VariantObjectBuilder StartObject(BuildTarget& target) { + return std::visit( + [](auto& destination) -> VariantObjectBuilder { + using Target = std::decay_t<decltype(destination)>; + auto& builder = destination.builder.get(); + if constexpr (std::is_same_v<Target, BuildTarget::ObjectField>) { + return builder.StartObject(destination.field_name); + } else { + return builder.StartObject(); + } + }, + target.destination); +} + +VariantListBuilder StartList(BuildTarget& target) { + return std::visit( + [](auto& destination) -> VariantListBuilder { + using Target = std::decay_t<decltype(destination)>; + auto& builder = destination.builder.get(); + if constexpr (std::is_same_v<Target, BuildTarget::ObjectField>) { + return builder.StartList(destination.field_name); + } else { + return builder.StartList(); + } + }, + target.destination); +} + +std::optional<std::string_view> GetValueSlot(const std::shared_ptr<Array>& value_array, + int64_t row) { + if (value_array == nullptr || value_array->IsNull(row)) { + return std::nullopt; + } + return BinaryFieldView(*value_array, row); +} + +CompiledVariantRowPlan CompileVariantRowPlanImpl( + const std::shared_ptr<Array>& value_array, const std::shared_ptr<Array>& typed_array, + std::string_view path); + +CompiledVariantRowPlan::FieldGroup CompileFieldGroupPlan( + const std::shared_ptr<Array>& field_group_array, std::string_view path, + CompiledVariantRowPlan& parent) { + if (field_group_array->type_id() != ::arrow::Type::STRUCT) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + path, ": expected struct storage, got ", + field_group_array->type()->ToString()); + } + const auto& field_struct = checked_cast<const StructArray&>(*field_group_array); + const auto child_plan_index = parent.children.size(); + parent.children.push_back( + CompileVariantRowPlanImpl(field_struct.GetFieldByName("value"), + field_struct.GetFieldByName("typed_value"), path)); + return { + .array = field_group_array, + .child_plan_index = child_plan_index, + }; +} + +CompiledTypedScalarPlan CompileTypedScalarPlan( + const std::shared_ptr<Array>& typed_array) { + switch (typed_array->type_id()) { + case ::arrow::Type::BOOL: + return {.kind = TypedScalarKind::kBoolean}; + case ::arrow::Type::INT8: + return {.kind = TypedScalarKind::kInt8}; + case ::arrow::Type::INT16: + return {.kind = TypedScalarKind::kInt16}; + case ::arrow::Type::INT32: + return {.kind = TypedScalarKind::kInt32}; + case ::arrow::Type::INT64: + return {.kind = TypedScalarKind::kInt64}; + case ::arrow::Type::FLOAT: + return {.kind = TypedScalarKind::kFloat}; + case ::arrow::Type::DOUBLE: + return {.kind = TypedScalarKind::kDouble}; + case ::arrow::Type::BINARY: + case ::arrow::Type::LARGE_BINARY: + case ::arrow::Type::BINARY_VIEW: + return {.kind = TypedScalarKind::kBinary}; + case ::arrow::Type::STRING: + case ::arrow::Type::LARGE_STRING: + case ::arrow::Type::STRING_VIEW: + return {.kind = TypedScalarKind::kString}; + case ::arrow::Type::DATE32: + return {.kind = TypedScalarKind::kDate}; + case ::arrow::Type::TIME64: { + const auto& type = checked_cast<const ::arrow::Time64Type&>(*typed_array->type()); + if (type.unit() == ::arrow::TimeUnit::MICRO) { + return {.kind = TypedScalarKind::kTimeNTZMicros}; + } + break; + } + case ::arrow::Type::TIMESTAMP: { + const auto& type = + checked_cast<const ::arrow::TimestampType&>(*typed_array->type()); + switch (type.unit()) { + case ::arrow::TimeUnit::MICRO: + return {.kind = TypedScalarKind::kTimestampMicros, + .adjusted_to_utc = !type.timezone().empty()}; + case ::arrow::TimeUnit::NANO: + return {.kind = TypedScalarKind::kTimestampNanos, + .adjusted_to_utc = !type.timezone().empty()}; + case ::arrow::TimeUnit::SECOND: + case ::arrow::TimeUnit::MILLI: + break; + } + break; + } + case ::arrow::Type::DECIMAL32: { + const auto& type = + checked_cast<const ::arrow::Decimal32Type&>(*typed_array->type()); + return {.kind = TypedScalarKind::kDecimal4, + .scale = static_cast<uint8_t>(type.scale())}; + } + case ::arrow::Type::DECIMAL64: { + const auto& type = + checked_cast<const ::arrow::Decimal64Type&>(*typed_array->type()); + return {.kind = TypedScalarKind::kDecimal8, + .scale = static_cast<uint8_t>(type.scale())}; + } + case ::arrow::Type::DECIMAL128: { + const auto& type = + checked_cast<const ::arrow::Decimal128Type&>(*typed_array->type()); + return {.kind = TypedScalarKind::kDecimal16, + .scale = static_cast<uint8_t>(type.scale())}; + } + case ::arrow::Type::EXTENSION: { + const auto& ext_type = + checked_cast<const ::arrow::ExtensionType&>(*typed_array->type()); + if (ext_type.extension_name() == "arrow.uuid") { + return {.kind = TypedScalarKind::kUuidExtension}; + } + break; + } + default: + break; + } + throw ParquetInvalidOrCorruptedFileException("Illegal shredded value type: ", + typed_array->type()->ToString()); +} + +CompiledVariantRowPlan CompileVariantRowPlanImpl( + const std::shared_ptr<Array>& value_array, const std::shared_ptr<Array>& typed_array, + std::string_view path) { + CompiledVariantRowPlan plan{ + .value_array = value_array, + .typed = std::nullopt, + .children = {}, + }; + if (typed_array == nullptr) { + return plan; + } + + switch (typed_array->type_id()) { + case ::arrow::Type::STRUCT: { + const auto& typed_struct = checked_cast<const StructArray&>(*typed_array); + CompiledVariantRowPlan::Object object; + object.fields.reserve(typed_struct.struct_type()->num_fields()); + object.field_names.reserve(typed_struct.struct_type()->num_fields()); + for (int i = 0; i < typed_struct.struct_type()->num_fields(); ++i) { + const auto& field_name = typed_struct.struct_type()->field(i)->name(); + if (!object.field_names.insert(field_name).second) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: duplicate shredded object field '", field_name, + "'"); + } + auto field_group = CompileFieldGroupPlan(typed_struct.field(i), field_name, plan); + object.fields.push_back( + {.field_name = field_name, .field_group = std::move(field_group)}); + } + plan.typed.emplace(CompiledVariantRowPlan::Typed{ + .array = typed_array, + .plan = std::move(object), + }); + return plan; + } + case ::arrow::Type::LIST: + case ::arrow::Type::LARGE_LIST: + case ::arrow::Type::LIST_VIEW: + case ::arrow::Type::LARGE_LIST_VIEW: + case ::arrow::Type::FIXED_SIZE_LIST: { + auto values = ValuesArray(*typed_array); + auto element = CompileFieldGroupPlan(values, path, plan); + plan.typed.emplace(CompiledVariantRowPlan::Typed{ + .array = typed_array, + .plan = CompiledVariantRowPlan::Array{.element = std::move(element)}, + }); + return plan; + } + default: + plan.typed.emplace(CompiledVariantRowPlan::Typed{ + .array = typed_array, + .plan = + CompiledVariantRowPlan::Primitive{ + .scalar = CompileTypedScalarPlan(typed_array), + }, + }); + return plan; + } +} + +void ValidateTypedFieldNames(const VariantMetadataView& metadata, + const CompiledVariantRowPlan::Object& object) { + for (const auto& field : object.fields) { + if (!metadata.FindString(field.field_name).has_value()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: shredded field '", field.field_name, + "' is not in metadata dictionary"); + } + } +} + +void AppendTypedScalar(BuildTarget* target, const CompiledTypedScalarPlan& plan, + const std::shared_ptr<Array>& typed_array, int64_t row) { + if (target == nullptr) { + return; + } + switch (plan.kind) { + case TypedScalarKind::kBoolean: + AppendBoolean(target, + checked_cast<const ::arrow::BooleanArray&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt8: + AppendInt8(target, + checked_cast<const ::arrow::Int8Array&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt16: + AppendInt16(target, + checked_cast<const ::arrow::Int16Array&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt32: + AppendInt32(target, + checked_cast<const ::arrow::Int32Array&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kInt64: + AppendInt64(target, + checked_cast<const ::arrow::Int64Array&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kFloat: + AppendFloat(target, + checked_cast<const ::arrow::FloatArray&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kDouble: + AppendDouble(target, + checked_cast<const ::arrow::DoubleArray&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kBinary: + AppendBinary(target, BinaryFieldView(*typed_array, row)); + return; + case TypedScalarKind::kString: { + auto value = StringFieldView(*typed_array, row); + if (value.size() <= kMaxShortStringSize) { + AppendShortString(target, value); + } else { + AppendString(target, value); + } + return; + } + case TypedScalarKind::kDate: + AppendDate(target, + checked_cast<const ::arrow::Date32Array&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kTimeNTZMicros: + AppendTimeNTZMicros( + target, checked_cast<const ::arrow::Time64Array&>(*typed_array).Value(row)); + return; + case TypedScalarKind::kTimestampMicros: + AppendTimestampMicros( + target, checked_cast<const ::arrow::TimestampArray&>(*typed_array).Value(row), + plan.adjusted_to_utc); + return; + case TypedScalarKind::kTimestampNanos: + AppendTimestampNanos( + target, checked_cast<const ::arrow::TimestampArray&>(*typed_array).Value(row), + plan.adjusted_to_utc); + return; + case TypedScalarKind::kDecimal4: { + const auto value = ::arrow::Decimal32( + checked_cast<const ::arrow::Decimal32Array&>(*typed_array).GetValue(row)); + AppendDecimal4(target, value, plan.scale); + return; + } + case TypedScalarKind::kDecimal8: { + const auto value = ::arrow::Decimal64( + checked_cast<const ::arrow::Decimal64Array&>(*typed_array).GetValue(row)); + AppendDecimal8(target, value, plan.scale); + return; + } + case TypedScalarKind::kDecimal16: { + const auto value = ::arrow::Decimal128( + checked_cast<const ::arrow::Decimal128Array&>(*typed_array).GetValue(row)); + AppendDecimal16(target, value, plan.scale); + return; + } + case TypedScalarKind::kUuidExtension: { + const auto& ext_array = checked_cast<const ExtensionArray&>(*typed_array); + auto storage = ext_array.storage(); + auto value = + checked_cast<const ::arrow::FixedSizeBinaryArray&>(*storage).GetView(row); + AppendUuid(target, value); + return; + } + } +} + +template <bool strict> +void ProcessListElement(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::FieldGroup& field_group_plan, + int64_t row, std::string_view path, BuildTarget* target) { + const auto& field_array = *field_group_plan.array; + if (field_array.IsNull(row)) { + if constexpr (strict) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant field at ", path, ": field group must be required"); + } else { + AppendVariantNull(target); + return; + } + } + + DCHECK_LT(field_group_plan.child_plan_index, parent_plan.children.size()); + const auto& row_plan = parent_plan.children[field_group_plan.child_plan_index]; + ProcessSlot<strict>(metadata, row_plan, row, target, path); +} + +template <bool strict> +void ProcessTypedArraySlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::Typed& typed, + const CompiledVariantRowPlan::Array& array, int64_t row, + std::string_view path, BuildTarget* target) { + DCHECK(!typed.array->IsNull(row)); + + std::optional<VariantListBuilder> list_builder; + std::optional<BuildTarget> child_target; + BuildTarget* child_target_ptr = nullptr; + if (target != nullptr) { + list_builder.emplace(StartList(*target)); + child_target.emplace(BuildTarget{BuildTarget::ListElement{*list_builder}}); + child_target_ptr = &*child_target; + } + const auto [offset, length] = ValuesRangeAt(*typed.array, row); + for (int64_t i = 0; i < length; ++i) { + ProcessListElement<strict>(metadata, parent_plan, array.element, offset + i, path, + child_target_ptr); + } + if (list_builder.has_value()) { + list_builder->Finish(); + } +} + +template <bool strict> +bool TryProcessSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& plan, int64_t row, + std::string_view path, BuildTarget* target); + +template <bool strict> +void ProcessObjectField(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::ObjectField& field, int64_t row, + BuildTarget* target) { + const auto& field_array = *field.field_group.array; + if (field_array.IsNull(row)) { + if constexpr (strict) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + field.field_name, + ": field group must be required"); + } else { + return; + } + } + + DCHECK_LT(field.field_group.child_plan_index, parent_plan.children.size()); + const auto& row_plan = parent_plan.children[field.field_group.child_plan_index]; + if constexpr (strict) { + if (row_plan.value_array == nullptr) { + throw ParquetInvalidOrCorruptedFileException("Invalid shredded Variant field at ", + field.field_name, + ": missing value field"); + } + } + + if (target != nullptr) { + std::get<BuildTarget::ObjectField>(target->destination).field_name = field.field_name; + } + TryProcessSlot<strict>(metadata, row_plan, row, field.field_name, target); +} + +template <bool strict> +void ProcessObjectResidual(const VariantValueView& value_view, + const CompiledVariantRowPlan::Object& object, + BuildTarget* target) { + DCHECK_EQ(value_view.basic_type(), VariantBasicType::kObject); + const auto& object_view = std::get<VariantObjectView>(value_view.data()); + for (const auto& field : object_view.fields()) { + if (object.field_names.contains(field.name)) { + if (strict || target != nullptr) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid shredded Variant: value object contains shredded field: ", + field.name); + } + continue; + } + if (target == nullptr) { + continue; + } + BuildTarget residual_target = *target; + std::get<BuildTarget::ObjectField>(residual_target.destination).field_name = + field.name; + AppendEncodedValue(&residual_target, field.value); + } +} + +template <bool strict> +void ProcessTypedObjectSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& parent_plan, + const CompiledVariantRowPlan::Typed& typed, + const CompiledVariantRowPlan::Object& object, int64_t row, + std::optional<std::string_view> value, BuildTarget* target) { + DCHECK(!typed.array->IsNull(row)); + + std::optional<VariantValueView> object_value_view; + if (value.has_value()) { + if (VariantValueView::PeekBasicType(*value) != VariantBasicType::kObject) { + throw ParquetInvalidOrCorruptedFileException( + "Expected object in value field for partially shredded struct"); + } + if (target != nullptr) { + object_value_view = VariantValueView::Make(*value, metadata); + } else if constexpr (strict) { + object_value_view = VariantValueView::MakeWithValidate(*value, metadata); + } else { + VariantValueView::Validate(*value, metadata); + } + } + + std::optional<VariantObjectBuilder> object_builder; + std::optional<BuildTarget> child_target; + BuildTarget* child_target_ptr = nullptr; + if (target != nullptr) { + object_builder.emplace(StartObject(*target)); + child_target.emplace(BuildTarget{BuildTarget::ObjectField{ + .builder = *object_builder, + .field_name = {}, + }}); + child_target_ptr = &*child_target; + } + for (const auto& field : object.fields) { + ProcessObjectField<strict>(metadata, parent_plan, field, row, child_target_ptr); + } + if (object_value_view.has_value()) { + ProcessObjectResidual<strict>(*object_value_view, object, child_target_ptr); + } + if (object_builder.has_value()) { + object_builder->Finish(); + } +} + +template <bool strict> +bool TryProcessSlot(const VariantMetadataView& metadata, + const CompiledVariantRowPlan& plan, int64_t row, + std::string_view path, BuildTarget* target) { + const auto value = GetValueSlot(plan.value_array, row); + if (!plan.typed.has_value()) { + if (!value.has_value()) { + return false; + } + if (target == nullptr) { + VariantValueView::Validate(*value, metadata); + } + AppendEncodedValue(target, *value); + return true; + } + + const auto& typed = *plan.typed; + return std::visit( + [&](const auto& typed_plan) -> bool { + using Plan = std::decay_t<decltype(typed_plan)>; + const bool typed_present = !typed.array->IsNull(row); + + if constexpr (std::is_same_v<Plan, CompiledVariantRowPlan::Object>) { + ValidateTypedFieldNames(metadata, typed_plan); + } Review Comment: `ValidateTypedFieldNames` calls `metadata.FindString()` for each shredded field on every row during unshredding. Since most rows share identical metadata bytes, caching the validation result per unique metadata blob (similar to `ObjectShredNode::ValidateMetadata` in `shred.cc`) would skip redundant lookups here too. ########## cpp/src/parquet/variant/decoding.cc: ########## @@ -0,0 +1,526 @@ +// 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. + +#include "parquet/variant/decoding.h" + +#include <algorithm> +#include <cstring> +#include <limits> +#include <utility> + +#include "arrow/util/endian.h" +#include "arrow/util/logging_internal.h" +#include "parquet/exception.h" +#include "parquet/variant/format_internal.h" + +namespace parquet::variant { + +namespace bit_util = ::arrow::bit_util; + +namespace internal { + +class ViewAccess { + public: + template <typename View, typename... Args> + static View Make(Args&&... args) { + return View(std::forward<Args>(args)...); + } +}; + +} // namespace internal + +namespace { + +using internal::ViewAccess; + +uint32_t ReadLittleEndian(std::string_view data, size_t offset, size_t width) { + DCHECK_LE(width, sizeof(uint32_t)); + uint32_t value = 0; + std::memcpy(&value, data.data() + offset, width); + return bit_util::FromLittleEndian(value); +} + +void CheckAvailable(std::string_view data, size_t offset, size_t size, + std::string_view context) { + if (offset > data.size() || data.size() - offset < size) { + throw ParquetInvalidOrCorruptedFileException("Invalid Variant encoding: truncated ", + context); + } +} + +size_t OffsetCount(uint32_t count) { + const auto converted_count = static_cast<size_t>(count); + if (converted_count == std::numeric_limits<size_t>::max()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: offset count overflow"); + } + return converted_count + 1; +} + +size_t OffsetBytes(size_t offset_count, uint8_t offset_size) { + if (offset_count > + std::numeric_limits<size_t>::max() / static_cast<size_t>(offset_size)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: offset table size overflow"); + } + return offset_count * static_cast<size_t>(offset_size); +} + +size_t PrimitivePayloadSize(std::string_view value, size_t offset, + VariantPrimitiveType primitive) { + switch (primitive) { + case VariantPrimitiveType::kNull: + case VariantPrimitiveType::kBooleanTrue: + case VariantPrimitiveType::kBooleanFalse: + return 0; + case VariantPrimitiveType::kInt8: + return 1; + case VariantPrimitiveType::kInt16: + return 2; + case VariantPrimitiveType::kInt32: + case VariantPrimitiveType::kDate: + case VariantPrimitiveType::kFloat: + return 4; + case VariantPrimitiveType::kInt64: + case VariantPrimitiveType::kDouble: + case VariantPrimitiveType::kTimestampMicros: + case VariantPrimitiveType::kTimestampNTZMicros: + case VariantPrimitiveType::kTimeNTZMicros: + case VariantPrimitiveType::kTimestampNanos: + case VariantPrimitiveType::kTimestampNTZNanos: + return 8; + case VariantPrimitiveType::kDecimal4: + return 5; + case VariantPrimitiveType::kDecimal8: + return 9; + case VariantPrimitiveType::kDecimal16: + return 17; + case VariantPrimitiveType::kUuid: + return 16; + case VariantPrimitiveType::kBinary: + case VariantPrimitiveType::kString: { + CheckAvailable(value, offset, 4, "variable-length size"); + const uint32_t length = ReadLittleEndian(value, offset, 4); + return 4 + static_cast<size_t>(length); + } + } + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown primitive type"); +} + +size_t ParsePrimitive(std::string_view value, size_t offset, + VariantPrimitiveType primitive) { + if (!internal::IsKnownVariantPrimitive(primitive)) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: unknown primitive type ", static_cast<int>(primitive)); + } + + const size_t payload_size = PrimitivePayloadSize(value, offset, primitive); + CheckAvailable(value, offset, payload_size, "primitive value"); + + if (internal::IsDecimalVariantPrimitive(primitive)) { + const auto scale = static_cast<uint8_t>(value[offset]); + internal::ValidateDecimalScale(scale); + } + + if (primitive == VariantPrimitiveType::kString) { + const uint32_t length = ReadLittleEndian(value, offset, 4); + internal::ValidateUtf8(value.substr(offset + 4, length), "primitive string value"); + } + + return payload_size; +} + +template <bool validate_children> +size_t ParseValue(std::string_view value, const VariantMetadataView& metadata, + VariantValueView* out); + +template <bool validate_children> +size_t ParseArray(std::string_view value, const VariantMetadataView& metadata, + uint8_t header, VariantValueView* out) { + const auto offset_size = static_cast<uint8_t>((header & 0x03) + 1); + const bool is_large = (header & 0x04) != 0; + const size_t count_size = is_large ? 4 : 1; + + size_t offset = 1; + CheckAvailable(value, offset, count_size, "array size"); + const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); + offset += count_size; + const size_t offset_count = OffsetCount(num_elements); + + CheckAvailable(value, offset, OffsetBytes(offset_count, offset_size), "array offsets"); + + std::vector<uint32_t> offsets(offset_count); + for (size_t i = 0; i < offset_count; ++i) { + offsets[i] = ReadLittleEndian(value, offset, offset_size); + offset += offset_size; + } + + if (offsets[0] != 0) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: first array offset must be 0"); + } + for (uint32_t i = 0; i < num_elements; ++i) { + if (offsets[i] > offsets[i + 1]) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: array offsets must be monotonic"); + } + } + + const size_t values_start = offset; + const size_t total_value_size = offsets[num_elements]; + CheckAvailable(value, values_start, total_value_size, "array values"); + + if constexpr (validate_children) { + for (uint32_t i = 0; i < num_elements; ++i) { + const size_t child_consumed = ParseValue<true>( + value.substr(values_start + offsets[i]), metadata, /*out=*/nullptr); + if (child_consumed != offsets[i + 1] - offsets[i]) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: array value does not end at next offset"); + } + } + } + + const size_t consumed = values_start + total_value_size; + if (out != nullptr) { + std::vector<std::string_view> array_elements; + array_elements.reserve(num_elements); + for (uint32_t i = 0; i < num_elements; ++i) { + array_elements.push_back( + value.substr(values_start + offsets[i], offsets[i + 1] - offsets[i])); + } + *out = ViewAccess::Make<VariantValueView>( + value.substr(0, consumed), + ViewAccess::Make<VariantArrayView>(std::move(array_elements), metadata)); + } + return consumed; +} + +template <bool validate_children> +size_t ParseObject(std::string_view value, const VariantMetadataView& metadata, + uint8_t header, VariantValueView* out) { + const auto offset_size = static_cast<uint8_t>((header & 0x03) + 1); + const auto id_size = static_cast<uint8_t>(((header >> 2) & 0x03) + 1); + const bool is_large = (header & 0x10) != 0; + const size_t count_size = is_large ? 4 : 1; + + size_t offset = 1; + CheckAvailable(value, offset, count_size, "object size"); + const uint32_t num_elements = ReadLittleEndian(value, offset, count_size); + offset += count_size; + + CheckAvailable(value, offset, static_cast<size_t>(num_elements) * id_size, + "object field ids"); + std::vector<uint32_t> field_ids(num_elements); + for (uint32_t i = 0; i < num_elements; ++i) { + field_ids[i] = ReadLittleEndian(value, offset, id_size); + offset += id_size; + } + + const size_t offset_count = OffsetCount(num_elements); + CheckAvailable(value, offset, OffsetBytes(offset_count, offset_size), + "object field offsets"); + std::vector<uint32_t> field_offsets(offset_count); + for (size_t i = 0; i < offset_count; ++i) { + field_offsets[i] = ReadLittleEndian(value, offset, offset_size); + offset += offset_size; + } + + const size_t values_start = offset; + const size_t total_value_size = field_offsets[num_elements]; + if (num_elements == 0 && total_value_size != 0) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: empty object must have zero value size"); + } + CheckAvailable(value, values_start, total_value_size, "object values"); + + std::vector<VariantObjectField> object_fields; + if (out != nullptr) { + object_fields.reserve(num_elements); + } + + for (uint32_t i = 0; i < num_elements; ++i) { + if (field_ids[i] >= metadata.dictionary_size()) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object field id ", field_ids[i], + " is outside metadata dictionary of size ", metadata.dictionary_size()); + } + if constexpr (validate_children) { + if (i > 0 && !(metadata.string(field_ids[i - 1]) < metadata.string(field_ids[i]))) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object field names must be sorted and unique"); + } + } + + const auto field_offset = field_offsets[i]; + if (field_offset >= total_value_size) { + throw ParquetInvalidOrCorruptedFileException( + "Invalid Variant encoding: object field offset is outside values"); + } + } + + std::vector<uint32_t> value_offsets = field_offsets; + std::ranges::sort(value_offsets); + if (std::ranges::adjacent_find(value_offsets) != value_offsets.end()) { Review Comment: This copies all field offsets into a new vector, sorts, and checks for duplicates -- per object, per row. For typical small objects (<10 fields) it's fine, but at scale a reusable scratch buffer (or a linear scan for small n, switching to sort only above a threshold) would reduce per-row allocation pressure. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
