orc.hpp

前往此文件的文档。

/*
 * Copyright (c) 2020-2025, 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
 *
 * https://apache.ac.cn/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 <cudf/utilities/memory_resource.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;

/**
 * @brief ORC reader options builder.
 * @ingroup io_readers
 */

class orc_reader_options {
 source_info _source;

 // 要读取的列名；`nullopt` 表示所有列
 std::optional<std::vector<std::string>> _columns;

 // 要读取的单独条带列表（如果为空则忽略）
 std::vector<std::vector<size_type>> _stripes;
 // 从开头跳过的行数
 int64_t _skip_rows = 0;
 // 要读取的行数；`nullopt` 表示所有行
 std::optional<int64_t> _num_rows;

 // 是否使用行索引加速读取
 bool _use_index = true;

 // 是否使用 numpy 兼容的数据类型
 bool _use_np_dtypes = true;
 // 将时间戳列转换为特定类型
 data_type _timestamp_type{type_id::EMPTY};

 // 应读取为 Decimal128 的列
 std::vector<std::string> _decimal128_columns;

 friend orc_reader_options_builder;

 /// @brief Construct an ORC reader options with a source.
 /// @param src The source to read from.
 explicit orc_reader_options(source_info src) : _source{std::move(src)} {}

 public
 /// @brief Default constructor.
 orc_reader_options() = default;

 /// @brief Return a builder object to create the options.
 /// @param src The source to read from.
 /// @return orc_reader_options_builder The builder object.
 static orc_reader_options_builder builder(source_info src);

 /// @brief Get the source.
 [[nodiscard]] source_info const& get_source() const { return _source; }

 /// @brief Get the requested column names.
 [[nodiscard]] auto const& get_columns() const { return _columns; }

 /// @brief Get the requested stripe list.
 [[nodiscard]] auto const& get_stripes() const { return _stripes; }

 /// @brief Get the number of rows to skip from the start.
 [[nodiscard]] int64_t get_skip_rows() const { return _skip_rows; }

 /// @brief Get the number of rows to read.
 /// nullopt implies reading till the end.
 [[nodiscard]] std::optional<int64_t> const& get_num_rows() const { return _num_rows; }

 /// @brief Whether row index is used for speed-up.
 [[nodiscard]] bool is_enabled_use_index() const { return _use_index; }

 /// @brief Whether numpy-compatible dtypes are used.
 /// This option will be removed in a future release.
 [[nodiscard]] bool is_enabled_use_np_dtypes() const { return _use_np_dtypes; }

 /// @brief Get the timestamp data type.
 [[nodiscard]] data_type get_timestamp_type() const { return _timestamp_type; }

 /// @brief Get the names of columns to be read as Decimal128.
 {
 return _decimal128_columns;
 }

 // Setters

 /// @brief Set the names of columns to read; `nullopt` is all.
 /// @param col_names The list of column names to read.
 void set_columns(std::vector<std::string> col_names) { _columns = std::move(col_names); }

 /// @brief Set the list of individual stripes to read (ignored if empty).
 /// Only one of `set_stripes`, `set_skip_rows`, or `set_num_rows` can be specified.
 /// @param stripes The list of individual stripes to read.
 void set_stripes(std::vector<std::vector<size_type>> stripes)
 {
 CUDF_EXPECTS(stripes.empty() or (_skip_rows == 0), "无法同时设置条带和跳过的行数");
 CUDF_EXPECTS(stripes.empty() or not _num_rows.has_value(),
 "无法同时设置条带和行数");
 _stripes = std::move(stripes);
 }

 /// @brief Set the number of rows to skip from the start.
 /// @param rows The number of rows to skip.
 void set_skip_rows(int64_t rows)
 {
 CUDF_EXPECTS(rows >= 0, "skip_rows 不能为负");
 CUDF_EXPECTS(rows == 0 or _stripes.empty(), "无法同时设置 skip_rows 和条带");
 _skip_rows = rows;
 }

