src/parquet/file/writer-internal.cc - parquet-cpp - Git at Google

 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
 // regarding copyright ownership.  The ASF licenses this file
 // to you under the Apache License, Version 2.0 (the
 // "License"); you may not use this file except in compliance
 // with the License.  You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing,
 // software distributed under the License is distributed on an
 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 // KIND, either express or implied.  See the License for the
 // specific language governing permissions and limitations
 // under the License.

 #include "parquet/file/writer-internal.h"

 #include <cstdint>
 #include <memory>

 #include "arrow/util/compression.h"

 #include "parquet/column_writer.h"
 #include "parquet/schema-internal.h"
 #include "parquet/schema.h"
 #include "parquet/thrift.h"
 #include "parquet/util/memory.h"

 using arrow::MemoryPool;

 using parquet::schema::GroupNode;
 using parquet::schema::SchemaFlattener;

 namespace parquet {

 // FIXME: copied from reader-internal.cc
 static constexpr uint8_t PARQUET_MAGIC[4] = {'P', 'A', 'R', '1'};

 // ----------------------------------------------------------------------
 // SerializedPageWriter

 SerializedPageWriter::SerializedPageWriter(OutputStream* sink, Compression::type codec,
                                            ColumnChunkMetaDataBuilder* metadata,
                                            MemoryPool* pool)
     : sink_(sink),
       metadata_(metadata),
       pool_(pool),
       num_values_(0),
       dictionary_page_offset_(0),
       data_page_offset_(0),
       total_uncompressed_size_(0),
       total_compressed_size_(0) {
   compressor_ = GetCodecFromArrow(codec);
 }

 static format::Statistics ToThrift(const EncodedStatistics& row_group_statistics) {
   format::Statistics statistics;
   if (row_group_statistics.has_min) statistics.__set_min(row_group_statistics.min());
   if (row_group_statistics.has_max) statistics.__set_max(row_group_statistics.max());
   if (row_group_statistics.has_null_count)
     statistics.__set_null_count(row_group_statistics.null_count);
   if (row_group_statistics.has_distinct_count)
     statistics.__set_distinct_count(row_group_statistics.distinct_count);
   return statistics;
 }

 void SerializedPageWriter::Close(bool has_dictionary, bool fallback) {
   // index_page_offset = 0 since they are not supported
   metadata_->Finish(num_values_, dictionary_page_offset_, 0, data_page_offset_,
                     total_compressed_size_, total_uncompressed_size_, has_dictionary,
                     fallback);

   // Write metadata at end of column chunk
   metadata_->WriteTo(sink_);
 }

 void SerializedPageWriter::Compress(const Buffer& src_buffer,
                                     ResizableBuffer* dest_buffer) {
   DCHECK(compressor_ != nullptr);

   // Compress the data
   int64_t max_compressed_size =
       compressor_->MaxCompressedLen(src_buffer.size(), src_buffer.data());

   // Use Arrow::Buffer::shrink_to_fit = false
   // underlying buffer only keeps growing. Resize to a smaller size does not reallocate.
   PARQUET_THROW_NOT_OK(dest_buffer->Resize(max_compressed_size, false));

   int64_t compressed_size;
   PARQUET_THROW_NOT_OK(
       compressor_->Compress(src_buffer.size(), src_buffer.data(), max_compressed_size,
                             dest_buffer->mutable_data(), &compressed_size));
   PARQUET_THROW_NOT_OK(dest_buffer->Resize(compressed_size, false));
 }

 int64_t SerializedPageWriter::WriteDataPage(const CompressedDataPage& page) {
   int64_t uncompressed_size = page.uncompressed_size();
   std::shared_ptr<Buffer> compressed_data = page.buffer();

   format::DataPageHeader data_page_header;
   data_page_header.__set_num_values(page.num_values());
   data_page_header.__set_encoding(ToThrift(page.encoding()));
   data_page_header.__set_definition_level_encoding(
       ToThrift(page.definition_level_encoding()));
   data_page_header.__set_repetition_level_encoding(
       ToThrift(page.repetition_level_encoding()));
   data_page_header.__set_statistics(ToThrift(page.statistics()));

   format::PageHeader page_header;
   page_header.__set_type(format::PageType::DATA_PAGE);
   page_header.__set_uncompressed_page_size(static_cast<int32_t>(uncompressed_size));
   page_header.__set_compressed_page_size(static_cast<int32_t>(compressed_data->size()));
   page_header.__set_data_page_header(data_page_header);
   // TODO(PARQUET-594) crc checksum

   int64_t start_pos = sink_->Tell();
   if (data_page_offset_ == 0) {
     data_page_offset_ = start_pos;
   }

   int64_t header_size =
       SerializeThriftMsg(&page_header, sizeof(format::PageHeader), sink_);
   sink_->Write(compressed_data->data(), compressed_data->size());

   total_uncompressed_size_ += uncompressed_size + header_size;
   total_compressed_size_ += compressed_data->size() + header_size;
   num_values_ += page.num_values();

   return sink_->Tell() - start_pos;
 }

 int64_t SerializedPageWriter::WriteDictionaryPage(const DictionaryPage& page) {
   int64_t uncompressed_size = page.size();
   std::shared_ptr<Buffer> compressed_data = nullptr;
   if (has_compressor()) {
     auto buffer = std::static_pointer_cast<ResizableBuffer>(
         AllocateBuffer(pool_, uncompressed_size));
     Compress(*(page.buffer().get()), buffer.get());
     compressed_data = std::static_pointer_cast<Buffer>(buffer);
   } else {
     compressed_data = page.buffer();
   }

   format::DictionaryPageHeader dict_page_header;
   dict_page_header.__set_num_values(page.num_values());
   dict_page_header.__set_encoding(ToThrift(page.encoding()));
   dict_page_header.__set_is_sorted(page.is_sorted());

   format::PageHeader page_header;
   page_header.__set_type(format::PageType::DICTIONARY_PAGE);
   page_header.__set_uncompressed_page_size(static_cast<int32_t>(uncompressed_size));
   page_header.__set_compressed_page_size(static_cast<int32_t>(compressed_data->size()));
   page_header.__set_dictionary_page_header(dict_page_header);
   // TODO(PARQUET-594) crc checksum

   int64_t start_pos = sink_->Tell();
   if (dictionary_page_offset_ == 0) {
     dictionary_page_offset_ = start_pos;
   }
   int64_t header_size =
       SerializeThriftMsg(&page_header, sizeof(format::PageHeader), sink_);
   sink_->Write(compressed_data->data(), compressed_data->size());

   total_uncompressed_size_ += uncompressed_size + header_size;
   total_compressed_size_ += compressed_data->size() + header_size;

   return sink_->Tell() - start_pos;
 }

 // ----------------------------------------------------------------------
 // RowGroupSerializer

 int RowGroupSerializer::num_columns() const { return metadata_->num_columns(); }

 int64_t RowGroupSerializer::num_rows() const { return num_rows_; }

 ColumnWriter* RowGroupSerializer::NextColumn() {
   // Throws an error if more columns are being written
   auto col_meta = metadata_->NextColumnChunk();

   if (current_column_writer_) {
     total_bytes_written_ += current_column_writer_->Close();
   }

   const ColumnDescriptor* column_descr = col_meta->descr();
   std::unique_ptr<PageWriter> pager(
       new SerializedPageWriter(sink_, properties_->compression(column_descr->path()),
                                col_meta, properties_->memory_pool()));
   current_column_writer_ =
       ColumnWriter::Make(col_meta, std::move(pager), num_rows_, properties_);
   return current_column_writer_.get();
 }

 int RowGroupSerializer::current_column() const { return metadata_->current_column(); }

 void RowGroupSerializer::Close() {
   if (!closed_) {
     closed_ = true;

     if (current_column_writer_) {
       total_bytes_written_ += current_column_writer_->Close();
       current_column_writer_.reset();
     }
     // Ensures all columns have been written
     metadata_->Finish(total_bytes_written_);
   }
 }

 // ----------------------------------------------------------------------
 // FileSerializer

 std::unique_ptr<ParquetFileWriter::Contents> FileSerializer::Open(
     const std::shared_ptr<OutputStream>& sink, const std::shared_ptr<GroupNode>& schema,
     const std::shared_ptr<WriterProperties>& properties,
     const std::shared_ptr<const KeyValueMetadata>& key_value_metadata) {
   std::unique_ptr<ParquetFileWriter::Contents> result(
       new FileSerializer(sink, schema, properties, key_value_metadata));

   return result;
 }

 void FileSerializer::Close() {
   if (is_open_) {
     if (row_group_writer_) {
       row_group_writer_->Close();
     }
     row_group_writer_.reset();

     // Write magic bytes and metadata
     WriteMetaData();

     sink_->Close();
     is_open_ = false;
   }
 }

 int FileSerializer::num_columns() const { return schema_.num_columns(); }

 int FileSerializer::num_row_groups() const { return num_row_groups_; }

 int64_t FileSerializer::num_rows() const { return num_rows_; }

 const std::shared_ptr<WriterProperties>& FileSerializer::properties() const {
   return properties_;
 }

 RowGroupWriter* FileSerializer::AppendRowGroup(int64_t num_rows) {
   if (row_group_writer_) {
     row_group_writer_->Close();
   }
   num_rows_ += num_rows;
   num_row_groups_++;
   auto rg_metadata = metadata_->AppendRowGroup(num_rows);
   std::unique_ptr<RowGroupWriter::Contents> contents(
       new RowGroupSerializer(num_rows, sink_.get(), rg_metadata, properties_.get()));
   row_group_writer_.reset(new RowGroupWriter(std::move(contents)));
   return row_group_writer_.get();
 }

 FileSerializer::~FileSerializer() {
   try {
     Close();
   } catch (...) {
   }
 }

 void FileSerializer::WriteMetaData() {
   // Write MetaData
   uint32_t metadata_len = static_cast<uint32_t>(sink_->Tell());

   // Get a FileMetaData
   auto metadata = metadata_->Finish();
   metadata->WriteTo(sink_.get());
   metadata_len = static_cast<uint32_t>(sink_->Tell()) - metadata_len;

   // Write Footer
   sink_->Write(reinterpret_cast<uint8_t*>(&metadata_len), 4);
   sink_->Write(PARQUET_MAGIC, 4);
 }

 FileSerializer::FileSerializer(
     const std::shared_ptr<OutputStream>& sink, const std::shared_ptr<GroupNode>& schema,
     const std::shared_ptr<WriterProperties>& properties,
     const std::shared_ptr<const KeyValueMetadata>& key_value_metadata)
     : ParquetFileWriter::Contents(schema, key_value_metadata),
       sink_(sink),
       is_open_(true),
       properties_(properties),
       num_row_groups_(0),
       num_rows_(0),
       metadata_(FileMetaDataBuilder::Make(&schema_, properties, key_value_metadata)) {
   StartFile();
 }

 void FileSerializer::StartFile() {
   // Parquet files always start with PAR1
   sink_->Write(PARQUET_MAGIC, 4);
 }

 }  // namespace parquet
	// Licensed to the Apache Software Foundation (ASF) under one
	// or more contributor license agreements. See the NOTICE file
	// distributed with this work for additional information
	// regarding copyright ownership. The ASF licenses this file
	// to you under the Apache License, Version 2.0 (the
	// "License"); you may not use this file except in compliance
	// with the License. You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing,
	// software distributed under the License is distributed on an
	// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
	// KIND, either express or implied. See the License for the
	// specific language governing permissions and limitations
	// under the License.

	#include "parquet/file/writer-internal.h"

	#include <cstdint>
	#include <memory>

	#include "arrow/util/compression.h"

	#include "parquet/column_writer.h"
	#include "parquet/schema-internal.h"
	#include "parquet/schema.h"
	#include "parquet/thrift.h"
	#include "parquet/util/memory.h"

	using arrow::MemoryPool;

	using parquet::schema::GroupNode;
	using parquet::schema::SchemaFlattener;

	namespace parquet {

	// FIXME: copied from reader-internal.cc
	static constexpr uint8_t PARQUET_MAGIC[4] = {'P', 'A', 'R', '1'};

	// ----------------------------------------------------------------------
	// SerializedPageWriter

	SerializedPageWriter::SerializedPageWriter(OutputStream* sink, Compression::type codec,
	ColumnChunkMetaDataBuilder* metadata,
	MemoryPool* pool)
	: sink_(sink),
	metadata_(metadata),
	pool_(pool),
	num_values_(0),
	dictionary_page_offset_(0),
	data_page_offset_(0),
	total_uncompressed_size_(0),
	total_compressed_size_(0) {
	compressor_ = GetCodecFromArrow(codec);
	}

	static format::Statistics ToThrift(const EncodedStatistics& row_group_statistics) {
	format::Statistics statistics;
	if (row_group_statistics.has_min) statistics.__set_min(row_group_statistics.min());
	if (row_group_statistics.has_max) statistics.__set_max(row_group_statistics.max());
	if (row_group_statistics.has_null_count)
	statistics.__set_null_count(row_group_statistics.null_count);
	if (row_group_statistics.has_distinct_count)
	statistics.__set_distinct_count(row_group_statistics.distinct_count);
	return statistics;
	}

	void SerializedPageWriter::Close(bool has_dictionary, bool fallback) {
	// index_page_offset = 0 since they are not supported
	metadata_->Finish(num_values_, dictionary_page_offset_, 0, data_page_offset_,
	total_compressed_size_, total_uncompressed_size_, has_dictionary,
	fallback);

	// Write metadata at end of column chunk
	metadata_->WriteTo(sink_);
	}

	void SerializedPageWriter::Compress(const Buffer& src_buffer,
	ResizableBuffer* dest_buffer) {
	DCHECK(compressor_ != nullptr);

	// Compress the data
	int64_t max_compressed_size =
	compressor_->MaxCompressedLen(src_buffer.size(), src_buffer.data());

	// Use Arrow::Buffer::shrink_to_fit = false
	// underlying buffer only keeps growing. Resize to a smaller size does not reallocate.
	PARQUET_THROW_NOT_OK(dest_buffer->Resize(max_compressed_size, false));

	int64_t compressed_size;
	PARQUET_THROW_NOT_OK(
	compressor_->Compress(src_buffer.size(), src_buffer.data(), max_compressed_size,
	dest_buffer->mutable_data(), &compressed_size));
	PARQUET_THROW_NOT_OK(dest_buffer->Resize(compressed_size, false));
	}

	int64_t SerializedPageWriter::WriteDataPage(const CompressedDataPage& page) {
	int64_t uncompressed_size = page.uncompressed_size();
	std::shared_ptr<Buffer> compressed_data = page.buffer();

	format::DataPageHeader data_page_header;
	data_page_header.__set_num_values(page.num_values());
	data_page_header.__set_encoding(ToThrift(page.encoding()));
	data_page_header.__set_definition_level_encoding(
	ToThrift(page.definition_level_encoding()));
	data_page_header.__set_repetition_level_encoding(
	ToThrift(page.repetition_level_encoding()));
	data_page_header.__set_statistics(ToThrift(page.statistics()));

	format::PageHeader page_header;
	page_header.__set_type(format::PageType::DATA_PAGE);
	page_header.__set_uncompressed_page_size(static_cast<int32_t>(uncompressed_size));
	page_header.__set_compressed_page_size(static_cast<int32_t>(compressed_data->size()));
	page_header.__set_data_page_header(data_page_header);
	// TODO(PARQUET-594) crc checksum

	int64_t start_pos = sink_->Tell();
	if (data_page_offset_ == 0) {
	data_page_offset_ = start_pos;
	}

	int64_t header_size =
	SerializeThriftMsg(&page_header, sizeof(format::PageHeader), sink_);
	sink_->Write(compressed_data->data(), compressed_data->size());

	total_uncompressed_size_ += uncompressed_size + header_size;
	total_compressed_size_ += compressed_data->size() + header_size;
	num_values_ += page.num_values();

	return sink_->Tell() - start_pos;
	}

	int64_t SerializedPageWriter::WriteDictionaryPage(const DictionaryPage& page) {
	int64_t uncompressed_size = page.size();
	std::shared_ptr<Buffer> compressed_data = nullptr;
	if (has_compressor()) {
	auto buffer = std::static_pointer_cast<ResizableBuffer>(
	AllocateBuffer(pool_, uncompressed_size));
	Compress(*(page.buffer().get()), buffer.get());
	compressed_data = std::static_pointer_cast<Buffer>(buffer);
	} else {
	compressed_data = page.buffer();
	}

	format::DictionaryPageHeader dict_page_header;
	dict_page_header.__set_num_values(page.num_values());
	dict_page_header.__set_encoding(ToThrift(page.encoding()));
	dict_page_header.__set_is_sorted(page.is_sorted());

	format::PageHeader page_header;
	page_header.__set_type(format::PageType::DICTIONARY_PAGE);
	page_header.__set_uncompressed_page_size(static_cast<int32_t>(uncompressed_size));
	page_header.__set_compressed_page_size(static_cast<int32_t>(compressed_data->size()));
	page_header.__set_dictionary_page_header(dict_page_header);
	// TODO(PARQUET-594) crc checksum

	int64_t start_pos = sink_->Tell();
	if (dictionary_page_offset_ == 0) {
	dictionary_page_offset_ = start_pos;
	}
	int64_t header_size =
	SerializeThriftMsg(&page_header, sizeof(format::PageHeader), sink_);
	sink_->Write(compressed_data->data(), compressed_data->size());

	total_uncompressed_size_ += uncompressed_size + header_size;
	total_compressed_size_ += compressed_data->size() + header_size;

	return sink_->Tell() - start_pos;
	}

	// ----------------------------------------------------------------------
	// RowGroupSerializer

	int RowGroupSerializer::num_columns() const { return metadata_->num_columns(); }

	int64_t RowGroupSerializer::num_rows() const { return num_rows_; }

	ColumnWriter* RowGroupSerializer::NextColumn() {
	// Throws an error if more columns are being written
	auto col_meta = metadata_->NextColumnChunk();

	if (current_column_writer_) {
	total_bytes_written_ += current_column_writer_->Close();
	}

	const ColumnDescriptor* column_descr = col_meta->descr();
	std::unique_ptr<PageWriter> pager(
	new SerializedPageWriter(sink_, properties_->compression(column_descr->path()),
	col_meta, properties_->memory_pool()));
	current_column_writer_ =
	ColumnWriter::Make(col_meta, std::move(pager), num_rows_, properties_);
	return current_column_writer_.get();
	}

	int RowGroupSerializer::current_column() const { return metadata_->current_column(); }

	void RowGroupSerializer::Close() {
	if (!closed_) {
	closed_ = true;

	if (current_column_writer_) {
	total_bytes_written_ += current_column_writer_->Close();
	current_column_writer_.reset();
	}
	// Ensures all columns have been written
	metadata_->Finish(total_bytes_written_);
	}
	}

	// ----------------------------------------------------------------------
	// FileSerializer

	std::unique_ptr<ParquetFileWriter::Contents> FileSerializer::Open(
	const std::shared_ptr<OutputStream>& sink, const std::shared_ptr<GroupNode>& schema,
	const std::shared_ptr<WriterProperties>& properties,
	const std::shared_ptr<const KeyValueMetadata>& key_value_metadata) {
	std::unique_ptr<ParquetFileWriter::Contents> result(
	new FileSerializer(sink, schema, properties, key_value_metadata));

	return result;
	}

	void FileSerializer::Close() {
	if (is_open_) {
	if (row_group_writer_) {
	row_group_writer_->Close();
	}
	row_group_writer_.reset();

	// Write magic bytes and metadata
	WriteMetaData();

	sink_->Close();
	is_open_ = false;
	}
	}

	int FileSerializer::num_columns() const { return schema_.num_columns(); }

	int FileSerializer::num_row_groups() const { return num_row_groups_; }

	int64_t FileSerializer::num_rows() const { return num_rows_; }

	const std::shared_ptr<WriterProperties>& FileSerializer::properties() const {
	return properties_;
	}

	RowGroupWriter* FileSerializer::AppendRowGroup(int64_t num_rows) {
	if (row_group_writer_) {
	row_group_writer_->Close();
	}
	num_rows_ += num_rows;
	num_row_groups_++;
	auto rg_metadata = metadata_->AppendRowGroup(num_rows);
	std::unique_ptr<RowGroupWriter::Contents> contents(
	new RowGroupSerializer(num_rows, sink_.get(), rg_metadata, properties_.get()));
	row_group_writer_.reset(new RowGroupWriter(std::move(contents)));
	return row_group_writer_.get();
	}

	FileSerializer::~FileSerializer() {
	try {
	Close();
	} catch (...) {
	}
	}

	void FileSerializer::WriteMetaData() {
	// Write MetaData
	uint32_t metadata_len = static_cast<uint32_t>(sink_->Tell());

	// Get a FileMetaData
	auto metadata = metadata_->Finish();
	metadata->WriteTo(sink_.get());
	metadata_len = static_cast<uint32_t>(sink_->Tell()) - metadata_len;

	// Write Footer
	sink_->Write(reinterpret_cast<uint8_t*>(&metadata_len), 4);
	sink_->Write(PARQUET_MAGIC, 4);
	}

	FileSerializer::FileSerializer(
	const std::shared_ptr<OutputStream>& sink, const std::shared_ptr<GroupNode>& schema,
	const std::shared_ptr<WriterProperties>& properties,
	const std::shared_ptr<const KeyValueMetadata>& key_value_metadata)
	: ParquetFileWriter::Contents(schema, key_value_metadata),
	sink_(sink),
	is_open_(true),
	properties_(properties),
	num_row_groups_(0),
	num_rows_(0),
	metadata_(FileMetaDataBuilder::Make(&schema_, properties, key_value_metadata)) {
	StartFile();
	}

	void FileSerializer::StartFile() {
	// Parquet files always start with PAR1
	sink_->Write(PARQUET_MAGIC, 4);
	}

	} // namespace parquet