feat: apply write defaults in Parquet and Avro writers

src/iceberg/arrow/arrow_c_data_util.cc

@@ -22,20 +22,34 @@
 #include <mutex>
 #include <span>
 #include <utility>
+#include <variant>
 #include <vector>
 
+#include <arrow/array/array_base.h>
 #include <arrow/array/array_primitive.h>
+#include <arrow/array/data.h>
+#include <arrow/array/util.h>
 #include <arrow/buffer.h>
 #include <arrow/c/bridge.h>
 #include <arrow/compute/api_vector.h>
+#include <arrow/memory_pool.h>
 #include <arrow/record_batch.h>
+#include <arrow/type.h>
 #include <nanoarrow/nanoarrow.h>
 
 #include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow/literal_util_internal.h"
 #include "iceberg/arrow_c_data_guard_internal.h"
 #include "iceberg/arrow_c_data_util_internal.h"
 #include "iceberg/nanoarrow_status_internal.h"
 #include "iceberg/result.h"
+#include "iceberg/schema.h"
+#include "iceberg/schema_field.h"
+#include "iceberg/schema_internal.h"
+#include "iceberg/schema_util.h"
+#include "iceberg/type.h"
+#include "iceberg/util/checked_cast.h"
+#include "iceberg/util/formatter.h"
 #include "iceberg/util/macros.h"
 
 namespace iceberg {
@@ -128,4 +142,154 @@ void RegisterArrowProjectBatch() {
   });
 }
 
