Skip to content

Commit cd4ca42

Browse files
huan233uschuan233usc
authored
feat(parquet): apply column default values when reading missing fields (#792)
1 parent a080eb0 commit cd4ca42

9 files changed

Lines changed: 564 additions & 1 deletion

‎src/iceberg/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,7 @@ if(ICEBERG_BUILD_BUNDLE)
263263
arrow/arrow_io.cc
264264
arrow/s3/arrow_s3_file_io.cc
265265
arrow/arrow_register.cc
266+
arrow/literal_util.cc
266267
arrow/metadata_column_util.cc
267268
avro/avro_data_util.cc
268269
avro/avro_direct_decoder.cc

‎src/iceberg/arrow/literal_util.cc‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#include <cstring>
21+
#include <string>
22+
#include <utility>
23+
#include <vector>
24+
25+
#include <arrow/array/util.h>
26+
#include <arrow/buffer.h>
27+
#include <arrow/compute/api.h>
28+
#include <arrow/extension_type.h>
29+
#include <arrow/scalar.h>
30+
#include <arrow/type.h>
31+
32+
#include "iceberg/arrow/arrow_status_internal.h"
33+
#include "iceberg/arrow/literal_util_internal.h"
34+
#include "iceberg/type.h"
35+
#include "iceberg/util/checked_cast.h"
36+
#include "iceberg/util/formatter.h" // IWYU pragma: keep
37+
#include "iceberg/util/macros.h"
38+
39+
namespace iceberg::arrow {
40+
41+
namespace {
42+
43+
Result<std::shared_ptr<::arrow::DataType>> ToArrowType(const PrimitiveType& type) {
44+
switch (type.type_id()) {
45+
case TypeId::kBoolean:
46+
return ::arrow::boolean();
47+
case TypeId::kInt:
48+
return ::arrow::int32();
49+
case TypeId::kLong:
50+
return ::arrow::int64();
51+
case TypeId::kFloat:
52+
return ::arrow::float32();
53+
case TypeId::kDouble:
54+
return ::arrow::float64();
55+
case TypeId::kDecimal: {
56+
const auto& decimal_type = internal::checked_cast<const DecimalType&>(type);
57+
return ::arrow::decimal128(decimal_type.precision(), decimal_type.scale());
58+
}
59+
case TypeId::kDate:
60+
return ::arrow::date32();
61+
case TypeId::kTime:
62+
return ::arrow::time64(::arrow::TimeUnit::MICRO);
63+
case TypeId::kTimestamp:
64+
return ::arrow::timestamp(::arrow::TimeUnit::MICRO);
65+
case TypeId::kTimestampTz:
66+
return ::arrow::timestamp(::arrow::TimeUnit::MICRO, "UTC");
67+
case TypeId::kTimestampNs:
68+
return ::arrow::timestamp(::arrow::TimeUnit::NANO);
69+
case TypeId::kTimestampTzNs:
70+
return ::arrow::timestamp(::arrow::TimeUnit::NANO, "UTC");
71+
case TypeId::kString:
72+
return ::arrow::utf8();
73+
case TypeId::kBinary:
74+
return ::arrow::binary();
75+
case TypeId::kFixed: {
76+
const auto& fixed_type = internal::checked_cast<const FixedType&>(type);
77+
return ::arrow::fixed_size_binary(static_cast<int32_t>(fixed_type.length()));
78+
}
79+
case TypeId::kUuid:
80+
return ::arrow::fixed_size_binary(16);
81+
default:
82+
return NotSupported("Cannot convert {} to an Arrow type", type);
83+
}
84+
}
85+
86+
Result<std::shared_ptr<::arrow::Buffer>> ToArrowBuffer(const std::vector<uint8_t>& bytes,
87+
::arrow::MemoryPool* pool) {
88+
ICEBERG_ARROW_ASSIGN_OR_RETURN(std::unique_ptr<::arrow::Buffer> buffer,
89+
::arrow::AllocateBuffer(bytes.size(), pool));
90+
std::memcpy(buffer->mutable_data(), bytes.data(), bytes.size());
91+
return std::shared_ptr<::arrow::Buffer>(std::move(buffer));
92+
}
93+
94+
} // namespace
95+
96+
Result<std::shared_ptr<::arrow::Scalar>> ToArrowScalar(const Literal& literal,
97+
::arrow::MemoryPool* pool) {
98+
if (literal.type() == nullptr) {
99+
return InvalidArgument("Cannot convert a literal without type to an Arrow scalar");
100+
}
101+
102+
if (literal.IsAboveMax() || literal.IsBelowMin()) {
103+
return NotSupported("Cannot convert {} to an Arrow scalar", literal);
104+
}
105+
106+
ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::DataType> arrow_type,
107+
ToArrowType(*literal.type()));
108+
if (literal.IsNull()) {
109+
return ::arrow::MakeNullScalar(std::move(arrow_type));
110+
}
111+
112+
const Literal::Value& value = literal.value();
113+
switch (literal.type()->type_id()) {
114+
case TypeId::kBoolean:
115+
return std::make_shared<::arrow::BooleanScalar>(std::get<bool>(value));
116+
case TypeId::kInt:
117+
return std::make_shared<::arrow::Int32Scalar>(std::get<int32_t>(value));
118+
case TypeId::kLong:
119+
return std::make_shared<::arrow::Int64Scalar>(std::get<int64_t>(value));
120+
case TypeId::kFloat:
121+
return std::make_shared<::arrow::FloatScalar>(std::get<float>(value));
122+
case TypeId::kDouble:
123+
return std::make_shared<::arrow::DoubleScalar>(std::get<double>(value));
124+
case TypeId::kDecimal: {
125+
const auto& decimal = std::get<Decimal>(value);
126+
::arrow::Decimal128 arrow_decimal(
127+
static_cast<int64_t>(decimal.value() >> 64),
128+
static_cast<uint64_t>(decimal.value() & ~uint64_t{0}));
129+
return std::make_shared<::arrow::Decimal128Scalar>(arrow_decimal,
130+
std::move(arrow_type));
131+
}
132+
case TypeId::kDate:
133+
return std::make_shared<::arrow::Date32Scalar>(std::get<int32_t>(value));
134+
case TypeId::kTime:
135+
return std::make_shared<::arrow::Time64Scalar>(std::get<int64_t>(value),
136+
std::move(arrow_type));
137+
case TypeId::kTimestamp:
138+
case TypeId::kTimestampTz:
139+
case TypeId::kTimestampNs:
140+
case TypeId::kTimestampTzNs:
141+
return std::make_shared<::arrow::TimestampScalar>(std::get<int64_t>(value),
142+
std::move(arrow_type));
143+
case TypeId::kString:
144+
return std::make_shared<::arrow::StringScalar>(std::get<std::string>(value));
145+
case TypeId::kBinary: {
146+
ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Buffer> buffer,
147+
ToArrowBuffer(std::get<std::vector<uint8_t>>(value), pool));
148+
return std::make_shared<::arrow::BinaryScalar>(std::move(buffer));
149+
}
150+
case TypeId::kFixed: {
151+
ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Buffer> buffer,
152+
ToArrowBuffer(std::get<std::vector<uint8_t>>(value), pool));
153+
return std::make_shared<::arrow::FixedSizeBinaryScalar>(std::move(buffer),
154+
std::move(arrow_type));
155+
}
156+
case TypeId::kUuid: {
157+
const Uuid& uuid = std::get<Uuid>(value);
158+
ICEBERG_ASSIGN_OR_RAISE(
159+
std::shared_ptr<::arrow::Buffer> buffer,
160+
ToArrowBuffer(std::vector<uint8_t>(uuid.bytes().begin(), uuid.bytes().end()),
161+
pool));
162+
return std::make_shared<::arrow::FixedSizeBinaryScalar>(std::move(buffer),
163+
std::move(arrow_type));
164+
}
165+
default:
166+
return NotSupported("Cannot convert {} literal to an Arrow scalar",
167+
*literal.type());
168+
}
169+
}
170+
171+
Result<std::shared_ptr<::arrow::Array>> MakeDefaultArray(
172+
const Literal& literal, const std::shared_ptr<::arrow::DataType>& type,
173+
int64_t num_rows, ::arrow::MemoryPool* pool) {
174+
// An extension type (e.g. `arrow.uuid` for an Iceberg UUID) is backed by a storage
175+
// type, and compute::Cast has no kernel that casts a storage array into an extension
176+
// type. Materialize the array as the storage type and wrap it in the extension type.
177+
if (type->id() == ::arrow::Type::EXTENSION) {
178+
const auto& extension_type =
179+
internal::checked_cast<const ::arrow::ExtensionType&>(*type);
180+
ICEBERG_ASSIGN_OR_RAISE(
181+
std::shared_ptr<::arrow::Array> storage,
182+
MakeDefaultArray(literal, extension_type.storage_type(), num_rows, pool));
183+
return ::arrow::ExtensionType::WrapArray(type, storage);
184+
}
185+
186+
ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr<::arrow::Scalar> scalar,
187+
ToArrowScalar(literal, pool));
188+
ICEBERG_ARROW_ASSIGN_OR_RETURN(std::shared_ptr<::arrow::Array> array,
189+
::arrow::MakeArrayFromScalar(*scalar, num_rows, pool));
190+
if (!array->type()->Equals(*type)) {
191+
ICEBERG_ARROW_ASSIGN_OR_RETURN(::arrow::Datum cast_result,
192+
::arrow::compute::Cast(array, type));
193+
return cast_result.make_array();
194+
}
195+
return array;
196+
}
197+
198+
} // namespace iceberg::arrow

