orc.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/io/detail/orc.hpp>
#include <cudf/io/types.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/export.hpp>
 
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
 
namespace CUDF_EXPORT cudf {
namespace io {
constexpr size_t default_stripe_size_bytes   = 64 * 1024 * 1024;  
constexpr size_type default_stripe_size_rows = 1000000;  
constexpr size_type default_row_index_stride = 10000;    
 
class orc_reader_options_builder;
 
class orc_reader_options {
  source_info _source;
 
  // Names of column to read; `nullopt` is all
  std::optional<std::vector<std::string>> _columns;
 
  // List of individual stripes to read (ignored if empty)
  std::vector<std::vector<size_type>> _stripes;
  // Rows to skip from the start
  int64_t _skip_rows = 0;
  // Rows to read; `nullopt` is all
  std::optional<int64_t> _num_rows;
 
  // Whether to use row index to speed-up reading
  bool _use_index = true;
 
  // Whether to use numpy-compatible dtypes
  bool _use_np_dtypes = true;
  // Cast timestamp columns to a specific type
  data_type _timestamp_type{type_id::EMPTY};
 
  // Columns that should be read as Decimal128
  std::vector<std::string> _decimal128_columns;
 
  friend orc_reader_options_builder;
 
  explicit orc_reader_options(source_info src) : _source{std::move(src)} {}
 
 public:
  orc_reader_options() = default;
 
  static orc_reader_options_builder builder(source_info src);
 
  [[nodiscard]] source_info const& get_source() const { return _source; }
 
  [[nodiscard]] auto const& get_columns() const { return _columns; }
 
  [[nodiscard]] auto const& get_stripes() const { return _stripes; }
 
  [[nodiscard]] int64_t get_skip_rows() const { return _skip_rows; }
 
  [[nodiscard]] std::optional<int64_t> const& get_num_rows() const { return _num_rows; }
 
  [[nodiscard]] bool is_enabled_use_index() const { return _use_index; }
 
  [[nodiscard]] bool is_enabled_use_np_dtypes() const { return _use_np_dtypes; }
 
  [[nodiscard]] data_type get_timestamp_type() const { return _timestamp_type; }
 
  [[nodiscard]] std::vector<std::string> const& get_decimal128_columns() const
  {
    return _decimal128_columns;
  }
 
  // Setters
 
  void set_columns(std::vector<std::string> col_names) { _columns = std::move(col_names); }
 
  void set_stripes(std::vector<std::vector<size_type>> stripes)
  {
    CUDF_EXPECTS(stripes.empty() or (_skip_rows == 0), "Can't set stripes along with skip_rows");
    CUDF_EXPECTS(stripes.empty() or not _num_rows.has_value(),
                 "Can't set stripes along with num_rows");
    _stripes = std::move(stripes);
  }
 
  void set_skip_rows(int64_t rows)
  {
    CUDF_EXPECTS(rows >= 0, "skip_rows cannot be negative");
    CUDF_EXPECTS(rows == 0 or _stripes.empty(), "Can't set both skip_rows along with stripes");
    _skip_rows = rows;
  }
 
  void set_num_rows(int64_t nrows)
  {
    CUDF_EXPECTS(nrows >= 0, "num_rows cannot be negative");
    CUDF_EXPECTS(_stripes.empty(), "Can't set both num_rows and stripes");
    _num_rows = nrows;
  }
 
  void enable_use_index(bool use) { _use_index = use; }
 
  void enable_use_np_dtypes(bool use) { _use_np_dtypes = use; }
 
  void set_timestamp_type(data_type type) { _timestamp_type = type; }
 
  void set_decimal128_columns(std::vector<std::string> val)
  {
    _decimal128_columns = std::move(val);
  }
};
 
class orc_reader_options_builder {
  orc_reader_options options;
 
 public:
  explicit orc_reader_options_builder() = default;
 
  explicit orc_reader_options_builder(source_info src) : options{std::move(src)} {}
 
  orc_reader_options_builder& columns(std::vector<std::string> col_names)
  {
    options._columns = std::move(col_names);
    return *this;
  }
 
  orc_reader_options_builder& stripes(std::vector<std::vector<size_type>> stripes)
  {
    options.set_stripes(std::move(stripes));
    return *this;
  }
 
  orc_reader_options_builder& skip_rows(int64_t rows)
  {
    options.set_skip_rows(rows);
    return *this;
  }
 
  orc_reader_options_builder& num_rows(int64_t nrows)
  {
    options.set_num_rows(nrows);
    return *this;
  }
 
  orc_reader_options_builder& use_index(bool use)
  {
    options._use_index = use;
    return *this;
  }
 
