parquet.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2020-2024, NVIDIA CORPORATION.
 *
 * Licensed 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.
 */
 
#pragma once
 
#include <cudf/ast/expressions.hpp>
#include <cudf/io/detail/parquet.hpp>
#include <cudf/io/types.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
 
#include <rmm/mr/device/per_device_resource.hpp>
 
#include <iostream>
#include <memory>
#include <optional>
#include <string>
#include <vector>
 
namespace cudf::io {
constexpr size_t default_row_group_size_bytes   = 128 * 1024 * 1024;  
constexpr size_type default_row_group_size_rows = 1000000;     
constexpr size_t default_max_page_size_bytes    = 512 * 1024;  
constexpr size_type default_max_page_size_rows  = 20000;       
constexpr int32_t default_column_index_truncate_length = 64;   
constexpr size_t default_max_dictionary_size           = 1024 * 1024;  
constexpr size_type default_max_page_fragment_size     = 5000;  
 
class parquet_reader_options_builder;
 
class parquet_reader_options {
  source_info _source;
 
  // Path in schema of column to read; `nullopt` is all
  std::optional<std::vector<std::string>> _columns;
 
  // List of individual row groups to read (ignored if empty)
  std::vector<std::vector<size_type>> _row_groups;
  // Number of rows to skip from the start; Parquet stores the number of rows as int64_t
  int64_t _skip_rows = 0;
  // Number of rows to read; `nullopt` is all
  std::optional<size_type> _num_rows;
 
  // Predicate filter as AST to filter output rows.
  std::optional<std::reference_wrapper<ast::expression const>> _filter;
 
  // Whether to store string data as categorical type
  bool _convert_strings_to_categories = false;
  // Whether to use PANDAS metadata to load columns
  bool _use_pandas_metadata = true;
  // Cast timestamp columns to a specific type
  data_type _timestamp_type{type_id::EMPTY};
 
  std::optional<std::vector<reader_column_schema>> _reader_column_schema;
 
  explicit parquet_reader_options(source_info src) : _source{std::move(src)} {}
 
  friend parquet_reader_options_builder;
 
 public:
  explicit parquet_reader_options() = default;
 
  static parquet_reader_options_builder builder(source_info src);
 
  [[nodiscard]] source_info const& get_source() const { return _source; }
 
  [[nodiscard]] bool is_enabled_convert_strings_to_categories() const
  {
    return _convert_strings_to_categories;
  }
 
  [[nodiscard]] bool is_enabled_use_pandas_metadata() const { return _use_pandas_metadata; }
 
  [[nodiscard]] std::optional<std::vector<reader_column_schema>> get_column_schema() const
  {
    return _reader_column_schema;
  }
 
  [[nodiscard]] int64_t get_skip_rows() const { return _skip_rows; }
 
  [[nodiscard]] std::optional<size_type> const& get_num_rows() const { return _num_rows; }
 
  [[nodiscard]] auto const& get_columns() const { return _columns; }
 
  [[nodiscard]] auto const& get_row_groups() const { return _row_groups; }
 
  [[nodiscard]] auto const& get_filter() const { return _filter; }
 
  data_type get_timestamp_type() const { return _timestamp_type; }
 
  void set_columns(std::vector<std::string> col_names) { _columns = std::move(col_names); }
 
  void set_row_groups(std::vector<std::vector<size_type>> row_groups);
 
  void set_filter(ast::expression const& filter) { _filter = filter; }
 
  void enable_convert_strings_to_categories(bool val) { _convert_strings_to_categories = val; }
 
  void enable_use_pandas_metadata(bool val) { _use_pandas_metadata = val; }
 
  void set_column_schema(std::vector<reader_column_schema> val)
  {
    _reader_column_schema = std::move(val);
  }
 
  void set_skip_rows(int64_t val);
 
  void set_num_rows(size_type val);
 
  void set_timestamp_type(data_type type) { _timestamp_type = type; }
};
 
class parquet_reader_options_builder {
  parquet_reader_options options;
 
 public:
  parquet_reader_options_builder() = default;
 