 /// @brief Set the number of rows to read.
 /// `nullopt` implies reading till the end.
 /// @param nrows The number of rows to read.
 void set_num_rows(int64_t nrows)
 {
 CUDF_EXPECTS(nrows >= 0, "num_rows 不能为负");
 CUDF_EXPECTS(_stripes.empty(), "无法同时设置 num_rows 和条带");
 _num_rows = nrows;
 }

 /// @brief Enable or disable row index usage.
 /// @param use Enable if `true`, disable if `false`.
 void enable_use_index(bool use) { _use_index = use; }

 /// @brief Enable or disable numpy-compatible dtypes.
 /// This option will be removed in a future release.
 void enable_use_np_dtypes(bool use) { _use_np_dtypes = use; }

 /// @brief Set the timestamp data type.
 /// @param type The desired timestamp type.
 void set_timestamp_type(data_type type) { _timestamp_type = type; }

 /// @brief Set the names of columns to be read as Decimal128.
 /// @param val The list of column names.
 void set_decimal128_columns(std::vector<std::string> val)
 {
 _decimal128_columns = std::move(val);
 }
};

class orc_reader_options_builder {
 orc_reader_options options;

 public
 /// @brief Default constructor.
 explicit orc_reader_options_builder() = default;

 /// @brief Construct with a source.
 /// @param src The source to read from.
 explicit orc_reader_options_builder(source_info src) : options{std::move(src)} {}

 /// @brief Set the names of columns to read.
 /// @param col_names The list of column names to read.
 orc_reader_options_builder& columns(std::vector<std::string> col_names)
 {
 options._columns = std::move(col_names);
 return *this;
 }

 /// @brief Set the list of individual stripes to read (ignored if empty).
 /// @param stripes The list of individual stripes to read.
 orc_reader_options_builder& stripes(std::vector<std::vector<size_type>> stripes)
 {
 options.set_stripes(std::move(stripes));
 return *this;
 }

 /// @brief Set the number of rows to skip from the start.
 /// @param rows The number of rows to skip.
 orc_reader_options_builder& skip_rows(int64_t rows)
 {
 options.set_skip_rows(rows);
 return *this;
 }

 /// @brief Set the number of rows to read.
 /// @param nrows The number of rows to read.
 orc_reader_options_builder& num_rows(int64_t nrows)
 {
 options.set_num_rows(nrows);
 return *this;
 }

 /// @brief Enable or disable row index usage.
 /// @param use Enable if `true`, disable if `false`.
 orc_reader_options_builder& use_index(bool use)
 {
 options._use_index = use;
 return *this;
 }

 /// @brief Enable or disable numpy-compatible dtypes.
 /// This option will be removed in a future release.
 /// @param use Enable if `true`, disable if `false`.
 orc_reader_options_builder& use_np_dtypes(bool use)
 {
 options._use_np_dtypes = use;
 return *this;
 }

 /// @brief Set the timestamp data type.
 /// @param type The desired timestamp type.
 orc_reader_options_builder& timestamp_type(data_type type)
 {
 options._timestamp_type = type;
 return *this;
 }

 /// @brief Set the names of columns to be read as Decimal128.
 /// @param val The list of column names.
 orc_reader_options_builder& decimal128_columns(std::vector<std::string> val)
 {
 options._decimal128_columns = std::move(val);
 return *this;
 }

 /// @brief Implicit conversion to orc_reader_options.
 operator orc_reader_options&&() { return std::move(options); }

 /// @brief Explicit conversion to orc_reader_options.
 /// @return The orc reader options object.
 orc_reader_options&& build() { return std::move(options); }
};

/// @brief Read a dataset from an ORC file/buffer into a cudf table.
///
/// @param options Settings for the read operation.
/// @param stream CUDA stream used for device memory operations and kernel launches.
/// @param mr Device memory resource to use for device memory allocation.
/// @return The table and any metadata.
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 = cudf::get_current_device_resource_ref());