  orc_reader_options_builder& use_np_dtypes(bool use)
  {
    options._use_np_dtypes = use;
    return *this;
  }
 
  orc_reader_options_builder& timestamp_type(data_type type)
  {
    options._timestamp_type = type;
    return *this;
  }
 
  orc_reader_options_builder& decimal128_columns(std::vector<std::string> val)
  {
    options._decimal128_columns = std::move(val);
    return *this;
  }
 
  operator orc_reader_options&&() { return std::move(options); }
 
  orc_reader_options&& build() { return std::move(options); }
};
 
table_with_metadata read_orc(
  orc_reader_options const& options,
  rmm::cuda_stream_view stream      = cudf::get_default_stream(),
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
class chunked_orc_reader {
 public:
  chunked_orc_reader();
 
  explicit chunked_orc_reader(
    std::size_t chunk_read_limit,
    std::size_t pass_read_limit,
    size_type output_row_granularity,
    orc_reader_options const& options,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  explicit chunked_orc_reader(
    std::size_t chunk_read_limit,
    std::size_t pass_read_limit,
    orc_reader_options const& options,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  explicit chunked_orc_reader(
    std::size_t chunk_read_limit,
    orc_reader_options const& options,
    rmm::cuda_stream_view stream      = cudf::get_default_stream(),
    rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
 
  ~chunked_orc_reader();
 
  [[nodiscard]] bool has_next() const;
 
  [[nodiscard]] table_with_metadata read_chunk() const;
 
 private:
  std::unique_ptr<cudf::io::orc::detail::chunked_reader> reader;
};
  // end of group
class orc_writer_options_builder;
 
static constexpr statistics_freq ORC_STATISTICS_STRIPE    = statistics_freq::STATISTICS_ROWGROUP;
static constexpr statistics_freq ORC_STATISTICS_ROW_GROUP = statistics_freq::STATISTICS_PAGE;
 
class orc_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 frequency of statistics collection
  statistics_freq _stats_freq = ORC_STATISTICS_ROW_GROUP;
  // Maximum size of each stripe (unless smaller than a single row group)
  size_t _stripe_size_bytes = default_stripe_size_bytes;
  // Maximum number of rows in stripe (unless smaller than a single row group)
  size_type _stripe_size_rows = default_stripe_size_rows;
  // Row index stride (maximum number of rows in each row group)
  size_type _row_index_stride = default_row_index_stride;
  // Set of columns to output
  table_view _table;
  // Optional associated metadata
  std::optional<table_input_metadata> _metadata;
  // Optional footer key_value_metadata
  std::map<std::string, std::string> _user_data;
  // Optional compression statistics
  std::shared_ptr<writer_compression_statistics> _compression_stats;
  // Specify whether string dictionaries should be alphabetically sorted
  bool _enable_dictionary_sort = true;
 
  friend orc_writer_options_builder;
 
  explicit orc_writer_options(sink_info sink, table_view table)
    : _sink(std::move(sink)), _table(std::move(table))
  {
  }
 
 public:
  explicit orc_writer_options() = default;
 
  static orc_writer_options_builder builder(sink_info const& sink, table_view const& table);
 
  [[nodiscard]] sink_info const& get_sink() const { return _sink; }
 
  [[nodiscard]] compression_type get_compression() const { return _compression; }
 
  [[nodiscard]] bool is_enabled_statistics() const
  {
    return _stats_freq != statistics_freq::STATISTICS_NONE;
  }
 
  [[nodiscard]] statistics_freq get_statistics_freq() const { return _stats_freq; }
 
  [[nodiscard]] auto get_stripe_size_bytes() const { return _stripe_size_bytes; }
 
  [[nodiscard]] auto get_stripe_size_rows() const { return _stripe_size_rows; }
 
  [[nodiscard]] auto get_row_index_stride() const
  {
    auto const unaligned_stride = std::min(_row_index_stride, get_stripe_size_rows());
    return unaligned_stride - unaligned_stride % 8;
  }
 
  [[nodiscard]] table_view get_table() const { return _table; }
 
  [[nodiscard]] auto const& get_metadata() const { return _metadata; }
 
  [[nodiscard]] std::map<std::string, std::string> const& get_key_value_metadata() const
  {
    return _user_data;
  }
 
  [[nodiscard]] std::shared_ptr<writer_compression_statistics> get_compression_statistics() const
  {
    return _compression_stats;
  }
 
  [[nodiscard]] bool get_enable_dictionary_sort() const { return _enable_dictionary_sort; }
 
  // Setters
 
  void set_compression(compression_type comp) { _compression = comp; }
 
  void enable_statistics(statistics_freq val) { _stats_freq = val; }
 
  void set_stripe_size_bytes(size_t size_bytes)
  {
    CUDF_EXPECTS(size_bytes >= 64 << 10, "64KB is the minimum stripe size");
    _stripe_size_bytes = size_bytes;
  }
 
  void set_stripe_size_rows(size_type size_rows)
  {
    CUDF_EXPECTS(size_rows >= 512, "Maximum stripe size cannot be smaller than 512");
    _stripe_size_rows = size_rows;
  }
 
  void set_row_index_stride(size_type stride)
  {
    CUDF_EXPECTS(stride >= 512, "Row index stride cannot be smaller than 512");
    _row_index_stride = stride;
  }
 
  void set_table(table_view tbl) { _table = tbl; }
 
  void set_metadata(table_input_metadata meta) { _metadata = std::move(meta); }
 
  void set_key_value_metadata(std::map<std::string, std::string> metadata)
  {
    _user_data = std::move(metadata);
  }
 
  void set_compression_statistics(std::shared_ptr<writer_compression_statistics> comp_stats)
  {
    _compression_stats = std::move(comp_stats);
  }
 
  void set_enable_dictionary_sort(bool val) { _enable_dictionary_sort = val; }
};
 
class orc_writer_options_builder {
  orc_writer_options options;
 
 public:
  orc_writer_options_builder() = default;
 
  orc_writer_options_builder(sink_info const& sink, table_view const& table) : options{sink, table}
  {
  }
 
  orc_writer_options_builder& compression(compression_type comp)
  {
    options._compression = comp;
    return *this;
  }
 
  orc_writer_options_builder& enable_statistics(statistics_freq val)
  {
    options._stats_freq = val;
    return *this;
  }
 
  orc_writer_options_builder& stripe_size_bytes(size_t val)
  {
    options.set_stripe_size_bytes(val);
    return *this;
  }
 
  orc_writer_options_builder& stripe_size_rows(size_type val)
  {
    options.set_stripe_size_rows(val);
    return *this;
  }
 
  orc_writer_options_builder& row_index_stride(size_type val)
  {
    options.set_row_index_stride(val);
    return *this;
  }
 
  orc_writer_options_builder& table(table_view tbl)
  {
    options._table = tbl;
    return *this;
  }
 
  orc_writer_options_builder& metadata(table_input_metadata meta)
  {
    options._metadata = std::move(meta);
    return *this;
  }
 
  orc_writer_options_builder& key_value_metadata(std::map<std::string, std::string> metadata)
  {
    options._user_data = std::move(metadata);
    return *this;
  }
 
  orc_writer_options_builder& compression_statistics(
    std::shared_ptr<writer_compression_statistics> const& comp_stats)
  {
    options._compression_stats = comp_stats;
    return *this;
  }
 
  orc_writer_options_builder& enable_dictionary_sort(bool val)
  {
    options._enable_dictionary_sort = val;
    return *this;
  }
 
  operator orc_writer_options&&() { return std::move(options); }
 
  orc_writer_options&& build() { return std::move(options); }
};
 
void write_orc(orc_writer_options const& options,
               rmm::cuda_stream_view stream = cudf::get_default_stream());
 
class chunked_orc_writer_options_builder;
 
class chunked_orc_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 granularity of statistics collection
  statistics_freq _stats_freq = ORC_STATISTICS_ROW_GROUP;
  // Maximum size of each stripe (unless smaller than a single row group)
  size_t _stripe_size_bytes = default_stripe_size_bytes;
  // Maximum number of rows in stripe (unless smaller than a single row group)
  size_type _stripe_size_rows = default_stripe_size_rows;
  // Row index stride (maximum number of rows in each row group)
  size_type _row_index_stride = default_row_index_stride;
  // Optional associated metadata
  std::optional<table_input_metadata> _metadata;
  // Optional footer key_value_metadata
  std::map<std::string, std::string> _user_data;
  // Optional compression statistics
  std::shared_ptr<writer_compression_statistics> _compression_stats;
  // Specify whether string dictionaries should be alphabetically sorted
  bool _enable_dictionary_sort = true;
 
  friend chunked_orc_writer_options_builder;
 
  chunked_orc_writer_options(sink_info sink) : _sink(std::move(sink)) {}
 
 public:
  explicit chunked_orc_writer_options() = default;
 
  static chunked_orc_writer_options_builder builder(sink_info const& sink);
 
  [[nodiscard]] sink_info const& get_sink() const { return _sink; }
 
  [[nodiscard]] compression_type get_compression() const { return _compression; }
 
  [[nodiscard]] statistics_freq get_statistics_freq() const { return _stats_freq; }
 
  [[nodiscard]] auto get_stripe_size_bytes() const { return _stripe_size_bytes; }
 
  [[nodiscard]] auto get_stripe_size_rows() const { return _stripe_size_rows; }
 
  [[nodiscard]] auto get_row_index_stride() const
  {
    auto const unaligned_stride = std::min(_row_index_stride, get_stripe_size_rows());
    return unaligned_stride - unaligned_stride % 8;
  }
 
  [[nodiscard]] auto const& get_metadata() const { return _metadata; }
 
  [[nodiscard]] std::map<std::string, std::string> const& get_key_value_metadata() const
  {
    return _user_data;
  }
 
  [[nodiscard]] std::shared_ptr<writer_compression_statistics> get_compression_statistics() const
  {
    return _compression_stats;
  }
 
  [[nodiscard]] bool get_enable_dictionary_sort() const { return _enable_dictionary_sort; }
 
  // Setters
 
  void set_compression(compression_type comp) { _compression = comp; }
 
  void enable_statistics(statistics_freq val) { _stats_freq = val; }
 
  void set_stripe_size_bytes(size_t size_bytes)
  {
    CUDF_EXPECTS(size_bytes >= 64 << 10, "64KB is the minimum stripe size");
    _stripe_size_bytes = size_bytes;
  }
 
  void set_stripe_size_rows(size_type size_rows)
  {
    CUDF_EXPECTS(size_rows >= 512, "maximum stripe size cannot be smaller than 512");
    _stripe_size_rows = size_rows;
  }
 
  void set_row_index_stride(size_type stride)
  {
    CUDF_EXPECTS(stride >= 512, "Row index stride cannot be smaller than 512");
    _row_index_stride = stride;
  }
 
  void metadata(table_input_metadata meta) { _metadata = std::move(meta); }
 
  void set_key_value_metadata(std::map<std::string, std::string> metadata)
  {
    _user_data = std::move(metadata);
  }
 
  void set_compression_statistics(std::shared_ptr<writer_compression_statistics> comp_stats)
  {
    _compression_stats = std::move(comp_stats);
  }
 
  void set_enable_dictionary_sort(bool val) { _enable_dictionary_sort = val; }
};
 
class chunked_orc_writer_options_builder {
  chunked_orc_writer_options options;
 
 public:
  chunked_orc_writer_options_builder() = default;
 
  explicit chunked_orc_writer_options_builder(sink_info const& sink) : options{sink} {}
 
  chunked_orc_writer_options_builder& compression(compression_type comp)
  {
    options._compression = comp;
    return *this;
  }
 
  chunked_orc_writer_options_builder& enable_statistics(statistics_freq val)
  {
    options._stats_freq = val;
    return *this;
  }
 
  chunked_orc_writer_options_builder& stripe_size_bytes(size_t val)
  {
    options.set_stripe_size_bytes(val);
    return *this;
  }
 
  chunked_orc_writer_options_builder& stripe_size_rows(size_type val)
  {
    options.set_stripe_size_rows(val);
    return *this;
  }
 
  chunked_orc_writer_options_builder& row_index_stride(size_type val)
  {
    options.set_row_index_stride(val);
    return *this;
  }
 
  chunked_orc_writer_options_builder& metadata(table_input_metadata meta)
  {
    options._metadata = std::move(meta);
    return *this;
  }
 
  chunked_orc_writer_options_builder& key_value_metadata(
    std::map<std::string, std::string> metadata)
  {
    options._user_data = std::move(metadata);
    return *this;
  }
 
  chunked_orc_writer_options_builder& compression_statistics(
    std::shared_ptr<writer_compression_statistics> const& comp_stats)
  {
    options._compression_stats = comp_stats;
    return *this;
  }
 
  chunked_orc_writer_options_builder& enable_dictionary_sort(bool val)
  {
    options._enable_dictionary_sort = val;
    return *this;
  }
 
  operator chunked_orc_writer_options&&() { return std::move(options); }
 
  chunked_orc_writer_options&& build() { return std::move(options); }
};
 
class orc_chunked_writer {
 public:
  orc_chunked_writer();
 
  ~orc_chunked_writer();
 
  orc_chunked_writer(chunked_orc_writer_options const& options,
                     rmm::cuda_stream_view stream = cudf::get_default_stream());
 
  orc_chunked_writer& write(table_view const& table);
 
  void close();
 
  std::unique_ptr<orc::detail::writer> writer;
};
  // end of group
}  // namespace io
}  // namespace CUDF_EXPORT cudf