  explicit parquet_reader_options_builder(source_info src) : options{std::move(src)} {}
 
  parquet_reader_options_builder& columns(std::vector<std::string> col_names)
  {
    options._columns = std::move(col_names);
    return *this;
  }
 
  parquet_reader_options_builder& row_groups(std::vector<std::vector<size_type>> row_groups)
  {
    options.set_row_groups(std::move(row_groups));
    return *this;
  }
 
  parquet_reader_options_builder& filter(ast::expression const& filter)
  {
    options.set_filter(filter);
    return *this;
  }
 
  parquet_reader_options_builder& convert_strings_to_categories(bool val)
  {
    options._convert_strings_to_categories = val;
    return *this;
  }
 
  parquet_reader_options_builder& use_pandas_metadata(bool val)
  {
    options._use_pandas_metadata = val;
    return *this;
  }
 
  parquet_reader_options_builder& set_column_schema(std::vector<reader_column_schema> val)
  {
    options._reader_column_schema = std::move(val);
    return *this;
  }
 
  parquet_reader_options_builder& skip_rows(int64_t val)
  {
    options.set_skip_rows(val);
    return *this;
  }
 
  parquet_reader_options_builder& num_rows(size_type val)
  {
    options.set_num_rows(val);
    return *this;
  }
 
  parquet_reader_options_builder& timestamp_type(data_type type)
  {
    options._timestamp_type = type;
    return *this;
  }
 
  operator parquet_reader_options&&() { return std::move(options); }
 
