io/types.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-2023, 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/table/table.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/span.hpp>
 
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <vector>
 
// Forward declarations
namespace arrow {
namespace io {
class RandomAccessFile;
}
}  // namespace arrow
 
namespace cudf {
namespace io {
class data_sink;
class datasource;
}  // namespace io
}  // namespace cudf
 
namespace cudf {
namespace io {
enum class compression_type {
  NONE,    
  AUTO,    
  SNAPPY,  
  GZIP,    
  BZIP2,   
  BROTLI,  
  ZIP,     
  XZ,      
  ZLIB,    
  LZ4,     
  LZO,     
  ZSTD     
};
 
enum class io_type {
  FILEPATH,          
  HOST_BUFFER,       
  DEVICE_BUFFER,     
  VOID,              
  USER_IMPLEMENTED,  
};
 
enum class quote_style {
  MINIMAL,     
  ALL,         
  NONNUMERIC,  
  NONE         
};
 
enum statistics_freq {
  STATISTICS_NONE     = 0,  
  STATISTICS_ROWGROUP = 1,  
  STATISTICS_PAGE     = 2,  
  STATISTICS_COLUMN   = 3,  
};
 
class writer_compression_statistics {
 public:
  writer_compression_statistics() = default;
 
  writer_compression_statistics(size_t num_compressed_bytes,
                                size_t num_failed_bytes,
                                size_t num_skipped_bytes,
                                size_t num_compressed_output_bytes)
    : _num_compressed_bytes(num_compressed_bytes),
      _num_failed_bytes(num_failed_bytes),
      _num_skipped_bytes(num_skipped_bytes),
      _num_compressed_output_bytes(num_compressed_output_bytes)
  {
  }
 
  writer_compression_statistics& operator+=(writer_compression_statistics const& other) noexcept
  {
    _num_compressed_bytes += other._num_compressed_bytes;
    _num_failed_bytes += other._num_failed_bytes;
    _num_skipped_bytes += other._num_skipped_bytes;
    _num_compressed_output_bytes += other._num_compressed_output_bytes;
    return *this;
  }
 
  [[nodiscard]] auto num_compressed_bytes() const noexcept { return _num_compressed_bytes; }
 
  [[nodiscard]] auto num_failed_bytes() const noexcept { return _num_failed_bytes; }
 
  [[nodiscard]] auto num_skipped_bytes() const noexcept { return _num_skipped_bytes; }
 
  [[nodiscard]] auto num_total_input_bytes() const noexcept
  {
    return num_compressed_bytes() + num_failed_bytes() + num_skipped_bytes();
  }
 
  [[nodiscard]] auto compression_ratio() const noexcept
  {
    return static_cast<double>(num_compressed_bytes()) / _num_compressed_output_bytes;
  }
 
 private:
  std::size_t _num_compressed_bytes = 0;  
  std::size_t _num_failed_bytes     = 0;  
  std::size_t _num_skipped_bytes = 0;  
  std::size_t _num_compressed_output_bytes = 0;  
};
 
enum dictionary_policy {
  NEVER,     
  ADAPTIVE,  
  ALWAYS     
};
 
struct column_name_info {
  std::string name;                        
  std::vector<column_name_info> children;  
 
  column_name_info(std::string const& _name) : name(_name) {}
  column_name_info() = default;
};
 
struct table_metadata {
  std::vector<column_name_info>
    schema_info;  
  std::map<std::string, std::string> user_data;  
  std::vector<std::unordered_map<std::string, std::string>>
    per_file_user_data;  
};
 
struct table_with_metadata {
  std::unique_ptr<table> tbl;  
  table_metadata metadata;     
};
 
struct host_buffer {
  // TODO: to be replaced by `host_span`
  char const* data = nullptr;  
  size_t size      = 0;        
  host_buffer()    = default;
  host_buffer(char const* data, size_t size) : data(data), size(size) {}
};
 
template <typename T>
constexpr inline auto is_byte_like_type()
{
  using non_cv_T = std::remove_cv_t<T>;
  return std::is_same_v<non_cv_T, int8_t> || std::is_same_v<non_cv_T, char> ||
         std::is_same_v<non_cv_T, uint8_t> || std::is_same_v<non_cv_T, unsigned char> ||
         std::is_same_v<non_cv_T, std::byte>;
}
 
struct source_info {
  std::vector<std::shared_ptr<arrow::io::RandomAccessFile>> _files;  
 
  source_info() = default;
 
  explicit source_info(std::vector<std::string> const& file_paths) : _filepaths(file_paths) {}
 
  explicit source_info(std::string const& file_path) : _filepaths({file_path}) {}
 