‎src/iceberg/arrow/literal_util_internal.h‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
#pragma once
21+
22+
#include <cstdint>
23+
#include <memory>
24+
25+
#include <arrow/type_fwd.h>
26+
27+
#include "iceberg/expression/literal.h"
28+
#include "iceberg/result.h"
29+
30+
namespace iceberg::arrow {
31+
32+
/// \brief Convert a primitive literal to an Arrow scalar of its canonical Arrow type.
33+
///
34+
/// A null literal converts to a null scalar of the corresponding Arrow type. `pool` is
35+
/// used for the backing buffer of binary/fixed/uuid scalars.
36+
Result<std::shared_ptr<::arrow::Scalar>> ToArrowScalar(const Literal& literal,
37+
::arrow::MemoryPool* pool);
38+
39+
/// \brief Create an Arrow array of `num_rows` rows where every row holds the literal
40+
/// value, e.g. to materialize a missing column with a default value.
41+
///
42+
/// The array is cast to `type` when the literal's canonical Arrow type differs.
43+
Result<std::shared_ptr<::arrow::Array>> MakeDefaultArray(
44+
const Literal& literal, const std::shared_ptr<::arrow::DataType>& type,
45+
int64_t num_rows, ::arrow::MemoryPool* pool);
46+
47+
} // namespace iceberg::arrow

