-
Notifications
You must be signed in to change notification settings - Fork 118
feat: implement DataWriter for Iceberg data files #552
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -19,20 +19,118 @@ | |||||
|
|
||||||
| #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: | ||||||
| static Result<std::unique_ptr<Impl>> Make(DataWriterOptions options) { | ||||||
| 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(auto writer, | ||||||
| WriterFactoryRegistry::Open(options.format, writer_options)); | ||||||
|
|
||||||
| return std::unique_ptr<Impl>(new Impl(std::move(options), std::move(writer))); | ||||||
| } | ||||||
|
|
||||||
| Status Write(ArrowArray* data) { | ||||||
| ICEBERG_PRECHECK(writer_, "Writer not initialized"); | ||||||
|
wgtmac marked this conversation as resolved.
Outdated
|
||||||
| return writer_->Write(data); | ||||||
| } | ||||||
|
|
||||||
| Result<int64_t> Length() const { | ||||||
| ICEBERG_PRECHECK(writer_, "Writer not initialized"); | ||||||
| return writer_->length(); | ||||||
| } | ||||||
|
|
||||||
| Status Close() { | ||||||
| ICEBERG_PRECHECK(writer_, "Writer not initialized"); | ||||||
| if (closed_) { | ||||||
| // Idempotent: no-op if already closed | ||||||
| return {}; | ||||||
| } | ||||||
|
Comment on lines
+58
to
+61
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I could see a case for making close idempotent, is there any strong reason why we want to return this error instead of no op for example?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed |
||||||
| ICEBERG_RETURN_UNEXPECTED(writer_->Close()); | ||||||
| closed_ = true; | ||||||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should this class address thread safety?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Good question! I've added explicit documentation that this class is not thread-safe:
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't think a single writer (or reader) should support thread safety so it is fine not to add comment like this.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @wgtmac out of curiosity for my own knowledge, what guarantees that a single writer/reader will be using the class?
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. These file writers are supposed to be used by a single write task, which for example can be an unit of a table sink operator in a sql job plan. Usually the writer is responsible for partitioned (and sometimes sorted) data chunks.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Agreed. Removed the thread safety comment from the header. |
||||||
| return {}; | ||||||
| } | ||||||
|
|
||||||
| Result<FileWriter::WriteResult> Metadata() { | ||||||
| ICEBERG_PRECHECK(closed_, "Cannot get metadata before closing the writer"); | ||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
We should return
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done |
||||||
|
|
||||||
| ICEBERG_ASSIGN_OR_RAISE(auto metrics, writer_->metrics()); | ||||||
| ICEBERG_ASSIGN_OR_RAISE(auto length, writer_->length()); | ||||||
| auto split_offsets = writer_->split_offsets(); | ||||||
|
|
||||||
| auto data_file = std::make_shared<DataFile>(); | ||||||
| data_file->content = DataFile::Content::kData; | ||||||
|
wgtmac marked this conversation as resolved.
Outdated
|
||||||
| data_file->file_path = options_.path; | ||||||
| data_file->file_format = options_.format; | ||||||
| data_file->partition = options_.partition; | ||||||
| data_file->record_count = metrics.row_count.value_or(0); | ||||||
|
wgtmac marked this conversation as resolved.
Outdated
|
||||||
| data_file->file_size_in_bytes = length; | ||||||
| data_file->sort_order_id = options_.sort_order_id; | ||||||
| data_file->split_offsets = std::move(split_offsets); | ||||||
|
|
||||||
| // Convert metrics maps from unordered_map to map | ||||||
| for (const auto& [col_id, size] : metrics.column_sizes) { | ||||||
| data_file->column_sizes[col_id] = size; | ||||||
|
wgtmac marked this conversation as resolved.
Outdated
|
||||||
| } | ||||||
| for (const auto& [col_id, count] : metrics.value_counts) { | ||||||
| data_file->value_counts[col_id] = count; | ||||||
| } | ||||||
| for (const auto& [col_id, count] : metrics.null_value_counts) { | ||||||
| data_file->null_value_counts[col_id] = count; | ||||||
| } | ||||||
| for (const auto& [col_id, count] : metrics.nan_value_counts) { | ||||||
| data_file->nan_value_counts[col_id] = count; | ||||||
| } | ||||||
|
|
||||||
| // Serialize literal bounds to binary format | ||||||
| for (const auto& [col_id, literal] : metrics.lower_bounds) { | ||||||
| ICEBERG_ASSIGN_OR_RAISE(auto serialized, literal.Serialize()); | ||||||
| data_file->lower_bounds[col_id] = std::move(serialized); | ||||||
| } | ||||||
| for (const auto& [col_id, literal] : metrics.upper_bounds) { | ||||||
| ICEBERG_ASSIGN_OR_RAISE(auto serialized, literal.Serialize()); | ||||||
| data_file->upper_bounds[col_id] = std::move(serialized); | ||||||
| } | ||||||
|
|
||||||
| FileWriter::WriteResult result; | ||||||
| result.data_files.push_back(std::move(data_file)); | ||||||
| return result; | ||||||
| } | ||||||
|
|
||||||
| private: | ||||||
| Impl(DataWriterOptions options, std::unique_ptr<Writer> writer) | ||||||
| : options_(std::move(options)), writer_(std::move(writer)) {} | ||||||
|
|
||||||
| DataWriterOptions options_; | ||||||
| std::unique_ptr<Writer> writer_; | ||||||
| bool closed_ = false; | ||||||
| }; | ||||||
|
|
||||||
| DataWriter::DataWriter(std::unique_ptr<Impl> impl) : impl_(std::move(impl)) {} | ||||||
|
|
||||||
| DataWriter::~DataWriter() = default; | ||||||
|
|
||||||
| Status DataWriter::Write(ArrowArray* data) { return NotImplemented(""); } | ||||||
| Result<std::unique_ptr<DataWriter>> DataWriter::Make(const DataWriterOptions& options) { | ||||||
| ICEBERG_ASSIGN_OR_RAISE(auto impl, Impl::Make(options)); | ||||||
| return std::unique_ptr<DataWriter>(new DataWriter(std::move(impl))); | ||||||
| } | ||||||
|
|
||||||
| Status DataWriter::Write(ArrowArray* data) { return impl_->Write(data); } | ||||||
|
|
||||||
| Result<int64_t> DataWriter::Length() const { return NotImplemented(""); } | ||||||
| Result<int64_t> DataWriter::Length() const { return impl_->Length(); } | ||||||
|
|
||||||
| Status DataWriter::Close() { return NotImplemented(""); } | ||||||
| Status DataWriter::Close() { return impl_->Close(); } | ||||||
|
|
||||||
| Result<FileWriter::WriteResult> DataWriter::Metadata() { return NotImplemented(""); } | ||||||
| Result<FileWriter::WriteResult> DataWriter::Metadata() { return impl_->Metadata(); } | ||||||
|
|
||||||
| } // namespace iceberg | ||||||
Uh oh!
There was an error while loading. Please reload this page.