  explicit source_info(std::vector<host_buffer> const& host_buffers) : _type(io_type::HOST_BUFFER)
  {
    _host_buffers.reserve(host_buffers.size());
    std::transform(host_buffers.begin(),
                   host_buffers.end(),
                   std::back_inserter(_host_buffers),
                   [](auto const hb) {
                     return cudf::host_span<std::byte const>{
                       reinterpret_cast<std::byte const*>(hb.data), hb.size};
                   });
  }
 
  explicit source_info(char const* host_data, size_t size)
    : _type(io_type::HOST_BUFFER),
      _host_buffers(
        {cudf::host_span<std::byte const>(reinterpret_cast<std::byte const*>(host_data), size)})
  {
  }
 
  template <typename T, CUDF_ENABLE_IF(is_byte_like_type<std::remove_cv_t<T>>())>
  explicit source_info(cudf::host_span<cudf::host_span<T>> const host_buffers)
    : _type(io_type::HOST_BUFFER)
  {
    if constexpr (not std::is_same_v<std::remove_cv_t<T>, std::byte>) {
      _host_buffers.reserve(host_buffers.size());
      std::transform(host_buffers.begin(),
                     host_buffers.end(),
                     std::back_inserter(_host_buffers),
                     [](auto const s) {
                       return cudf::host_span<std::byte const>{
                         reinterpret_cast<std::byte const*>(s.data()), s.size()};
                     });
    } else {
      _host_buffers.assign(host_buffers.begin(), host_buffers.end());
    }
  }
 
  template <typename T, CUDF_ENABLE_IF(is_byte_like_type<std::remove_cv_t<T>>())>
  explicit source_info(cudf::host_span<T> host_data)
    : _type(io_type::HOST_BUFFER),
      _host_buffers{cudf::host_span<std::byte const>(
        reinterpret_cast<std::byte const*>(host_data.data()), host_data.size())}
  {
  }
 
  explicit source_info(cudf::host_span<cudf::device_span<std::byte const>> device_buffers)
    : _type(io_type::DEVICE_BUFFER), _device_buffers(device_buffers.begin(), device_buffers.end())
  {
  }
 
  explicit source_info(cudf::device_span<std::byte const> d_buffer)
    : _type(io_type::DEVICE_BUFFER), _device_buffers({{d_buffer}})
  {
  }
 
  explicit source_info(std::vector<cudf::io::datasource*> const& sources)
    : _type(io_type::USER_IMPLEMENTED), _user_sources(sources)
  {
  }
 
  explicit source_info(cudf::io::datasource* source)
    : _type(io_type::USER_IMPLEMENTED), _user_sources({source})
  {
  }
 
  [[nodiscard]] auto type() const { return _type; }
  [[nodiscard]] auto const& filepaths() const { return _filepaths; }
  [[nodiscard]] auto const& host_buffers() const { return _host_buffers; }
  [[nodiscard]] auto const& device_buffers() const { return _device_buffers; }
  [[nodiscard]] auto const& files() const { return _files; }
  [[nodiscard]] auto const& user_sources() const { return _user_sources; }
 
 private:
  io_type _type = io_type::FILEPATH;
  std::vector<std::string> _filepaths;
  std::vector<cudf::host_span<std::byte const>> _host_buffers;
  std::vector<cudf::device_span<std::byte const>> _device_buffers;
  std::vector<cudf::io::datasource*> _user_sources;
};
 
struct sink_info {
  sink_info() = default;
  sink_info(size_t num_sinks) : _num_sinks(num_sinks) {}
 
  explicit sink_info(std::vector<std::string> const& file_paths)
    : _type(io_type::FILEPATH), _num_sinks(file_paths.size()), _filepaths(file_paths)
  {
  }
 
  explicit sink_info(std::string const& file_path)
    : _type(io_type::FILEPATH), _filepaths({file_path})
  {
  }
 
  explicit sink_info(std::vector<std::vector<char>*> const& buffers)
    : _type(io_type::HOST_BUFFER), _num_sinks(buffers.size()), _buffers(buffers)
  {
  }
  explicit sink_info(std::vector<char>* buffer) : _type(io_type::HOST_BUFFER), _buffers({buffer}) {}
 
  explicit sink_info(std::vector<cudf::io::data_sink*> const& user_sinks)
    : _type(io_type::USER_IMPLEMENTED), _num_sinks(user_sinks.size()), _user_sinks(user_sinks)
  {
  }
 
  explicit sink_info(class cudf::io::data_sink* user_sink)
    : _type(io_type::USER_IMPLEMENTED), _user_sinks({user_sink})
  {
  }
 
  [[nodiscard]] auto type() const { return _type; }
  [[nodiscard]] auto num_sinks() const { return _num_sinks; }
  [[nodiscard]] auto const& filepaths() const { return _filepaths; }
  [[nodiscard]] auto const& buffers() const { return _buffers; }
  [[nodiscard]] auto const& user_sinks() const { return _user_sinks; }
 
