| /* |
| * 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_reassembler.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/generic_variant.h" |
| #include "paimon/common/data/variant/variant_builder.h" |
| #include "paimon/common/data/variant/variant_type_utils.h" |
| #include "paimon/common/utils/arrow/status_utils.h" |
| #include "paimon/memory/memory_pool.h" |
| |
| namespace paimon { |
| |
| namespace { |
| |
| // A row view over a shredded struct array, the Arrow analog of the Java `ShreddedRow`. |
| struct Cursor { |
| const arrow::StructArray* array; |
| int64_t row; |
| |
| bool IsNullAt(int32_t field) const { |
| return array->field(field)->IsNull(row); |
| } |
| |
| std::string_view GetBinary(int32_t field) const { |
| return static_cast<const arrow::BinaryArray&>(*array->field(field)).GetView(row); |
| } |
| }; |
| |
| Status Rebuild(const Cursor& cursor, std::string_view metadata, const VariantSchema& schema, |
| const std::shared_ptr<MemoryPool>& pool, VariantBuilder* builder); |
| |
| Status RebuildTypedScalar(const Cursor& cursor, const VariantSchema& schema, |
| VariantBuilder* builder) { |
| int32_t typed_idx = schema.typed_idx; |
| const arrow::Array& typed_array = *cursor.array->field(typed_idx); |
| int64_t row = cursor.row; |
| const VariantSchema::ScalarType& scalar = schema.scalar_schema.value(); |
| switch (scalar.kind) { |
| case VariantSchema::ScalarKind::kString: |
| return builder->AppendString( |
| static_cast<const arrow::StringArray&>(typed_array).GetView(row)); |
| case VariantSchema::ScalarKind::kByte: |
| return builder->AppendLong( |
| static_cast<const arrow::Int8Array&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kShort: |
| return builder->AppendLong( |
| static_cast<const arrow::Int16Array&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kInt: |
| return builder->AppendLong( |
| static_cast<const arrow::Int32Array&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kLong: |
| return builder->AppendLong( |
| static_cast<const arrow::Int64Array&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kFloat: |
| return builder->AppendFloat( |
| static_cast<const arrow::FloatArray&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kDouble: |
| return builder->AppendDouble( |
| static_cast<const arrow::DoubleArray&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kBoolean: |
| return builder->AppendBoolean( |
| static_cast<const arrow::BooleanArray&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kBinary: |
| return builder->AppendBinary( |
| static_cast<const arrow::BinaryArray&>(typed_array).GetView(row)); |
| case VariantSchema::ScalarKind::kDecimal: { |
| const auto& decimal_array = static_cast<const arrow::Decimal128Array&>(typed_array); |
| arrow::Decimal128 value(decimal_array.GetValue(row)); |
| VariantDecimal decimal; |
| decimal.unscaled = (static_cast<__int128_t>(value.high_bits()) << 64) | |
| static_cast<__int128_t>(static_cast<__uint128_t>(value.low_bits())); |
| decimal.scale = scalar.scale; |
| return builder->AppendDecimal(decimal); |
| } |
| case VariantSchema::ScalarKind::kDate: |
| return builder->AppendDate( |
| static_cast<const arrow::Date32Array&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kTimestampLtz: |
| return builder->AppendTimestamp( |
| static_cast<const arrow::TimestampArray&>(typed_array).Value(row)); |
| case VariantSchema::ScalarKind::kTimestampNtz: |
| return builder->AppendTimestampNtz( |
| static_cast<const arrow::TimestampArray&>(typed_array).Value(row)); |
| default: |
| return Status::NotImplemented("unsupported variant scalar kind in reassembly"); |
| } |
| } |
| |
| // Rebuilds a variant value from the shredded data according to the reconstruction algorithm in |
| // the parquet-format VariantShredding.md specification, appending the result to `builder`. |
| Status Rebuild(const Cursor& cursor, std::string_view metadata, const VariantSchema& schema, |
| const std::shared_ptr<MemoryPool>& pool, VariantBuilder* builder) { |
| int32_t typed_idx = schema.typed_idx; |
| int32_t variant_idx = schema.variant_idx; |
| if (typed_idx >= 0 && !cursor.IsNullAt(typed_idx)) { |
| if (schema.scalar_schema.has_value()) { |
| return RebuildTypedScalar(cursor, schema, builder); |
| } else if (schema.array_schema != nullptr) { |
| const auto& list_array = |
| static_cast<const arrow::ListArray&>(*cursor.array->field(typed_idx)); |
| const auto& element_array = |
| static_cast<const arrow::StructArray&>(*list_array.values()); |
| int64_t element_start = list_array.value_offset(cursor.row); |
| int64_t element_end = list_array.value_offset(cursor.row + 1); |
| int32_t start = builder->GetWritePos(); |
| std::vector<int32_t> offsets; |
| offsets.reserve(element_end - element_start); |
| for (int64_t i = element_start; i < element_end; ++i) { |
| offsets.push_back(builder->GetWritePos() - start); |
| PAIMON_RETURN_NOT_OK(Rebuild(Cursor{&element_array, i}, metadata, |
| *schema.array_schema, pool, builder)); |
| } |
| return builder->FinishWritingArray(start, offsets); |
| } else { |
| const auto& object_array = |
| static_cast<const arrow::StructArray&>(*cursor.array->field(typed_idx)); |
| Cursor object_cursor{&object_array, cursor.row}; |
| std::vector<VariantBuilder::FieldEntry> fields; |
| int32_t start = builder->GetWritePos(); |
| for (size_t field_idx = 0; field_idx < schema.object_schema.size(); ++field_idx) { |
| // Shredded fields must not be null. |
| if (object_cursor.IsNullAt(static_cast<int32_t>(field_idx))) { |
| return VariantBinaryUtil::MalformedVariant( |
| "a shredded object field group is null"); |
| } |
| const std::string& field_name = schema.object_schema[field_idx].name; |
| const VariantSchema& field_schema = *schema.object_schema[field_idx].schema; |
| const auto& field_array = static_cast<const arrow::StructArray&>( |
| *object_array.field(static_cast<int32_t>(field_idx))); |
| Cursor field_cursor{&field_array, cursor.row}; |
| // If the field doesn't have a non-null `typed_value` or `value`, it is missing. |
| if ((field_schema.typed_idx >= 0 && |
| !field_cursor.IsNullAt(field_schema.typed_idx)) || |
| (field_schema.variant_idx >= 0 && |
| !field_cursor.IsNullAt(field_schema.variant_idx))) { |
| int32_t id = builder->AddKey(field_name); |
| fields.emplace_back(field_name, id, builder->GetWritePos() - start); |
| PAIMON_RETURN_NOT_OK( |
| Rebuild(field_cursor, metadata, field_schema, pool, builder)); |
| } |
| } |
| if (variant_idx >= 0 && !cursor.IsNullAt(variant_idx)) { |
| // Add the leftover fields in the variant binary. |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<GenericVariant> leftover, |
| GenericVariant::Create(cursor.GetBinary(variant_idx), metadata, pool)); |
| PAIMON_ASSIGN_OR_RAISE(VariantValueType leftover_type, leftover->GetType()); |
| if (leftover_type != VariantValueType::kObject) { |
| return VariantBinaryUtil::MalformedVariant( |
| "the value column of a shredded object is not an object"); |
| } |
| PAIMON_ASSIGN_OR_RAISE(int32_t leftover_size, leftover->ObjectSize()); |
| for (int32_t i = 0; i < leftover_size; ++i) { |
| PAIMON_ASSIGN_OR_RAISE(std::optional<GenericVariant::ObjectField> field, |
| leftover->GetFieldAtIndex(i)); |
| if (!field.has_value()) { |
| return VariantBinaryUtil::MalformedVariant( |
| "a leftover object field is missing"); |
| } |
| // `value` must not contain any shredded field. |
| if (schema.object_schema_map.count(field->key) > 0) { |
| return VariantBinaryUtil::MalformedVariant(fmt::format( |
| "the value column duplicates the shredded field '{}'", field->key)); |
| } |
| int32_t id = builder->AddKey(field->key); |
| fields.emplace_back(field->key, id, builder->GetWritePos() - start); |
| PAIMON_RETURN_NOT_OK(builder->AppendVariant(*field->value)); |
| } |
| } |
| return builder->FinishWritingObject(start, &fields); |
| } |
| } else if (variant_idx >= 0 && !cursor.IsNullAt(variant_idx)) { |
| // `typed_value` doesn't exist or is null. Read from `value`. |
| PAIMON_ASSIGN_OR_RAISE( |
| std::shared_ptr<GenericVariant> variant, |
| GenericVariant::Create(cursor.GetBinary(variant_idx), metadata, pool)); |
| return builder->AppendVariant(*variant); |
| } else { |
| // The variant is missing in a context where it must be present; the data is invalid. |
| return VariantBinaryUtil::MalformedVariant( |
| "both typed_value and value of a required variant are null"); |
| } |
| } |
| |
| } // namespace |
| |
| Status VariantReassembler::RebuildValue(const arrow::StructArray& shredded, int64_t row, |
| std::string_view metadata, const VariantSchema& schema, |
| const std::shared_ptr<MemoryPool>& pool, |
| VariantBuilder* builder) { |
| return Rebuild(Cursor{&shredded, row}, metadata, schema, pool, builder); |
| } |
| |
| Result<std::shared_ptr<arrow::Array>> VariantReassembler::AssembleVariantArray( |
| const std::shared_ptr<arrow::StructArray>& shredded, |
| const std::shared_ptr<VariantSchema>& schema, const std::shared_ptr<MemoryPool>& pool, |
| arrow::MemoryPool* arrow_pool) { |
| if (schema->top_level_metadata_idx < 0) { |
| return VariantBinaryUtil::MalformedVariant("a shredded file column misses metadata"); |
| } |
| auto output_type = VariantTypeUtils::UnshreddedStructType(); |
| std::unique_ptr<arrow::ArrayBuilder> output_builder; |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::MakeBuilder(arrow_pool, output_type, &output_builder)); |
| auto* struct_builder = static_cast<arrow::StructBuilder*>(output_builder.get()); |
| auto* value_builder = static_cast<arrow::BinaryBuilder*>(struct_builder->field_builder(0)); |
| auto* metadata_builder = static_cast<arrow::BinaryBuilder*>(struct_builder->field_builder(1)); |
| |
| bool unshredded = schema->IsUnshredded(); |
| for (int64_t row = 0; row < shredded->length(); ++row) { |
| if (shredded->IsNull(row)) { |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->AppendNull()); |
| continue; |
| } |
| Cursor cursor{shredded.get(), row}; |
| if (cursor.IsNullAt(schema->top_level_metadata_idx)) { |
| return VariantBinaryUtil::MalformedVariant("the variant metadata column is null"); |
| } |
| std::string_view metadata = cursor.GetBinary(schema->top_level_metadata_idx); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(struct_builder->Append()); |
| if (unshredded) { |
| // Rebuilding is unnecessary for unshredded variants. |
| // TODO(nicholas): avoid copying the value/metadata binaries through the builder for |
| // unshredded files; the physical arrays could be returned directly with at most a |
| // per-row malformed-variant check. |
| if (cursor.IsNullAt(schema->variant_idx)) { |
| return VariantBinaryUtil::MalformedVariant( |
| "the value column of an unshredded variant is null"); |
| } |
| PAIMON_RETURN_NOT_OK_FROM_ARROW( |
| value_builder->Append(cursor.GetBinary(schema->variant_idx))); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata_builder->Append(metadata)); |
| } else { |
| VariantBuilder builder(/*allow_duplicate_keys=*/false); |
| PAIMON_RETURN_NOT_OK(Rebuild(cursor, metadata, *schema, pool, &builder)); |
| PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<GenericVariant> variant, builder.Build(pool)); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(value_builder->Append(variant->RawValue())); |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(metadata_builder->Append(variant->Metadata())); |
| } |
| } |
| std::shared_ptr<arrow::Array> result; |
| PAIMON_RETURN_NOT_OK_FROM_ARROW(output_builder->Finish(&result)); |
| return result; |
| } |
| |
| } // namespace paimon |
