zhjwpku commented on code in PR #552:
URL: https://github.com/apache/iceberg-cpp/pull/552#discussion_r2757455710
##########
src/iceberg/data/data_writer.cc:
##########
@@ -19,20 +19,124 @@
#include "iceberg/data/data_writer.h"
+#include "iceberg/file_writer.h"
+#include "iceberg/manifest/manifest_entry.h"
+#include "iceberg/util/macros.h"
+
namespace iceberg {
class DataWriter::Impl {
public:
+ explicit Impl(DataWriterOptions options) : options_(std::move(options)) {}
+
+ Status Initialize() {
+ WriterOptions writer_options;
+ writer_options.path = options_.path;
+ writer_options.schema = options_.schema;
+ writer_options.io = options_.io;
+ writer_options.properties = WriterProperties::FromMap(options_.properties);
+
+ ICEBERG_ASSIGN_OR_RAISE(writer_,
+ WriterFactoryRegistry::Open(options_.format,
writer_options));
+ return {};
+ }
+
+ Status Write(ArrowArray* data) {
+ if (!writer_) {
+ return InvalidArgument("Writer not initialized");
+ }
Review Comment:
```suggestion
ICEBERG_PRECHECK(writer_, "Writer not initialized");
```
nit, this should make the code shorter.
##########
src/iceberg/data/data_writer.cc:
##########
@@ -19,20 +19,124 @@
#include "iceberg/data/data_writer.h"
+#include "iceberg/file_writer.h"
+#include "iceberg/manifest/manifest_entry.h"
+#include "iceberg/util/macros.h"
+
namespace iceberg {
class DataWriter::Impl {
public:
+ explicit Impl(DataWriterOptions options) : options_(std::move(options)) {}
+
+ Status Initialize() {
+ WriterOptions writer_options;
+ writer_options.path = options_.path;
+ writer_options.schema = options_.schema;
+ writer_options.io = options_.io;
+ writer_options.properties = WriterProperties::FromMap(options_.properties);
+
+ ICEBERG_ASSIGN_OR_RAISE(writer_,
+ WriterFactoryRegistry::Open(options_.format,
writer_options));
+ return {};
+ }
+
+ Status Write(ArrowArray* data) {
+ if (!writer_) {
+ return InvalidArgument("Writer not initialized");
+ }
+ return writer_->Write(data);
+ }
+
+ Result<int64_t> Length() const {
+ if (!writer_) {
+ return InvalidArgument("Writer not initialized");
+ }
+ return writer_->length();
+ }
+
+ Status Close() {
+ if (!writer_) {
+ return InvalidArgument("Writer not initialized");
+ }
+ if (closed_) {
+ return InvalidArgument("Writer already closed");
+ }
+ ICEBERG_RETURN_UNEXPECTED(writer_->Close());
+ closed_ = true;
+ return {};
+ }
+
+ Result<FileWriter::WriteResult> Metadata() {
+ if (!closed_) {
Review Comment:
nit: use ICEBERG_CHECK here
##########
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:
nit: move this closing namespace curly before the first TEST_F?
--
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]