/// @brief ORC chunked reader class.
class chunked_orc_reader {
 public
 /// @brief Default constructor (disabled).
 chunked_orc_reader() = delete;

 /// @brief Constructor for chunked ORC reader.
 ///
 /// @param chunk_read_limit Minimum size, in bytes, of dataset chunk to be read at once.
 ///   Actual chunks can be larger due to minimum stripe/rowgroup size.
 /// @param pass_read_limit Minimum size, in bytes, of dataset part to be read in one pass.
 /// @param output_row_granularity Number of rows to read as a single unit (row group).
 ///   Has to be multiple of 1024 or 10000 rows typically. Default is 10000.
 /// @param options Settings for the read operation.
 /// @param stream CUDA stream used for device memory operations and kernel launches.
 /// @param mr Device memory resource to use for device memory allocation.
 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 = cudf::get_current_device_resource_ref());

 /// @brief Constructor for chunked ORC reader with default output row granularity.
 ///
 /// @param chunk_read_limit Minimum size, in bytes, of dataset chunk to be read at once.
 ///   Actual chunks can be larger due to minimum stripe/rowgroup size.
 /// @param pass_read_limit Minimum size, in bytes, of dataset part to be read in one pass.
 /// @param options Settings for the read operation.
 /// @param stream CUDA stream used for device memory operations and kernel launches.
 /// @param mr Device memory resource to use for device memory allocation.
 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 = cudf::get_current_device_resource_ref());

 /// @brief Constructor for chunked ORC reader with default pass read limit and output row
 /// granularity.
 /// @param chunk_read_limit Minimum size, in bytes, of dataset chunk to be read at once.
 ///   Actual chunks can be larger due to minimum stripe/rowgroup size.
 /// @param options Settings for the read operation.
 /// @param stream CUDA stream used for device memory operations and kernel launches.
 /// @param mr Device memory resource to use for device memory allocation.
 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 = cudf::get_current_device_resource_ref());

 /// @brief Destructor for chunked ORC reader.

 /// @brief Check if there is more data to read.
 ~chunked_orc_reader();

 /// @return bool True if there is more data to read, false otherwise.
 [[nodiscard]] bool has_next() const;

 /// @brief Read the next chunk from the dataset.
 /// @return The table and any metadata for the next chunk.
 [[nodiscard]] table_with_metadata read_chunk() const;

 private
 std::unique_ptr<cudf::io::orc::detail::chunked_reader> reader;
};
 // end of group

/**
 * @brief ORC writer options builder.
 * @ingroup io_writers
 */

/// @brief Statistics granularity values for ORC writers.
/// These are deprecated in favor of the general `statistics_freq`.
static constexpr statistics_freq ORC_STATISTICS_STRIPE = statistics_freq::STATISTICS_ROWGROUP;
static constexpr statistics_freq ORC_STATISTICS_ROW_GROUP = statistics_freq::STATISTICS_PAGE;

/// @brief ORC writer options.
class orc_writer_options {
 // 指定写入器输出使用的接收器
 sink_info _sink;
 // 指定要使用的压缩格式
 compression_type _compression = compression_type::SNAPPY;
 // 指定统计信息收集的频率
 statistics_freq _stats_freq = ORC_STATISTICS_ROW_GROUP;
 // 每个条带的最大大小（除非小于单个行组）
 size_t _stripe_size_bytes = default_stripe_size_bytes;
 // 条带中的最大行数（除非小于单个行组）
 size_type _stripe_size_rows = default_stripe_size_rows;
 // 行索引步长（每个行组中的最大行数）
 size_type _row_index_stride = default_row_index_stride;
 // 要输出的列集
 table_view _table;
 // 可选的关联元数据
 std::optional<table_input_metadata> _metadata;
 // 可选的页脚 key_value_metadata
 std::map<std::string, std::string> _user_data;
 // 可选的压缩统计信息
 std::shared_ptr<writer_compression_statistics> _compression_stats;
 // 指定字符串字典是否应按字母顺序排序
 bool _enable_dictionary_sort = true;

