HuaHuaY commented on code in PR #50252: URL: https://github.com/apache/arrow/pull/50252#discussion_r4357738331
########## 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 is an improvement, but given that its implementation is rather complex, I think it’s better to leave it for a future PR. -- 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]
