✨ feat(data): Apply write defaults in writers

DataWriter aligns batches to the write schema when the input schema
differs, via a new Write(data, input_schema) overload; schemas that
already match keep the zero-conversion Write(data) path. Adds Parquet
and Avro end-to-end tests where a column added with write-default
reads back the default for rows written without it.

Author: huan233usc

src/iceberg/avro/avro_writer.cc

@@ -32,6 +32,7 @@
 
 #include "iceberg/arrow/arrow_io_internal.h"
 #include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow_c_data_util_internal.h"
 #include "iceberg/avro/avro_data_util_internal.h"
 #include "iceberg/avro/avro_direct_encoder_internal.h"
 #include "iceberg/avro/avro_metrics.h"
@@ -241,6 +242,14 @@ class AvroWriter::Impl {
     return {};
   }
 
+  Status Write(ArrowArray* data, const Schema& input_schema) {
+    ICEBERG_ASSIGN_OR_RAISE(
+        auto aligned_data,
+        arrow::AlignBatchForWrite(data, input_schema, *write_schema_,
+                                  ::arrow::default_memory_pool()));
+    return Write(&aligned_data);
+  }
+
   Status Close() {
     if (!backend_->Closed()) {
       backend_->Close();
@@ -290,6 +299,10 @@ AvroWriter::~AvroWriter() = default;
 
 Status AvroWriter::Write(ArrowArray* data) { return impl_->Write(data); }
 
+Status AvroWriter::Write(ArrowArray* data, const Schema& input_schema) {
+  return impl_->Write(data, input_schema);
+}
+
 Status AvroWriter::Open(const WriterOptions& options) {
   impl_ = std::make_unique<Impl>();
   return impl_->Open(options);

src/iceberg/avro/avro_writer.h

@@ -40,6 +40,8 @@ class ICEBERG_BUNDLE_EXPORT AvroWriter : public Writer {
 
   Status Write(ArrowArray* data) final;
 
+  Status Write(ArrowArray* data, const Schema& input_schema) final;
+
   Result<Metrics> metrics() final;
 
   Result<int64_t> length() final;

src/iceberg/data/data_writer.cc

@@ -26,13 +26,16 @@
 #include "iceberg/file_writer.h"
 #include "iceberg/manifest/manifest_entry.h"
 #include "iceberg/partition_spec.h"
+#include "iceberg/schema.h"
 #include "iceberg/util/macros.h"
 
 namespace iceberg {
 
 class DataWriter::Impl {
  public:
   static Result<std::unique_ptr<Impl>> Make(DataWriterOptions options) {
+    ICEBERG_PRECHECK(options.schema != nullptr, "Write schema must not be null");
+    ICEBERG_PRECHECK(options.input_schema != nullptr, "Input schema must not be null");
     WriterOptions writer_options{
         .path = options.path,
         .schema = options.schema,
@@ -46,7 +49,12 @@ class DataWriter::Impl {
     return std::unique_ptr<Impl>(new Impl(std::move(options), std::move(writer)));
   }
 
-  Status Write(ArrowArray* data) { return writer_->Write(data); }
+  Status Write(ArrowArray* data) {
+    if (*options_.input_schema == *options_.schema) {
+      return writer_->Write(data);
+    }
+    return writer_->Write(data, *options_.input_schema);
+  }
 
   Result<int64_t> Length() const { return writer_->length(); }
 

src/iceberg/data/data_writer.h

@@ -42,6 +42,8 @@ namespace iceberg {
 struct ICEBERG_DATA_EXPORT DataWriterOptions {
   std::string path;
   std::shared_ptr<Schema> schema;
+  /// Required schema of Arrow batches passed to Write.
+  std::shared_ptr<Schema> input_schema;
   std::shared_ptr<PartitionSpec> spec;
   PartitionValues partition;
   FileFormatType format = FileFormatType::kParquet;

src/iceberg/file_writer.h

@@ -75,7 +75,7 @@ class ICEBERG_EXPORT WriterProperties : public ConfigBase<WriterProperties> {
 struct ICEBERG_EXPORT WriterOptions {
   /// \brief The path to the file to write.
   std::string path;
-  /// \brief The schema of the data to write.
+  /// \brief The target schema to write to the file.
   std::shared_ptr<Schema> schema;
   /// \brief FileIO instance to create the file.
   std::shared_ptr<class FileIO> io;
@@ -107,6 +107,12 @@ class ICEBERG_EXPORT Writer {
   /// \note Ownership of the data is transferred to the writer.
   virtual Status Write(ArrowArray* data) = 0;
 
+  /// \brief Write Arrow data described by an input schema that may differ from the
+  /// writer schema.
+  ///
+  /// \note Ownership of the data is transferred to the writer.
+  virtual Status Write(ArrowArray* data, const Schema& input_schema) = 0;
+
   /// \brief Get the file statistics.
   /// Only valid after the file is closed.
   virtual Result<Metrics> metrics() = 0;

src/iceberg/parquet/parquet_writer.cc

@@ -39,6 +39,7 @@
 
 #include "iceberg/arrow/arrow_io_internal.h"
 #include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow_c_data_util_internal.h"
 #include "iceberg/parquet/parquet_metrics_internal.h"
 #include "iceberg/parquet/parquet_schema_util_internal.h"
 #include "iceberg/schema_internal.h"
@@ -295,6 +296,13 @@ class ParquetWriter::Impl {
     return {};
   }
 
+  Status Write(ArrowArray* array, const Schema& input_schema) {
+    ICEBERG_ASSIGN_OR_RAISE(
+        auto aligned_array,
+        arrow::AlignBatchForWrite(array, input_schema, *schema_, pool_));
+    return Write(&aligned_array);
+  }
+
   // Close the writer and release resources
   Status Close() {
     if (writer_ == nullptr) {
@@ -371,6 +379,10 @@ Status ParquetWriter::Open(const WriterOptions& options) {
 
 Status ParquetWriter::Write(ArrowArray* array) { return impl_->Write(array); }
 
+Status ParquetWriter::Write(ArrowArray* array, const Schema& input_schema) {
+  return impl_->Write(array, input_schema);
+}
+
 Status ParquetWriter::Close() { return impl_->Close(); }
 
 Result<Metrics> ParquetWriter::metrics() { return impl_->metrics(); }

src/iceberg/parquet/parquet_writer.h

@@ -40,6 +40,8 @@ class ICEBERG_BUNDLE_EXPORT ParquetWriter : public Writer {
 
   Status Write(ArrowArray* array) final;
 
+  Status Write(ArrowArray* array, const Schema& input_schema) final;
+
   Result<Metrics> metrics() final;
 
   Result<int64_t> length() final;

src/iceberg/test/data_writer_test.cc

@@ -71,6 +71,7 @@ class DataWriterTest : public ::testing::Test {
     return DataWriterOptions{
         .path = "test_data.parquet",
         .schema = schema_,
+        .input_schema = schema_,
         .spec = partition_spec_,
         .partition = std::move(partition),
         .format = FileFormatType::kParquet,
@@ -138,11 +139,22 @@ class DataWriterFormatTest
     : public DataWriterTest,
       public ::testing::WithParamInterface<std::pair<FileFormatType, std::string>> {};
 
+TEST_F(DataWriterTest, RejectsMissingInputSchema) {
+  auto options = MakeDefaultOptions();
+  options.input_schema.reset();
+
+  auto result = DataWriter::Make(options);
+
+  EXPECT_THAT(result, IsError(ErrorKind::kInvalidArgument));
+  EXPECT_THAT(result, HasErrorMessage("Input schema must not be null"));
+}
+
 TEST_P(DataWriterFormatTest, CreateWithFormat) {
   auto [format, path] = GetParam();
   DataWriterOptions options{
       .path = path,
       .schema = schema_,
+      .input_schema = schema_,
       .spec = partition_spec_,
       .partition = PartitionValues{},
       .format = format,
@@ -162,6 +174,7 @@ TEST_P(DataWriterFormatTest, WriteRowLineage) {
   DataWriterOptions options{
       .path = path,
       .schema = schema,
+      .input_schema = schema,
       .spec = partition_spec_,
       .partition = PartitionValues{},
       .format = format,

src/iceberg/test/default_value_test.cc

@@ -23,6 +23,8 @@
 #include <memory>
 #include <string>
 #include <string_view>
+#include <unordered_map>
+#include <utility>
 #include <vector>
 
 #include <arrow/array.h>
@@ -38,11 +40,14 @@
 
 #include "iceberg/arrow/arrow_io_internal.h"
 #include "iceberg/avro/avro_register.h"
+#include "iceberg/data/data_writer.h"
 #include "iceberg/expression/literal.h"
 #include "iceberg/file_format.h"
 #include "iceberg/file_reader.h"
 #include "iceberg/file_writer.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"
@@ -72,6 +77,12 @@ struct DefaultValueEndToEndParam {
   bool avro_skip_datum = true;
 };
 
+struct WriteDefaultEndToEndParam {
+  std::string name;
+  FileFormatType format;
+  std::string path;
+};
+
 class DefaultValueEndToEndTest
     : public UpdateTestBase,
       public ::testing::WithParamInterface<DefaultValueEndToEndParam> {
@@ -86,6 +97,37 @@ class DefaultValueEndToEndTest
   }
 };
 
+class WriteDefaultEndToEndTest
+    : public ::testing::TestWithParam<WriteDefaultEndToEndParam> {
+ protected:
+  static void SetUpTestSuite() {
+    parquet::RegisterAll();
+    avro::RegisterAll();
+  }
+
+  void SetUp() override { file_io_ = arrow::ArrowFileSystemFileIO::MakeMockFileIO(); }
+
+  std::shared_ptr<::arrow::Array> CreateArray(const Schema& schema,
+                                              std::string_view json) {
+    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()),
+                                              std::string(json))
+        .ValueOrDie();
+  }
+
+  static std::unordered_map<std::string, std::string> FormatProperties(
+      FileFormatType format) {
+    if (format == FileFormatType::kParquet) {
+      return {{"write.parquet.compression-codec", "uncompressed"}};
+    }
+    return {};
+  }
+
+  std::shared_ptr<FileIO> file_io_;
+};
+
 TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) {
   const auto& param = GetParam();
   ICEBERG_UNWRAP_OR_FAIL(auto original_schema, table_->schema());
@@ -186,6 +228,50 @@ TEST_P(DefaultValueEndToEndTest, WriteEvolveReadFillsInitialDefault) {
   ASSERT_FALSE(next_batch.has_value());
 }
 
+TEST_P(WriteDefaultEndToEndTest, MissingColumnUsesWriteDefault) {
+  const auto& param = GetParam();
+  auto input_schema = std::make_shared<Schema>(
+      std::vector<SchemaField>{SchemaField::MakeRequired(1, "id", int32())});
+  auto write_schema = std::make_shared<Schema>(std::vector<SchemaField>{
+      SchemaField::MakeRequired(1, "id", int32()),
+      SchemaField(2, "added", int32(), /*optional=*/false, /*doc=*/{},
+                  std::make_shared<const Literal>(Literal::Int(42)),
+                  std::make_shared<const Literal>(Literal::Int(7))),
+  });
+  DataWriterOptions options{
+      .path = param.path,
+      .schema = write_schema,
+      .input_schema = input_schema,
+      .spec = PartitionSpec::Unpartitioned(),
+      .partition = PartitionValues{},
+      .format = param.format,
+      .io = file_io_,
+      .properties = FormatProperties(param.format),
+  };
+
+  ICEBERG_UNWRAP_OR_FAIL(auto writer, DataWriter::Make(options));
+  auto input = CreateArray(*input_schema, R"([[1], [2]])");
+  ArrowArray arrow_array;
+  ASSERT_TRUE(::arrow::ExportArray(*input, &arrow_array).ok());
+  ASSERT_THAT(writer->Write(&arrow_array), IsOk());
+  ASSERT_THAT(writer->Close(), IsOk());
+
+  ICEBERG_UNWRAP_OR_FAIL(
+      auto reader, ReaderFactoryRegistry::Open(
+                       param.format,
+                       {.path = param.path, .io = file_io_, .projection = write_schema}));
+  ICEBERG_UNWRAP_OR_FAIL(auto batch, reader->Next());
+  ASSERT_TRUE(batch.has_value());
+
+  auto expected = CreateArray(*write_schema, R"([[1, 7], [2, 7]])");
+  auto actual = ::arrow::ImportArray(&batch.value(), expected->type()).ValueOrDie();
+
+  // This file is written after the column exists, so the missing input column uses
+  // write-default (7), not initial-default (42).
+  ASSERT_TRUE(actual->Equals(expected))
+      << "actual: " << actual->ToString() << "\nexpected: " << expected->ToString();
+}
+
 namespace {
 
 // Two-row column of `json` values at `type`, for the simple cases.
@@ -261,6 +347,16 @@ INSTANTIATE_TEST_SUITE_P(
       return info.param.name;
     });
 
+INSTANTIATE_TEST_SUITE_P(
+    Formats, WriteDefaultEndToEndTest,
+    ::testing::Values(WriteDefaultEndToEndParam{"parquet", FileFormatType::kParquet,
+                                                "write-default.parquet"},
+                      WriteDefaultEndToEndParam{"avro", FileFormatType::kAvro,
+                                                "write-default.avro"}),
+    [](const ::testing::TestParamInfo<WriteDefaultEndToEndParam>& info) {
+      return info.param.name;
+    });
+
 }  // namespace
 
 }  // namespace iceberg