+namespace arrow {
+namespace {
+
+Result<std::shared_ptr<::arrow::Schema>> ImportIcebergSchema(const Schema& schema) {
+  ArrowSchema c_schema;
+  ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(schema, &c_schema));
+  internal::ArrowSchemaGuard c_schema_guard(&c_schema);
+  ICEBERG_ARROW_ASSIGN_OR_RETURN(auto arrow_schema, ::arrow::ImportSchema(&c_schema));
+  return arrow_schema;
+}
+
+// Writers must not silently drop data at any nesting level.
+Status RejectUnknownInputFields(const NestedType& input_type,
+                                const NestedType& write_type) {
+  for (const auto& input_field : input_type.fields()) {
+    ICEBERG_ASSIGN_OR_RAISE(auto write_field,
+                            write_type.GetFieldById(input_field.field_id()));
+    if (!write_field.has_value()) {
+      return InvalidSchema("Input field id {} is not present in write schema",
+                           input_field.field_id());
+    }
+    if (input_field.type()->is_nested() &&
+        input_field.type()->type_id() == write_field->get().type()->type_id()) {
+      ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(
+          internal::checked_cast<const NestedType&>(*input_field.type()),
+          internal::checked_cast<const NestedType&>(*write_field->get().type())));
+    }
+  }
+  return {};
+}
+
+Result<std::vector<std::shared_ptr<::arrow::Array>>> MaterializeProjection(
+    std::span<const FieldProjection> projections,
+    const std::vector<std::shared_ptr<::arrow::Array>>& input_arrays,
+    const std::vector<std::shared_ptr<::arrow::Field>>& write_arrow_fields,
+    int64_t length, ::arrow::MemoryPool* pool) {
+  std::vector<std::shared_ptr<::arrow::Array>> result;
+  result.reserve(projections.size());
+  for (size_t index = 0; index < projections.size(); ++index) {
+    const auto& projection = projections[index];
+    const auto& write_arrow_type = write_arrow_fields[index]->type();
+    switch (projection.kind) {
+      case FieldProjection::Kind::kProjected: {
+        const auto& input_array = input_arrays[std::get<size_t>(projection.from)];
+        if (projection.children.empty()) {
+          result.emplace_back(input_array);
+          break;
+        }
+        auto data = input_array->data()->Copy();
+        data->type = write_arrow_type;
+        auto child_data = data;
+        auto child_type = write_arrow_type;
+        int64_t child_length = data->offset + data->length;
+        switch (write_arrow_type->id()) {
+          case ::arrow::Type::STRUCT:
+            break;
+          case ::arrow::Type::LIST:
+            child_length = data->child_data[0]->length;
+            break;
+          case ::arrow::Type::MAP:
+            // Arrow maps have an entries struct between the map and its key/value.
+            child_data = data->child_data[0]->Copy();
+            child_type = write_arrow_type->field(0)->type();
+            child_length = child_data->offset + child_data->length;
+            break;
+          default:
+            return InvalidArrowData("Expected nested array for projected field {}",
+                                    index);
+        }
+        std::vector<std::shared_ptr<::arrow::Array>> input_children;
+        input_children.reserve(child_data->child_data.size());
+        for (const auto& child : child_data->child_data) {
+          input_children.emplace_back(::arrow::MakeArray(child));
+        }
+        // Raw child data retains the parent's physical row positions.
+        // StructArray::fields() would apply the parent slice, causing its offset to be
+        // applied twice.
+        ICEBERG_ASSIGN_OR_RAISE(
+            auto children,
+            MaterializeProjection(projection.children, input_children,
+                                  child_type->fields(), child_length, pool));
+        child_data->type = child_type;
+        child_data->child_data.clear();
+        for (const auto& child : children) {
+          child_data->child_data.emplace_back(child->data());
+        }
+        if (write_arrow_type->id() == ::arrow::Type::MAP) {
+          data->child_data = {std::move(child_data)};
+        }
+        result.emplace_back(::arrow::MakeArray(std::move(data)));
+        break;
+      }
+      case FieldProjection::Kind::kDefault: {
+        ICEBERG_ASSIGN_OR_RAISE(auto default_array,
+                                MakeDefaultArray(std::get<Literal>(projection.from),
+                                                 write_arrow_type, length, pool));
+        result.emplace_back(std::move(default_array));
+        break;
+      }
+      case FieldProjection::Kind::kNull: {
+        ICEBERG_ARROW_ASSIGN_OR_RETURN(
+            auto null_array, ::arrow::MakeArrayOfNull(write_arrow_type, length, pool));
+        result.emplace_back(std::move(null_array));
+        break;
+      }
+      default:
+        return NotSupported("Projection kind {} is not supported for write alignment",
+                            ToString(projection.kind));
+    }
+  }
+  return result;
+}
+
+}  // namespace
+
+Result<ArrowArray> AlignBatchForWrite(ArrowArray* input_batch, const Schema& input_schema,
+                                      const Schema& write_schema) {
+  ICEBERG_PRECHECK(input_batch != nullptr, "input_batch must not be null");
+  internal::ArrowArrayGuard input_guard(input_batch);
+  if (input_schema == write_schema) {
+    return std::exchange(*input_batch, ArrowArray{});
+  }
+
+  ICEBERG_ASSIGN_OR_RAISE(auto input_arrow_schema, ImportIcebergSchema(input_schema));
+  ICEBERG_ARROW_ASSIGN_OR_RETURN(
+      auto input_record_batch,
+      ::arrow::ImportRecordBatch(input_batch, input_arrow_schema));
+
+  ICEBERG_ASSIGN_OR_RAISE(auto write_arrow_schema, ImportIcebergSchema(write_schema));
+
+  ICEBERG_RETURN_UNEXPECTED(RejectUnknownInputFields(input_schema, write_schema));
+  ICEBERG_ASSIGN_OR_RAISE(auto projection,
+                          Project(write_schema, input_schema, /*prune_source=*/false,
+                                  ProjectionMode::kWrite));
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto columns,
+      MaterializeProjection(projection.fields, input_record_batch->columns(),
+                            write_arrow_schema->fields(), input_record_batch->num_rows(),
+                            ::arrow::default_memory_pool()));
+  auto output_record_batch = ::arrow::RecordBatch::Make(
+      write_arrow_schema, input_record_batch->num_rows(), std::move(columns));
+
+  ArrowArray output{};
+  ICEBERG_ARROW_RETURN_NOT_OK(::arrow::ExportRecordBatch(*output_record_batch, &output));
+  internal::ArrowArrayGuard output_guard(&output);
+  return std::exchange(output, ArrowArray{});
+}
+
+}  // namespace arrow
+
 }  // namespace iceberg

src/iceberg/arrow_c_data_util_internal.h

@@ -31,6 +31,7 @@
 #include <vector>
 
 #include "iceberg/arrow_c_data.h"
+#include "iceberg/iceberg_bundle_export.h"
 #include "iceberg/iceberg_export.h"
 #include "iceberg/result.h"
 #include "iceberg/type_fwd.h"
@@ -242,4 +243,20 @@ ICEBERG_EXPORT Result<ArrowArray> ProjectBatch(ArrowArray* input_batch,
                                                std::span<const int32_t> row_indices,
                                                ProjectionContext& projection);
 
+namespace arrow {
+
+/// \brief Align an Arrow batch to an Iceberg write schema.
+///
+/// Fields are matched by Iceberg field id. Missing fields are materialized from their
+/// `write-default`, or as null when optional. Missing required fields without a
+/// `write-default` are rejected. Existing values, including explicit nulls, are
+/// preserved. Structs, lists, and maps are aligned recursively.
+///
+/// `input_batch` is consumed even when alignment fails.
+ICEBERG_BUNDLE_EXPORT Result<ArrowArray> AlignBatchForWrite(ArrowArray* input_batch,
+                                                            const Schema& input_schema,
+                                                            const Schema& write_schema);
+
+}  // namespace arrow
+
 }  // namespace iceberg

