Skip to content

Commit 6dff5ac

Browse files
huan233uschuan233usc
authored
test: cover write, evolve and read back of a defaulted column end to end (#874)
1 parent d998c8f commit 6dff5ac

2 files changed

Lines changed: 267 additions & 0 deletions

File tree

‎src/iceberg/test/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -249,6 +249,7 @@ if(ICEBERG_BUILD_BUNDLE)
249249
arrow_c_data_util_test.cc
250250
arrow_row_builder_test.cc
251251
data_writer_test.cc
252+
default_value_test.cc
252253
delete_filter_test.cc
253254
delete_loader_test.cc
254255
deletion_vector_writer_test.cc

‎src/iceberg/test/default_value_test.cc‎

Lines changed: 266 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,266 @@
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 <array>
21+
#include <format>
22+
#include <functional>
23+
#include <memory>
24+
#include <string>
25+
#include <string_view>
26+
#include <vector>
27+
28+
#include <arrow/array.h>
29+
#include <arrow/array/builder_binary.h>
30+
#include <arrow/c/bridge.h>
31+
#include <arrow/extension/uuid.h>
32+
#include <arrow/extension_type.h>
33+
#include <arrow/filesystem/mockfs.h>
34+
#include <arrow/json/from_string.h>
35+
#include <arrow/type.h>
36+
#include <arrow/util/logging.h>
37+
#include <gtest/gtest.h>
38+
39+
#include "iceberg/arrow/arrow_io_internal.h"
40+
#include "iceberg/avro/avro_register.h"
41+
#include "iceberg/expression/literal.h"
42+
#include "iceberg/file_format.h"
43+
#include "iceberg/file_reader.h"
44+
#include "iceberg/file_writer.h"
45+
#include "iceberg/parquet/parquet_register.h"
46+
#include "iceberg/schema.h"
47+
#include "iceberg/schema_field.h"
48+
#include "iceberg/schema_internal.h"
49+
#include "iceberg/table.h"
50+
#include "iceberg/table_identifier.h"
51+
#include "iceberg/table_metadata.h"
52+
#include "iceberg/test/matchers.h"
53+
#include "iceberg/test/update_test_base.h"
54+
#include "iceberg/type.h"
55+
#include "iceberg/update/update_schema.h"
56+
#include "iceberg/util/macros.h"
57+
#include "iceberg/util/uuid.h"
58+
59+
namespace iceberg {
60+
61+
namespace {
62+
63+
struct DefaultValueEndToEndParam {
64+
std::string name;
65+
FileFormatType format;
66+
std::shared_ptr<Type> type;
67+
Literal default_value;
68+
// Builds the two-row column the added field is expected to read back as. A factory
69+
// rather than JSON so extension types (uuid) can be constructed explicitly.
70+
std::function<std::shared_ptr<::arrow::Array>()> expected_column;
71+
// Avro reads through one of two backends depending on this; unused for Parquet.
72+
bool avro_skip_datum = true;
73+
};
74+
75+
class DefaultValueEndToEndTest
76+
: public UpdateTestBase,
77+
public ::testing::WithParamInterface<DefaultValueEndToEndParam> {
78+
protected:
79+
static void SetUpTestSuite() {
80+
parquet::RegisterAll();
81+
avro::RegisterAll();
82+
}
83+
84+
std::string MetadataResource() const override {
85+
return "TableMetadataV3ValidMinimal.json";
86+
}
87+
};
88+
89+
TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) {
90+
const auto& param = GetParam();
91+
ICEBERG_UNWRAP_OR_FAIL(auto original_schema, table_->schema());
92+
93+
ArrowSchema arrow_c_schema;
94+
ASSERT_THAT(ToArrowSchema(*original_schema, &arrow_c_schema), IsOk());
95+
auto arrow_type = ::arrow::ImportType(&arrow_c_schema).ValueOrDie();
96+
auto array = ::arrow::json::ArrayFromJSONString(::arrow::struct_(arrow_type->fields()),
97+
R"([[1, 2, 3],[4, 5, 6]])")
98+
.ValueOrDie();
99+
100+
// Arrow's mock filesystem does not create directories on demand.
101+
auto arrow_fs = std::dynamic_pointer_cast<::arrow::fs::internal::MockFileSystem>(
102+
static_cast<arrow::ArrowFileSystemFileIO&>(*file_io_).fs());
103+
ASSERT_TRUE(arrow_fs != nullptr);
104+
ASSERT_TRUE(arrow_fs->CreateDir(table_location_ + "/data").ok());
105+
106+
const std::string data_file =
107+
std::format("{}/data/evolve-default-{}", table_location_, param.name);
108+
WriterProperties writer_properties;
109+
writer_properties.Set(WriterProperties::kParquetCompression,
110+
std::string("uncompressed"));
111+
ICEBERG_UNWRAP_OR_FAIL(
112+
auto writer,
113+
WriterFactoryRegistry::Open(param.format, {.path = data_file,
114+
.schema = original_schema,
115+
.io = file_io_,
116+
.properties = writer_properties}));
117+
ArrowArray exported;
118+
ASSERT_TRUE(::arrow::ExportArray(*array, &exported).ok());
119+
ASSERT_THAT(writer->Write(&exported), IsOk());
120+
ASSERT_THAT(writer->Close(), IsOk());
121+
122+
ICEBERG_UNWRAP_OR_FAIL(auto update, table_->NewUpdateSchema());
123+
update->AddColumn("added", param.type, "Added after the data was written",
124+
param.default_value);
125+
ASSERT_THAT(update->Commit(), IsOk());
126+
127+
ICEBERG_UNWRAP_OR_FAIL(auto reloaded, catalog_->LoadTable(table_ident_));
128+
ICEBERG_UNWRAP_OR_FAIL(auto evolved_schema, reloaded->schema());
129+
ICEBERG_UNWRAP_OR_FAIL(auto added_field, evolved_schema->FindFieldByName("added"));
130+
ASSERT_TRUE(added_field.has_value());
131+
// AddColumn sets both defaults, so evolution must carry both values through the commit
132+
// and the metadata round trip.
133+
ASSERT_NE(added_field->get().initial_default(), nullptr);
134+
ASSERT_NE(added_field->get().write_default(), nullptr);
135+
EXPECT_EQ(*added_field->get().initial_default(), param.default_value);
136+
EXPECT_EQ(*added_field->get().write_default(), param.default_value);
137+
138+
ReaderProperties reader_properties;
139+
reader_properties.Set(ReaderProperties::kAvroSkipDatum, param.avro_skip_datum);
140+
ICEBERG_UNWRAP_OR_FAIL(
141+
auto reader,
142+
ReaderFactoryRegistry::Open(param.format, {.path = data_file,
143+
.io = file_io_,
144+
.projection = evolved_schema,
145+
.properties = reader_properties}));
146+
147+
// The reader reports the evolved schema, including the added column.
148+
ICEBERG_UNWRAP_OR_FAIL(auto reported_c_schema, reader->Schema());
149+
auto reported_schema = ::arrow::ImportSchema(&reported_c_schema).ValueOrDie();
150+
ASSERT_EQ(reported_schema->num_fields(), 4);
151+
EXPECT_EQ(reported_schema->field(3)->name(), "added");
152+
ICEBERG_UNWRAP_OR_FAIL(auto batch, reader->Next());
153+
ASSERT_TRUE(batch.has_value());
154+
155+
ArrowSchema read_c_schema;
156+
ASSERT_THAT(ToArrowSchema(*evolved_schema, &read_c_schema), IsOk());
157+
auto read_type = ::arrow::ImportType(&read_c_schema).ValueOrDie();
158+
auto actual = ::arrow::ImportArray(&batch.value(), read_type).ValueOrDie();
159+
auto actual_struct = std::static_pointer_cast<::arrow::StructArray>(actual);
160+
ASSERT_EQ(actual_struct->length(), 2);
161+
162+
// The pre-existing columns are unchanged.
163+
auto expected_existing = ::arrow::json::ArrayFromJSONString(
164+
::arrow::struct_({read_type->field(0), read_type->field(1),
165+
read_type->field(2)}),
166+
R"([[1, 2, 3],[4, 5, 6]])")
167+
.ValueOrDie();
168+
auto actual_existing =
169+
::arrow::StructArray::Make(
170+
{actual_struct->field(0), actual_struct->field(1), actual_struct->field(2)},
171+
{read_type->field(0), read_type->field(1), read_type->field(2)})
172+
.ValueOrDie();
173+
ASSERT_TRUE(actual_existing->Equals(*expected_existing))
174+
<< "actual: " << actual_existing->ToString()
175+
<< "\nexpected: " << expected_existing->ToString();
176+
177+
// The added column materializes its initial-default for every row.
178+
auto expected_added = param.expected_column();
179+
auto actual_added = actual_struct->field(3);
180+
ASSERT_TRUE(actual_added->Equals(*expected_added))
181+
<< "actual: " << actual_added->ToString()
182+
<< "\nexpected: " << expected_added->ToString();
183+
184+
// The file held exactly those two rows.
185+
ICEBERG_UNWRAP_OR_FAIL(auto next_batch, reader->Next());
186+
ASSERT_FALSE(next_batch.has_value());
187+
}
188+
189+
namespace {
190+
191+
// Two-row column of `json` values at `type`, for the simple cases.
192+
std::function<std::shared_ptr<::arrow::Array>()> Rows(
193+
std::shared_ptr<::arrow::DataType> type, std::string json) {
194+
return [type = std::move(type), json = std::move(json)]() {
195+
return ::arrow::json::ArrayFromJSONString(type, json).ValueOrDie();
196+
};
197+
}
198+
199+
constexpr std::array<uint8_t, 16> kUuidBytes = {0xF7, 0x9C, 0x3E, 0x09, 0x67, 0x7C,
200+
0x4B, 0xBD, 0x90, 0x38, 0x1E, 0x7F,
201+
0x1A, 0x0C, 0x8B, 0x2D};
202+
203+
std::shared_ptr<::arrow::Array> UuidRows() {
204+
// An Iceberg uuid surfaces in Arrow as the `arrow.uuid` extension type over
205+
// fixed_size_binary(16), so the storage array is built then wrapped.
206+
::arrow::FixedSizeBinaryBuilder builder(::arrow::fixed_size_binary(Uuid::kLength));
207+
ARROW_CHECK_OK(builder.Append(kUuidBytes.data()));
208+
ARROW_CHECK_OK(builder.Append(kUuidBytes.data()));
209+
return ::arrow::ExtensionType::WrapArray(::arrow::extension::uuid(),
210+
builder.Finish().ValueOrDie());
211+
}
212+
213+
} // namespace
214+
215+
INSTANTIATE_TEST_SUITE_P(
216+
FormatsAndTypes, DefaultValueEndToEndTest,
217+
::testing::Values(
218+
DefaultValueEndToEndParam{"parquet_long", FileFormatType::kParquet, int64(),
219+
Literal::Long(42), Rows(::arrow::int64(), "[42, 42]")},
220+
DefaultValueEndToEndParam{"avro_long", FileFormatType::kAvro, int64(),
221+
Literal::Long(42), Rows(::arrow::int64(), "[42, 42]")},
222+
DefaultValueEndToEndParam{"parquet_string", FileFormatType::kParquet, string(),
223+
Literal::String("iceberg"),
224+
Rows(::arrow::utf8(), R"(["iceberg", "iceberg"])")},
225+
DefaultValueEndToEndParam{"avro_string", FileFormatType::kAvro, string(),
226+
Literal::String("iceberg"),
227+
Rows(::arrow::utf8(), R"(["iceberg", "iceberg"])")},
228+
DefaultValueEndToEndParam{
229+
"parquet_decimal", FileFormatType::kParquet, decimal(9, 2),
230+
Literal::Decimal(12345, 9, 2),
231+
Rows(::arrow::decimal128(9, 2), R"(["123.45", "123.45"])")},
232+
DefaultValueEndToEndParam{
233+
"avro_decimal", FileFormatType::kAvro, decimal(9, 2),
234+
Literal::Decimal(12345, 9, 2),
235+
Rows(::arrow::decimal128(9, 2), R"(["123.45", "123.45"])")},
236+
DefaultValueEndToEndParam{"parquet_timestamp", FileFormatType::kParquet,
237+
timestamp(), Literal::Timestamp(1672531200000000),
238+
Rows(::arrow::timestamp(::arrow::TimeUnit::MICRO),
239+
"[1672531200000000, 1672531200000000]")},
240+
DefaultValueEndToEndParam{"avro_timestamp", FileFormatType::kAvro, timestamp(),
241+
Literal::Timestamp(1672531200000000),
242+
Rows(::arrow::timestamp(::arrow::TimeUnit::MICRO),
243+
"[1672531200000000, 1672531200000000]")},
244+
DefaultValueEndToEndParam{"parquet_uuid", FileFormatType::kParquet, uuid(),
245+
Literal::UUID(Uuid::FromBytes(kUuidBytes).value()),
246+
UuidRows},
247+
DefaultValueEndToEndParam{"avro_uuid", FileFormatType::kAvro, uuid(),
248+
Literal::UUID(Uuid::FromBytes(kUuidBytes).value()),
249+
UuidRows},
250+
// The Avro reader has two backends; the cases above use the default
251+
// (skip-datum), so cover the GenericDatum path too.
252+
DefaultValueEndToEndParam{"avro_long_generic_datum", FileFormatType::kAvro,
253+
int64(), Literal::Long(42),
254+
Rows(::arrow::int64(), "[42, 42]"),
255+
/*avro_skip_datum=*/false},
256+
DefaultValueEndToEndParam{"avro_uuid_generic_datum", FileFormatType::kAvro,
257+
uuid(),
258+
Literal::UUID(Uuid::FromBytes(kUuidBytes).value()),
259+
UuidRows, /*avro_skip_datum=*/false}),
260+
[](const ::testing::TestParamInfo<DefaultValueEndToEndParam>& info) {
261+
return info.param.name;
262+
});
263+
264+
} // namespace
265+
266+
} // namespace iceberg

0 commit comments

Comments
 (0)