src/parquet/file/writer-internal.h

// 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.

#ifndef PARQUET_FILE_WRITER_INTERNAL_H
#define PARQUET_FILE_WRITER_INTERNAL_H

#include <memory>
#include <vector>

#include "parquet/column_page.h"
#include "parquet/file/metadata.h"
#include "parquet/file/writer.h"
#include "parquet/parquet_types.h"
#include "parquet/util/memory.h"

namespace arrow {

class Codec;
};

namespace parquet {

// This subclass delimits pages appearing in a serialized stream, each preceded
// by a serialized Thrift format::PageHeader indicating the type of each page
// and the page metadata.
class SerializedPageWriter : public PageWriter {
 public:
  SerializedPageWriter(OutputStream* sink, Compression::type codec,
                       ColumnChunkMetaDataBuilder* metadata,
                       ::arrow::MemoryPool* pool = ::arrow::default_memory_pool());

  virtual ~SerializedPageWriter() {}

  int64_t WriteDataPage(const CompressedDataPage& page) override;

  int64_t WriteDictionaryPage(const DictionaryPage& page) override;

  /**
   * Compress a buffer.
   */
  void Compress(const Buffer& src_buffer, ResizableBuffer* dest_buffer) override;

  bool has_compressor() override { return (compressor_ != nullptr); }

  void Close(bool has_dictionary, bool fallback) override;

 private:
  OutputStream* sink_;
  ColumnChunkMetaDataBuilder* metadata_;
  ::arrow::MemoryPool* pool_;
  int64_t num_values_;
  int64_t dictionary_page_offset_;
  int64_t data_page_offset_;
  int64_t total_uncompressed_size_;
  int64_t total_compressed_size_;

  // Compression codec to use.
  std::unique_ptr<::arrow::Codec> compressor_;
};

// RowGroupWriter::Contents implementation for the Parquet file specification
class RowGroupSerializer : public RowGroupWriter::Contents {
 public:
  RowGroupSerializer(OutputStream* sink, RowGroupMetaDataBuilder* metadata,
                     const WriterProperties* properties)
      : sink_(sink),
        metadata_(metadata),
        properties_(properties),
        total_bytes_written_(0),
        closed_(false),
        current_column_index_(0),
        num_rows_(-1) {}

  int num_columns() const override;
  int64_t num_rows() const override;

  ColumnWriter* NextColumn() override;
  int current_column() const override;
  void Close() override;

 private:
  OutputStream* sink_;
  mutable RowGroupMetaDataBuilder* metadata_;
  const WriterProperties* properties_;
  int64_t total_bytes_written_;
  bool closed_;
  int current_column_index_;
  mutable int64_t num_rows_;

  void CheckRowsWritten() const;

  std::shared_ptr<ColumnWriter> current_column_writer_;
};

// An implementation of ParquetFileWriter::Contents that deals with the Parquet
// file structure, Thrift serialization, and other internal matters

class FileSerializer : public ParquetFileWriter::Contents {
 public:
  static std::unique_ptr<ParquetFileWriter::Contents> Open(
      const std::shared_ptr<OutputStream>& sink,
      const std::shared_ptr<schema::GroupNode>& schema,
      const std::shared_ptr<WriterProperties>& properties = default_writer_properties(),
      const std::shared_ptr<const KeyValueMetadata>& key_value_metadata = nullptr);

  void Close() override;

  RowGroupWriter* AppendRowGroup() override;

  const std::shared_ptr<WriterProperties>& properties() const override;

  int num_columns() const override;
  int num_row_groups() const override;
  int64_t num_rows() const override;

  virtual ~FileSerializer();

 private:
  explicit FileSerializer(
      const std::shared_ptr<OutputStream>& sink,
      const std::shared_ptr<schema::GroupNode>& schema,
      const std::shared_ptr<WriterProperties>& properties,
      const std::shared_ptr<const KeyValueMetadata>& key_value_metadata);

  std::shared_ptr<OutputStream> sink_;
  bool is_open_;
  const std::shared_ptr<WriterProperties> properties_;
  int num_row_groups_;
  int64_t num_rows_;
  std::unique_ptr<FileMetaDataBuilder> metadata_;
  std::unique_ptr<RowGroupWriter> row_group_writer_;

  void StartFile();
  void WriteMetaData();
};

}  // namespace parquet

#endif  // PARQUET_FILE_WRITER_INTERNAL_H