src/iceberg/avro/avro_writer.cc

@@ -32,6 +32,8 @@
 
 #include "iceberg/arrow/arrow_io_internal.h"
 #include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow_c_data_guard_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"
@@ -181,6 +183,7 @@ class AvroWriter::Impl {
  public:
   Status Open(const WriterOptions& options) {
     write_schema_ = options.schema;
+    input_schema_ = options.input_schema ? options.input_schema : write_schema_;
 
     ::avro::NodePtr root;
     ICEBERG_RETURN_UNEXPECTED(ToAvroNodeVisitor{}.Visit(*write_schema_, &root));
@@ -229,6 +232,11 @@ class AvroWriter::Impl {
   }
 
   Status Write(ArrowArray* data) {
+    if (*input_schema_ != *write_schema_) {
+      ICEBERG_ASSIGN_OR_RAISE(
+          *data, arrow::AlignBatchForWrite(data, *input_schema_, *write_schema_));
+    }
+    internal::ArrowArrayGuard data_guard(data);
     ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch,
                                    ::arrow::ImportRecordBatch(data, arrow_schema_));
 
@@ -272,6 +280,7 @@ class AvroWriter::Impl {
  private:
   // The schema to write.
   std::shared_ptr<::iceberg::Schema> write_schema_;
+  std::shared_ptr<::iceberg::Schema> input_schema_;
   // The avro schema to write.
   std::shared_ptr<::avro::ValidSchema> avro_schema_;
   // Arrow output stream of the Avro file to write

src/iceberg/data/data_writer.cc

@@ -38,6 +38,7 @@ class DataWriter::Impl {
         .schema = options.schema,
         .io = options.io,
         .properties = WriterProperties::FromMap(options.properties),
+        .input_schema = options.input_schema,
     };
 
     ICEBERG_ASSIGN_OR_RAISE(auto writer,

src/iceberg/data/data_writer.h

@@ -48,6 +48,8 @@ struct ICEBERG_DATA_EXPORT DataWriterOptions {
   std::shared_ptr<FileIO> io;
   std::optional<int32_t> sort_order_id;
   std::unordered_map<std::string, std::string> properties;
+  /// \brief Schema of input batches; when omitted, batches use the target schema.
+  std::shared_ptr<Schema> input_schema;
 };
 
 /// \brief Writer for Iceberg data files.

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;
@@ -85,6 +85,8 @@ struct ICEBERG_EXPORT WriterOptions {
   std::shared_ptr<MetricsConfig> metrics_config = MetricsConfig::Default();
   /// \brief Format-specific or implementation-specific properties.
   WriterProperties properties;
+  /// \brief Schema of input batches; when omitted, batches use the target schema.
+  std::shared_ptr<Schema> input_schema;
 };
 
 /// \brief Base writer class to write data from different file formats.

src/iceberg/parquet/parquet_writer.cc

@@ -39,8 +39,11 @@
 
 #include "iceberg/arrow/arrow_io_internal.h"
 #include "iceberg/arrow/arrow_status_internal.h"
+#include "iceberg/arrow_c_data_guard_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.h"
 #include "iceberg/schema_internal.h"
 #include "iceberg/type.h"
 #include "iceberg/util/macros.h"
@@ -239,6 +242,7 @@ class ParquetWriter::Impl {
  public:
   Status Open(const WriterOptions& options) {
     schema_ = options.schema;
+    input_schema_ = options.input_schema ? options.input_schema : schema_;
 
     ICEBERG_ASSIGN_OR_RAISE(auto compression, ParseCompression(options.properties));
     ICEBERG_ASSIGN_OR_RAISE(auto compression_level, ParseCodecLevel(options.properties));
@@ -281,6 +285,11 @@ class ParquetWriter::Impl {
   }
 
   Status Write(ArrowArray* array) {
+    if (*input_schema_ != *schema_) {
+      ICEBERG_ASSIGN_OR_RAISE(*array,
+                              arrow::AlignBatchForWrite(array, *input_schema_, *schema_));
+    }
+    internal::ArrowArrayGuard array_guard(array);
     ICEBERG_ARROW_ASSIGN_OR_RETURN(auto batch,
                                    ::arrow::ImportRecordBatch(array, arrow_schema_));
 
@@ -342,6 +351,7 @@ class ParquetWriter::Impl {
   ::arrow::MemoryPool* pool_ = ::arrow::default_memory_pool();
   // Schema to write from the Iceberg table.
   std::shared_ptr<Schema> schema_;
+  std::shared_ptr<Schema> input_schema_;
   // Schema to write from the Parquet file.
   std::shared_ptr<::arrow::Schema> arrow_schema_;
   // Parquet schema descriptor generated from the Arrow schema.