  parquet_reader_options&& build() { return std::move(options); }
};
 
table_with_metadata read_parquet(
  parquet_reader_options const& options,
  rmm::cuda_stream_view stream        = cudf::get_default_stream(),
  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
class chunked_parquet_reader {
 public:
  chunked_parquet_reader() = default;
 
  chunked_parquet_reader(
    std::size_t chunk_read_limit,
    parquet_reader_options const& options,
    rmm::cuda_stream_view stream        = cudf::get_default_stream(),
    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
  chunked_parquet_reader(
    std::size_t chunk_read_limit,
    std::size_t pass_read_limit,
    parquet_reader_options const& options,
    rmm::cuda_stream_view stream        = cudf::get_default_stream(),
    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
  ~chunked_parquet_reader();
 
  [[nodiscard]] bool has_next() const;
 
  [[nodiscard]] table_with_metadata read_chunk() const;
 
 private:
  std::unique_ptr<cudf::io::parquet::detail::chunked_reader> reader;
};
  // end of group
class parquet_writer_options_builder;
 
class parquet_writer_options {
  // Specify the sink to use for writer output
  sink_info _sink;
  // Specify the compression format to use
  compression_type _compression = compression_type::SNAPPY;
  // Specify the level of statistics in the output file
  statistics_freq _stats_level = statistics_freq::STATISTICS_ROWGROUP;
  // Sets of columns to output
  table_view _table;
  // Partitions described as {start_row, num_rows} pairs
  std::vector<partition_info> _partitions;
  // Optional associated metadata
  std::optional<table_input_metadata> _metadata;
  // Optional footer key_value_metadata
  std::vector<std::map<std::string, std::string>> _user_data;
  // Parquet writer can write INT96 or TIMESTAMP_MICROS. Defaults to TIMESTAMP_MICROS.
  // If true then overrides any per-column setting in _metadata.
  bool _write_timestamps_as_int96 = false;
  // Parquet writer can write timestamps as UTC
  // Defaults to true because libcudf timestamps are implicitly UTC
  bool _write_timestamps_as_UTC = true;
  // Column chunks file paths to be set in the raw output metadata. One per output file
  std::vector<std::string> _column_chunks_file_paths;
  // Maximum size of each row group (unless smaller than a single page)
  size_t _row_group_size_bytes = default_row_group_size_bytes;
  // Maximum number of rows in row group (unless smaller than a single page)
  size_type _row_group_size_rows = default_row_group_size_rows;
  // Maximum size of each page (uncompressed)
  size_t _max_page_size_bytes = default_max_page_size_bytes;
  // Maximum number of rows in a page
  size_type _max_page_size_rows = default_max_page_size_rows;
  // Maximum size of min or max values in column index
  int32_t _column_index_truncate_length = default_column_index_truncate_length;
  // When to use dictionary encoding for data
  dictionary_policy _dictionary_policy = dictionary_policy::ALWAYS;
  // Maximum size of column chunk dictionary (in bytes)
  size_t _max_dictionary_size = default_max_dictionary_size;
  // Maximum number of rows in a page fragment
  std::optional<size_type> _max_page_fragment_size;
  // Optional compression statistics
  std::shared_ptr<writer_compression_statistics> _compression_stats;
  // write V2 page headers?
  bool _v2_page_headers = false;
 
  explicit parquet_writer_options(sink_info const& sink, table_view const& table)
    : _sink(sink), _table(table)
  {
  }
 
  friend parquet_writer_options_builder;
 
 public:
  parquet_writer_options() = default;
 
  static parquet_writer_options_builder builder(sink_info const& sink, table_view const& table);
 
  static parquet_writer_options_builder builder();
 
  [[nodiscard]] sink_info const& get_sink() const { return _sink; }
 
  [[nodiscard]] compression_type get_compression() const { return _compression; }
 
  [[nodiscard]] statistics_freq get_stats_level() const { return _stats_level; }
 
  [[nodiscard]] table_view get_table() const { return _table; }
 
  [[nodiscard]] std::vector<partition_info> const& get_partitions() const { return _partitions; }
 
  [[nodiscard]] auto const& get_metadata() const { return _metadata; }
 
  std::vector<std::map<std::string, std::string>> const& get_key_value_metadata() const
  {
    return _user_data;
  }
 
  bool is_enabled_int96_timestamps() const { return _write_timestamps_as_int96; }
 
  [[nodiscard]] auto is_enabled_utc_timestamps() const { return _write_timestamps_as_UTC; }
 
  std::vector<std::string> const& get_column_chunks_file_paths() const
  {
    return _column_chunks_file_paths;
  }
 
  auto get_row_group_size_bytes() const { return _row_group_size_bytes; }
 
  auto get_row_group_size_rows() const { return _row_group_size_rows; }
 
  auto get_max_page_size_bytes() const
  {
    return std::min(_max_page_size_bytes, get_row_group_size_bytes());
  }
 
  auto get_max_page_size_rows() const
  {
    return std::min(_max_page_size_rows, get_row_group_size_rows());
  }
 
  auto get_column_index_truncate_length() const { return _column_index_truncate_length; }
 
  [[nodiscard]] dictionary_policy get_dictionary_policy() const { return _dictionary_policy; }
 
  [[nodiscard]] auto get_max_dictionary_size() const { return _max_dictionary_size; }
 
  [[nodiscard]] auto get_max_page_fragment_size() const { return _max_page_fragment_size; }
 
  [[nodiscard]] std::shared_ptr<writer_compression_statistics> get_compression_statistics() const
  {
    return _compression_stats;
  }
 
  [[nodiscard]] auto is_enabled_write_v2_headers() const { return _v2_page_headers; }
 
  void set_partitions(std::vector<partition_info> partitions);
 
  void set_metadata(table_input_metadata metadata) { _metadata = std::move(metadata); }
 
  void set_key_value_metadata(std::vector<std::map<std::string, std::string>> metadata);
 
  void set_stats_level(statistics_freq sf) { _stats_level = sf; }
 
  void set_compression(compression_type compression) { _compression = compression; }
 
  void enable_int96_timestamps(bool req) { _write_timestamps_as_int96 = req; }
 
  void enable_utc_timestamps(bool val) { _write_timestamps_as_UTC = val; }
 
  void set_column_chunks_file_paths(std::vector<std::string> file_paths);
 
  void set_row_group_size_bytes(size_t size_bytes);
 
  void set_row_group_size_rows(size_type size_rows);
 
  void set_max_page_size_bytes(size_t size_bytes);
 
  void set_max_page_size_rows(size_type size_rows);
 
  void set_column_index_truncate_length(int32_t size_bytes);
 
  void set_dictionary_policy(dictionary_policy policy);
 
  void set_max_dictionary_size(size_t size_bytes);
 
  void set_max_page_fragment_size(size_type size_rows);
 
  void set_compression_statistics(std::shared_ptr<writer_compression_statistics> comp_stats)
  {
    _compression_stats = std::move(comp_stats);
  }
 
  void enable_write_v2_headers(bool val) { _v2_page_headers = val; }
};
 
class parquet_writer_options_builder {
  parquet_writer_options options;
 
 public:
  explicit parquet_writer_options_builder() = default;
 
  explicit parquet_writer_options_builder(sink_info const& sink, table_view const& table)
    : options(sink, table)
  {
  }
 
  parquet_writer_options_builder& partitions(std::vector<partition_info> partitions);
 
  parquet_writer_options_builder& metadata(table_input_metadata metadata)
  {
    options._metadata = std::move(metadata);
    return *this;
  }
 
  parquet_writer_options_builder& key_value_metadata(
    std::vector<std::map<std::string, std::string>> metadata);
 
  parquet_writer_options_builder& stats_level(statistics_freq sf)
  {
    options._stats_level = sf;
    return *this;
  }
 
  parquet_writer_options_builder& compression(compression_type compression)
  {
    options._compression = compression;
    return *this;
  }
 
  parquet_writer_options_builder& column_chunks_file_paths(std::vector<std::string> file_paths);
 
  parquet_writer_options_builder& row_group_size_bytes(size_t val)
  {
    options.set_row_group_size_bytes(val);
    return *this;
  }
 
  parquet_writer_options_builder& row_group_size_rows(size_type val)
  {
    options.set_row_group_size_rows(val);
    return *this;
  }
 
  parquet_writer_options_builder& max_page_size_bytes(size_t val)
  {
    options.set_max_page_size_bytes(val);
    return *this;
  }
 
  parquet_writer_options_builder& max_page_size_rows(size_type val)
  {
    options.set_max_page_size_rows(val);
    return *this;
  }
 
  parquet_writer_options_builder& column_index_truncate_length(int32_t val)
  {
    options.set_column_index_truncate_length(val);
    return *this;
  }
 
  parquet_writer_options_builder& dictionary_policy(enum dictionary_policy val);
 
  parquet_writer_options_builder& max_dictionary_size(size_t val);
 
  parquet_writer_options_builder& max_page_fragment_size(size_type val);
 
  parquet_writer_options_builder& compression_statistics(
    std::shared_ptr<writer_compression_statistics> const& comp_stats)
  {
    options._compression_stats = comp_stats;
    return *this;
  }
 
  parquet_writer_options_builder& int96_timestamps(bool enabled)
  {
    options._write_timestamps_as_int96 = enabled;
    return *this;
  }
 
  parquet_writer_options_builder& utc_timestamps(bool enabled)
  {
    options._write_timestamps_as_UTC = enabled;
    return *this;
  }
 
  parquet_writer_options_builder& write_v2_headers(bool enabled);
 
  operator parquet_writer_options&&() { return std::move(options); }
 
  parquet_writer_options&& build() { return std::move(options); }
};
 
std::unique_ptr<std::vector<uint8_t>> write_parquet(
  parquet_writer_options const& options, rmm::cuda_stream_view stream = cudf::get_default_stream());
 
std::unique_ptr<std::vector<uint8_t>> merge_row_group_metadata(
  std::vector<std::unique_ptr<std::vector<uint8_t>>> const& metadata_list);
 
class chunked_parquet_writer_options_builder;
 
class chunked_parquet_writer_options {
  // Specify the sink to use for writer output
  sink_info _sink;
  // Specify the compression format to use
  compression_type _compression = compression_type::AUTO;
  // Specify the level of statistics in the output file
  statistics_freq _stats_level = statistics_freq::STATISTICS_ROWGROUP;
  // Optional associated metadata.
  std::optional<table_input_metadata> _metadata;
  // Optional footer key_value_metadata
  std::vector<std::map<std::string, std::string>> _user_data;
  // Parquet writer can write INT96 or TIMESTAMP_MICROS. Defaults to TIMESTAMP_MICROS.
  // If true then overrides any per-column setting in _metadata.
  bool _write_timestamps_as_int96 = false;
  // Parquet writer can write timestamps as UTC. Defaults to true.
  bool _write_timestamps_as_UTC = true;
  // Maximum size of each row group (unless smaller than a single page)
  size_t _row_group_size_bytes = default_row_group_size_bytes;
  // Maximum number of rows in row group (unless smaller than a single page)
  size_type _row_group_size_rows = default_row_group_size_rows;
  // Maximum size of each page (uncompressed)
  size_t _max_page_size_bytes = default_max_page_size_bytes;
  // Maximum number of rows in a page
  size_type _max_page_size_rows = default_max_page_size_rows;
  // Maximum size of min or max values in column index
  int32_t _column_index_truncate_length = default_column_index_truncate_length;
  // When to use dictionary encoding for data
  dictionary_policy _dictionary_policy = dictionary_policy::ALWAYS;
  // Maximum size of column chunk dictionary (in bytes)
  size_t _max_dictionary_size = default_max_dictionary_size;
  // Maximum number of rows in a page fragment
  std::optional<size_type> _max_page_fragment_size;
  // Optional compression statistics
  std::shared_ptr<writer_compression_statistics> _compression_stats;
  // write V2 page headers?
  bool _v2_page_headers = false;
 
  explicit chunked_parquet_writer_options(sink_info const& sink) : _sink(sink) {}
 
  friend chunked_parquet_writer_options_builder;
 
 public:
  chunked_parquet_writer_options() = default;
 
  [[nodiscard]] sink_info const& get_sink() const { return _sink; }
 
  [[nodiscard]] compression_type get_compression() const { return _compression; }
 
  [[nodiscard]] statistics_freq get_stats_level() const { return _stats_level; }
 
  [[nodiscard]] auto const& get_metadata() const { return _metadata; }
 
  std::vector<std::map<std::string, std::string>> const& get_key_value_metadata() const
  {
    return _user_data;
  }
 
  bool is_enabled_int96_timestamps() const { return _write_timestamps_as_int96; }
 
  [[nodiscard]] auto is_enabled_utc_timestamps() const { return _write_timestamps_as_UTC; }
 
  auto get_row_group_size_bytes() const { return _row_group_size_bytes; }
 
  auto get_row_group_size_rows() const { return _row_group_size_rows; }
 
  auto get_max_page_size_bytes() const
  {
    return std::min(_max_page_size_bytes, get_row_group_size_bytes());
  }
 
  auto get_max_page_size_rows() const
  {
    return std::min(_max_page_size_rows, get_row_group_size_rows());
  }
 
  auto get_column_index_truncate_length() const { return _column_index_truncate_length; }
 
  [[nodiscard]] dictionary_policy get_dictionary_policy() const { return _dictionary_policy; }
 
  [[nodiscard]] auto get_max_dictionary_size() const { return _max_dictionary_size; }
 
  [[nodiscard]] auto get_max_page_fragment_size() const { return _max_page_fragment_size; }
 
  [[nodiscard]] std::shared_ptr<writer_compression_statistics> get_compression_statistics() const
  {
    return _compression_stats;
  }
 
  [[nodiscard]] auto is_enabled_write_v2_headers() const { return _v2_page_headers; }
 
  void set_metadata(table_input_metadata metadata) { _metadata = std::move(metadata); }
 
  void set_key_value_metadata(std::vector<std::map<std::string, std::string>> metadata);
 
  void set_stats_level(statistics_freq sf) { _stats_level = sf; }
 
  void set_compression(compression_type compression) { _compression = compression; }
 
  void enable_int96_timestamps(bool req) { _write_timestamps_as_int96 = req; }
 
  void enable_utc_timestamps(bool val) { _write_timestamps_as_UTC = val; }
 
  void set_row_group_size_bytes(size_t size_bytes);
 
  void set_row_group_size_rows(size_type size_rows);
 
  void set_max_page_size_bytes(size_t size_bytes);
 
  void set_max_page_size_rows(size_type size_rows);
 
  void set_column_index_truncate_length(int32_t size_bytes);
 
  void set_dictionary_policy(dictionary_policy policy);
 
  void set_max_dictionary_size(size_t size_bytes);
 
  void set_max_page_fragment_size(size_type size_rows);
 
  void set_compression_statistics(std::shared_ptr<writer_compression_statistics> comp_stats)
  {
    _compression_stats = std::move(comp_stats);
  }
 
  void enable_write_v2_headers(bool val) { _v2_page_headers = val; }
 
  static chunked_parquet_writer_options_builder builder(sink_info const& sink);
};
 
class chunked_parquet_writer_options_builder {
  chunked_parquet_writer_options options;
 
 public:
  chunked_parquet_writer_options_builder() = default;
 
  chunked_parquet_writer_options_builder(sink_info const& sink) : options(sink){};
 
  chunked_parquet_writer_options_builder& metadata(table_input_metadata metadata)
  {
    options._metadata = std::move(metadata);
    return *this;
  }
 
  chunked_parquet_writer_options_builder& key_value_metadata(
    std::vector<std::map<std::string, std::string>> metadata);
 
  chunked_parquet_writer_options_builder& stats_level(statistics_freq sf)
  {
    options._stats_level = sf;
    return *this;
  }
 
  chunked_parquet_writer_options_builder& compression(compression_type compression)
  {
    options._compression = compression;
    return *this;
  }
 
  chunked_parquet_writer_options_builder& int96_timestamps(bool enabled)
  {
    options._write_timestamps_as_int96 = enabled;
    return *this;
  }
 
  chunked_parquet_writer_options_builder& utc_timestamps(bool enabled)
  {
    options._write_timestamps_as_UTC = enabled;
    return *this;
  }
 
  chunked_parquet_writer_options_builder& write_v2_headers(bool enabled);
 
  chunked_parquet_writer_options_builder& row_group_size_bytes(size_t val)
  {
    options.set_row_group_size_bytes(val);
    return *this;
  }
 
  chunked_parquet_writer_options_builder& row_group_size_rows(size_type val)
  {
    options.set_row_group_size_rows(val);
    return *this;
  }
 
  chunked_parquet_writer_options_builder& max_page_size_bytes(size_t val)
  {
    options.set_max_page_size_bytes(val);
    return *this;
  }
 
  chunked_parquet_writer_options_builder& max_page_size_rows(size_type val)
  {
    options.set_max_page_size_rows(val);
    return *this;
  }
 
  chunked_parquet_writer_options_builder& column_index_truncate_length(int32_t val)
  {
    options.set_column_index_truncate_length(val);
    return *this;
  }
 
  chunked_parquet_writer_options_builder& dictionary_policy(enum dictionary_policy val);
 
  chunked_parquet_writer_options_builder& max_dictionary_size(size_t val);
 
  chunked_parquet_writer_options_builder& max_page_fragment_size(size_type val);
 
  chunked_parquet_writer_options_builder& compression_statistics(
    std::shared_ptr<writer_compression_statistics> const& comp_stats)
  {
    options._compression_stats = comp_stats;
    return *this;
  }
 
  operator chunked_parquet_writer_options&&() { return std::move(options); }
 
  chunked_parquet_writer_options&& build() { return std::move(options); }
};
 
class parquet_chunked_writer {
 public:
  parquet_chunked_writer() = default;
 
  parquet_chunked_writer(chunked_parquet_writer_options const& options,
                         rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  parquet_chunked_writer& write(table_view const& table,
                                std::vector<partition_info> const& partitions = {});
 
  std::unique_ptr<std::vector<uint8_t>> close(
    std::vector<std::string> const& column_chunks_file_paths = {});
 
  std::unique_ptr<parquet::detail::writer> writer;
};
  // end of group
 
}  // namespace cudf::io