‎src/iceberg/parquet/parquet_data_util.cc‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
#include <arrow/type.h>
2525

2626
#include "iceberg/arrow/arrow_status_internal.h"
27+
#include "iceberg/arrow/literal_util_internal.h"
2728
#include "iceberg/metadata_columns.h"
2829
#include "iceberg/parquet/parquet_data_util_internal.h"
2930
#include "iceberg/schema.h"
@@ -119,6 +120,11 @@ Result<std::shared_ptr<::arrow::Array>> ProjectStructArray(
119120
ICEBERG_ASSIGN_OR_RAISE(
120121
projected_array,
121122
MakeNullArray(output_arrow_type, struct_array->length(), pool));
123+
} else if (field_projection.kind == FieldProjection::Kind::kDefault) {
124+
ICEBERG_ASSIGN_OR_RAISE(
125+
projected_array,
126+
arrow::MakeDefaultArray(std::get<Literal>(field_projection.from),
127+
output_arrow_type, struct_array->length(), pool));
122128
} else if (field_projection.kind == FieldProjection::Kind::kMetadata) {
123129
int32_t field_id = projected_field.field_id();
124130
if (field_id == MetadataColumns::kFilePathColumnId) {

‎src/iceberg/parquet/parquet_schema_util.cc‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -342,6 +342,10 @@ Result<FieldProjection> ProjectStruct(
342342
child_projection, ProjectField(field, parquet_field, iter->second.local_index));
343343
} else if (MetadataColumns::IsMetadataColumn(field_id)) {
344344
child_projection.kind = FieldProjection::Kind::kMetadata;
345+
} else if (field.initial_default() != nullptr) {
346+
// Rows written before the field existed assume its `initial-default` value.
347+
child_projection.kind = FieldProjection::Kind::kDefault;
348+
child_projection.from = *field.initial_default();
345349
} else if (field.optional()) {
346350
child_projection.kind = FieldProjection::Kind::kNull;
347351
} else {

‎src/iceberg/test/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -256,7 +256,8 @@ if(ICEBERG_BUILD_BUNDLE)
256256
data_writer_test.cc
257257
delete_filter_test.cc
258258
delete_loader_test.cc
259-
file_scan_task_reader_test.cc)
259+
file_scan_task_reader_test.cc
260+
literal_util_test.cc)
260261

261262
endif()
262263

0 commit comments

Comments
 (0)