shangxinli commented on code in PR #552:
URL: https://github.com/apache/iceberg-cpp/pull/552#discussion_r2777600403
##########
src/iceberg/test/data_writer_test.cc:
##########
@@ -17,7 +17,414 @@
* under the License.
*/
+#include "iceberg/data/data_writer.h"
+
+#include <arrow/array.h>
+#include <arrow/c/bridge.h>
+#include <arrow/json/from_string.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
-namespace iceberg {} // namespace iceberg
+#include "iceberg/arrow/arrow_fs_file_io_internal.h"
+#include "iceberg/avro/avro_register.h"
+#include "iceberg/file_format.h"
+#include "iceberg/manifest/manifest_entry.h"
+#include "iceberg/parquet/parquet_register.h"
+#include "iceberg/partition_spec.h"
+#include "iceberg/row/partition_values.h"
+#include "iceberg/schema.h"
+#include "iceberg/schema_field.h"
+#include "iceberg/schema_internal.h"
+#include "iceberg/test/matchers.h"
+#include "iceberg/type.h"
+#include "iceberg/util/macros.h"
+
+namespace iceberg {
+namespace {
+
+using ::testing::HasSubstr;
+
+class DataWriterTest : public ::testing::Test {
+ protected:
+ static void SetUpTestSuite() {
+ parquet::RegisterAll();
+ avro::RegisterAll();
+ }
+
+ void SetUp() override {
+ file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO();
+ schema_ = std::make_shared<Schema>(std::vector<SchemaField>{
+ SchemaField::MakeRequired(1, "id", std::make_shared<IntType>()),
+ SchemaField::MakeOptional(2, "name", std::make_shared<StringType>())});
+ partition_spec_ = PartitionSpec::Unpartitioned();
+ }
+
+ std::shared_ptr<::arrow::Array> CreateTestData() {
+ ArrowSchema arrow_c_schema;
+ ICEBERG_THROW_NOT_OK(ToArrowSchema(*schema_, &arrow_c_schema));
+ auto arrow_schema = ::arrow::ImportType(&arrow_c_schema).ValueOrDie();
+
+ return ::arrow::json::ArrayFromJSONString(
+ ::arrow::struct_(arrow_schema->fields()),
+ R"([[1, "Alice"], [2, "Bob"], [3, "Charlie"]])")
+ .ValueOrDie();
+ }
+
+ std::shared_ptr<FileIO> file_io_;
+ std::shared_ptr<Schema> schema_;
+ std::shared_ptr<PartitionSpec> partition_spec_;
+};
+
+TEST_F(DataWriterTest, CreateWithParquetFormat) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+ ASSERT_NE(writer, nullptr);
+}
+
+TEST_F(DataWriterTest, CreateWithAvroFormat) {
+ DataWriterOptions options{
+ .path = "test_data.avro",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kAvro,
+ .io = file_io_,
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+ ASSERT_NE(writer, nullptr);
+}
+
+TEST_F(DataWriterTest, WriteAndClose) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+
+ // Check length before close
+ auto length_result = writer->Length();
+ ASSERT_THAT(length_result, IsOk());
+ EXPECT_GT(length_result.value(), 0);
+
+ // Close
+ ASSERT_THAT(writer->Close(), IsOk());
+}
+
+TEST_F(DataWriterTest, GetMetadataAfterClose) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+
+ // Close
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Get metadata
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsOk());
+
+ const auto& write_result = metadata_result.value();
+ ASSERT_EQ(write_result.data_files.size(), 1);
+
+ const auto& data_file = write_result.data_files[0];
+ EXPECT_EQ(data_file->content, DataFile::Content::kData);
+ EXPECT_EQ(data_file->file_path, "test_data.parquet");
+ EXPECT_EQ(data_file->file_format, FileFormatType::kParquet);
+ // Record count may be 0 or 3 depending on Parquet writer metrics support
+ EXPECT_GE(data_file->record_count, 0);
+ EXPECT_GT(data_file->file_size_in_bytes, 0);
+}
+
+TEST_F(DataWriterTest, MetadataBeforeCloseReturnsError) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Try to get metadata before closing
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsError(ErrorKind::kInvalidArgument));
+ EXPECT_THAT(metadata_result,
+ HasErrorMessage("Cannot get metadata before closing the
writer"));
+}
+
+TEST_F(DataWriterTest, DoubleCloseReturnsError) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+
+ // Close once
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Try to close again
+ auto second_close = writer->Close();
+ ASSERT_THAT(second_close, IsError(ErrorKind::kInvalidArgument));
+ EXPECT_THAT(second_close, HasErrorMessage("Writer already closed"));
+}
+
+TEST_F(DataWriterTest, SortOrderIdPreserved) {
+ const int32_t sort_order_id = 42;
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .sort_order_id = sort_order_id,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Check metadata
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsOk());
+ const auto& data_file = metadata_result.value().data_files[0];
+ ASSERT_TRUE(data_file->sort_order_id.has_value());
+ EXPECT_EQ(data_file->sort_order_id.value(), sort_order_id);
+}
+
+TEST_F(DataWriterTest, SortOrderIdNullByDefault) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ // sort_order_id not set
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Check metadata
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsOk());
+ const auto& data_file = metadata_result.value().data_files[0];
+ EXPECT_FALSE(data_file->sort_order_id.has_value());
+}
+
+TEST_F(DataWriterTest, MetadataContainsColumnMetrics) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Check metadata
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsOk());
+ const auto& data_file = metadata_result.value().data_files[0];
+
+ // Metrics availability depends on the underlying writer implementation
+ // Just verify the maps exist (they may be empty depending on writer config)
+ EXPECT_GE(data_file->column_sizes.size(), 0);
+ EXPECT_GE(data_file->value_counts.size(), 0);
+ EXPECT_GE(data_file->null_value_counts.size(), 0);
+}
+
+TEST_F(DataWriterTest, PartitionValuesPreserved) {
+ // Create partition values with a sample value
+ PartitionValues partition_values({Literal::Int(42),
Literal::String("test")});
+
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = partition_values,
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Check metadata
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsOk());
+ const auto& data_file = metadata_result.value().data_files[0];
+
+ // Verify partition values are preserved
+ EXPECT_EQ(data_file->partition.num_fields(), partition_values.num_fields());
+ EXPECT_EQ(data_file->partition.num_fields(), 2);
+}
+
+TEST_F(DataWriterTest, WriteMultipleBatches) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write first batch
+ auto test_data1 = CreateTestData();
+ ArrowArray arrow_array1;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data1, &arrow_array1).ok());
+ ASSERT_THAT(writer->Write(&arrow_array1), IsOk());
+
+ // Write second batch
+ auto test_data2 = CreateTestData();
+ ArrowArray arrow_array2;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data2, &arrow_array2).ok());
+ ASSERT_THAT(writer->Write(&arrow_array2), IsOk());
+
+ ASSERT_THAT(writer->Close(), IsOk());
+
+ // Check metadata - file should exist with data
+ auto metadata_result = writer->Metadata();
+ ASSERT_THAT(metadata_result, IsOk());
+ const auto& data_file = metadata_result.value().data_files[0];
+ // Record count depends on writer metrics support
+ EXPECT_GE(data_file->record_count, 0);
+ EXPECT_GT(data_file->file_size_in_bytes, 0);
+}
+
+TEST_F(DataWriterTest, LengthIncreasesAfterWrite) {
+ DataWriterOptions options{
+ .path = "test_data.parquet",
+ .schema = schema_,
+ .spec = partition_spec_,
+ .partition = PartitionValues{},
+ .format = FileFormatType::kParquet,
+ .io = file_io_,
+ .properties = {{"write.parquet.compression-codec", "uncompressed"}},
+ };
+
+ auto writer_result = DataWriter::Make(options);
+ ASSERT_THAT(writer_result, IsOk());
+ auto writer = std::move(writer_result.value());
+
+ // Write data
+ auto test_data = CreateTestData();
+ ArrowArray arrow_array;
+ ASSERT_TRUE(::arrow::ExportArray(*test_data, &arrow_array).ok());
+ ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+
+ // Length should be greater than 0 after write
+ auto length = writer->Length();
+ ASSERT_THAT(length, IsOk());
+ EXPECT_GT(length.value(), 0);
+}
+
+} // namespace
Review Comment:
done
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