 friend orc_writer_options_builder;

 /// @brief Construct an ORC writer options with a sink and table.
 /// @param sink The sink to write to.
 /// @param table The table to write.
 explicit orc_writer_options(sink_info sink, table_view table)
 : _sink(std::move(sink)), _table(std::move(table))
 {
 }

 public
 /// @brief Default constructor.
 explicit orc_writer_options() = default;

 /// @brief Return a builder object to create the options.
 /// @param sink The sink to write to.
 /// @param table The table to write.
 /// @return orc_writer_options_builder The builder object.
 static orc_writer_options_builder builder(sink_info const& sink, table_view const& table);

 /// @brief Get the sink.
 [[nodiscard]] sink_info const& get_sink() const { return _sink; }

 /// @brief Get the compression type.
 [[nodiscard]] compression_type get_compression() const { return _compression; }

 /// @brief Whether statistics collection is enabled.
 {
 return _stats_freq != statistics_freq::STATISTICS_NONE;
 }

 /// @brief Get the frequency of statistics collection.
 [[nodiscard]] statistics_freq get_statistics_freq() const { return _stats_freq; }

 /// @brief Get the maximum size of each stripe (unless smaller than a single row group).
 [[nodiscard]] auto get_stripe_size_bytes() const { return _stripe_size_bytes; }

 /// @brief Get the maximum number of rows in stripe (unless smaller than a single row group).
 [[nodiscard]] auto get_stripe_size_rows() const { return _stripe_size_rows; }

 /// @brief Get the row index stride (maximum number of rows in each row group).
 {
 auto const unaligned_stride = std::min(_row_index_stride, get_stripe_size_rows());
 return unaligned_stride - unaligned_stride % 8;
 }

 /// @brief Get the table to write.
 [[nodiscard]] table_view get_table() const { return _table; }

 /// @brief Get the optional associated metadata.
 [[nodiscard]] auto const& get_metadata() const { return _metadata; }

 /// @brief Get optional key_value_metadata.
 {
 return _user_data;
 }

 /// @brief Get the optional compression statistics.
 {
 return _compression_stats;
 }

 /// @brief Get whether string dictionaries should be alphabetically sorted.
 [[nodiscard]] bool get_enable_dictionary_sort() const { return _enable_dictionary_sort; }

 // Setters

 /// @brief Set the compression type.
 /// @param comp The compression type.
 void set_compression(compression_type comp)
 {
 _compression = comp;
 if (comp == compression_type::AUTO) { _compression = compression_type::SNAPPY; }
 }

 /// @brief Enable or disable statistics collection and set frequency.
 /// @param val Statistics collection frequency.
 void enable_statistics(statistics_freq val) { _stats_freq = val; }

 size_type _stripe_size_rows = default_stripe_size_rows;
 // 行索引步长（每个行组中的最大行数）
 size_type _row_index_stride = default_row_index_stride;
 // 可选的相关元数据
 std::optional<table_input_metadata> _metadata;
 // 可选的页脚键值元数据
 std::map<std::string, std::string> _user_data;
 // 可选的压缩统计信息
 std::shared_ptr<writer_compression_statistics> _compression_stats;
 // 指定字符串字典是否应按字母排序
 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; }

 // Setter 方法

 void set_compression(compression_type comp)
 {
 _compression = comp;
 if (comp == compression_type::AUTO) { _compression = compression_type::SNAPPY; }
 }

 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");
 _stripe_size_bytes = size_bytes;
 }

 void set_stripe_size_rows(size_type size_rows)
 {
 CUDF_EXPECTS(size_rows >= 512, "最大条带大小不能小于 512");
 _stripe_size_rows = size_rows;
 }

 void set_row_index_stride(size_type stride)
 {
 CUDF_EXPECTS(stride >= 512, "行索引步长不能小于 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.set_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;
};
 // 组结束
} // 命名空间 io
} // 命名空间 CUDF_EXPORT cudf