 private:
  io_type _type     = io_type::VOID;
  size_t _num_sinks = 1;
  std::vector<std::string> _filepaths;
  std::vector<std::vector<char>*> _buffers;
  std::vector<cudf::io::data_sink*> _user_sinks;
};
 
class table_input_metadata;
 
class column_in_metadata {
  friend table_input_metadata;
  std::string _name = "";
  std::optional<bool> _nullable;
  bool _list_column_is_map  = false;
  bool _use_int96_timestamp = false;
  bool _output_as_binary    = false;
  std::optional<uint8_t> _decimal_precision;
  std::optional<int32_t> _parquet_field_id;
  std::vector<column_in_metadata> children;
 
 public:
  column_in_metadata() = default;
  column_in_metadata(std::string_view name) : _name{name} {}
  column_in_metadata& add_child(column_in_metadata const& child)
  {
    children.push_back(child);
    return *this;
  }
 
  column_in_metadata& set_name(std::string const& name) noexcept
  {
    _name = name;
    return *this;
  }
 
  column_in_metadata& set_nullability(bool nullable) noexcept
  {
    _nullable = nullable;
    return *this;
  }
 
  column_in_metadata& set_list_column_as_map() noexcept
  {
    _list_column_is_map = true;
    return *this;
  }
 
  column_in_metadata& set_int96_timestamps(bool req) noexcept
  {
    _use_int96_timestamp = req;
    return *this;
  }
 
  column_in_metadata& set_decimal_precision(uint8_t precision) noexcept
  {
    _decimal_precision = precision;
    return *this;
  }
 
  column_in_metadata& set_parquet_field_id(int32_t field_id) noexcept
  {
    _parquet_field_id = field_id;
    return *this;
  }
 
  column_in_metadata& set_output_as_binary(bool binary) noexcept
  {
    _output_as_binary = binary;
    return *this;
  }
 
  column_in_metadata& child(size_type i) noexcept { return children[i]; }
 
  [[nodiscard]] column_in_metadata const& child(size_type i) const noexcept { return children[i]; }
 
  [[nodiscard]] std::string get_name() const noexcept { return _name; }
 
  [[nodiscard]] bool is_nullability_defined() const noexcept { return _nullable.has_value(); }
 
  [[nodiscard]] bool nullable() const { return _nullable.value(); }
 
  [[nodiscard]] bool is_map() const noexcept { return _list_column_is_map; }
 
  [[nodiscard]] bool is_enabled_int96_timestamps() const noexcept { return _use_int96_timestamp; }
 
  [[nodiscard]] bool is_decimal_precision_set() const noexcept
  {
    return _decimal_precision.has_value();
  }
 
  [[nodiscard]] uint8_t get_decimal_precision() const { return _decimal_precision.value(); }
 
  [[nodiscard]] bool is_parquet_field_id_set() const noexcept
  {
    return _parquet_field_id.has_value();
  }
 
  [[nodiscard]] int32_t get_parquet_field_id() const { return _parquet_field_id.value(); }
 
  [[nodiscard]] size_type num_children() const noexcept { return children.size(); }
 
  [[nodiscard]] bool is_enabled_output_as_binary() const noexcept { return _output_as_binary; }
};
 
class table_input_metadata {
 public:
  table_input_metadata() = default;  // Required by cython
 
  table_input_metadata(table_view const& table);
 
  std::vector<column_in_metadata> column_metadata;  
};
 
struct partition_info {
  size_type start_row;  
  size_type num_rows;   
 
  partition_info() = default;
  partition_info(size_type start_row, size_type num_rows) : start_row(start_row), num_rows(num_rows)
  {
  }
};
 
class reader_column_schema {
  // Whether to read binary data as a string column
  bool _convert_binary_to_strings{true};
 
  std::vector<reader_column_schema> children;
 
 public:
  reader_column_schema() = default;
 
  reader_column_schema(size_type number_of_children) { children.resize(number_of_children); }
 
  reader_column_schema(host_span<reader_column_schema> const& child_span)
  {
    children.assign(child_span.begin(), child_span.end());
  }
 
  reader_column_schema& add_child(reader_column_schema const& child)
  {
    children.push_back(child);
    return *this;
  }
 
  [[nodiscard]] reader_column_schema& child(size_type i) { return children[i]; }
 
  [[nodiscard]] reader_column_schema const& child(size_type i) const { return children[i]; }
 
  reader_column_schema& set_convert_binary_to_strings(bool convert_to_string)
  {
    _convert_binary_to_strings = convert_to_string;
    return *this;
  }
 
  [[nodiscard]] bool is_enabled_convert_binary_to_strings() const
  {
    return _convert_binary_to_strings;
  }
 
  [[nodiscard]] size_t get_num_children() const { return children.size(); }
};
 
}  // namespace io
}  // namespace cudf