parquet.hpp

转到此文件的文档。

/*
 * 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
 *
 * 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/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 <cudf/utilities/export.hpp>
#include <cudf/utilities/memory_resource.hpp>

#include <iostream>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

namespace CUDF_EXPORT cudf {
namespace io {
constexpr size_t default_row_group_size_bytes =
 std::numeric_limits<size_t>::max();
constexpr size_type default_row_group_size_rows = 1'000'000;
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;

/*! \brief Parquet 读取器选项
 *  \param src 数据源
 */
class parquet_reader_options {
 source_info _source;

 // 要读取的列在 schema 中的路径；`nullopt` 表示全部
 std::optional<std::vector<std::string>> _columns;

 // 要读取的行组列表（如果为空则忽略）
 std::vector<std::vector<size_type>> _row_groups;
 // 从开头跳过的行数；Parquet 将行数存储为 int64_t
 int64_t _skip_rows = 0;
 // 要读取的行数；`nullopt` 表示全部
 std::optional<size_type> _num_rows;

 // 用作抽象语法树 (AST) 的谓词过滤器，用于过滤输出行。
 std::optional<std::reference_wrapper<ast::expression const>> _filter;

 // 是否将字符串数据存储为分类类型
 bool _convert_strings_to_categories = false;
 // 是否使用 PANDAS 元数据加载列
 bool _use_pandas_metadata = true;
 // 是否读取并使用 ARROW schema
 bool _use_arrow_schema = true;
 // 是否允许从不匹配的 Parquet 文件中读取匹配的选中列。
 bool _allow_mismatched_pq_schemas = false;
 // 将时间戳列转换为特定类型
 data_type _timestamp_type{type_id::EMPTY};

 std::optional<std::vector<reader_column_schema>> _reader_column_schema;

 /*! \brief Parquet reader options
 *  \param src Data source
 */
 explicit parquet_reader_options(source_info src) : _source{std::move(src)} {}

 friend parquet_reader_options_builder;

 public
 /*! \brief 默认构造函数 */
 explicit parquet_reader_options() = default;

 /*! \brief 创建 builder 的工厂函数
 *  \param src 源
 */
 static parquet_reader_options_builder builder(source_info src);

 /*! \brief 返回数据源
 *  \return 源
 */
 [[nodiscard]] source_info const& get_source() const { return _source; }

 /*! \brief 返回是否将字符串转换为分类类型
 *  \return 如果启用则为 true
 */
 [[nodiscard]] bool is_enabled_convert_strings_to_categories() const
 {
 return _convert_strings_to_categories;
 }

 /*! \brief 返回是否使用 PANDAS 元数据
 *  \return 如果启用则为 true
 */
 [[nodiscard]] bool is_enabled_use_pandas_metadata() const { return _use_pandas_metadata; }

 /*! \brief 返回是否使用 ARROW schema
 *  \return 如果启用则为 true
 */
 [[nodiscard]] bool is_enabled_use_arrow_schema() const { return _use_arrow_schema; }

 /*! \brief 返回是否允许从不匹配的 Parquet 文件中读取匹配的选中列。
 *  \return 如果启用则为 true
 */
 [[nodiscard]] bool is_enabled_allow_mismatched_pq_schemas() const
 {
 return _allow_mismatched_pq_schemas;
 }

 /*! \brief 返回请求的列 schema
 *  \return 请求的列 schema
 */
 [[nodiscard]] std::optional<std::vector<reader_column_schema>> get_column_schema() const
 {
 return _reader_column_schema;
 }

 /*! \brief 返回要跳过的起始行数
 *  \return 行数
 */
 [[nodiscard]] int64_t get_skip_rows() const { return _skip_rows; }

 /*! \brief 返回要读取的行数
 *  \return 行数
 */
 [[nodiscard]] std::optional<size_type> const& get_num_rows() const { return _num_rows; }

 /*! \brief 返回要读取的列列表
 *  \return 列名列表
 */
 [[nodiscard]] auto const& get_columns() const { return _columns; }

 /*! \brief 返回要读取的行组列表
 *  \return 行组索引列表
 */
 [[nodiscard]] auto const& get_row_groups() const { return _row_groups; }

 /*! \brief 返回谓词过滤器
 *  \return 谓词过滤器
 */
 [[nodiscard]] auto const& get_filter() const { return _filter; }

 /*! \brief 返回时间戳类型
 *  \return 时间戳类型
 */
 [[nodiscard]] data_type get_timestamp_type() const { return _timestamp_type; }

 /*! \brief 设置要读取的列
 *  \param col_names 列名列表
 */
 void set_columns(std::vector<std::string> col_names) { _columns = std::move(col_names); }

 /*! \brief 设置要读取的单个行组
 *  \param row_groups 行组索引列表
 */
 void set_row_groups(std::vector<std::vector<size_type>> row_groups);

 /*! \brief 设置应用于输出行的过滤器
 *  \param filter 作为 AST 的谓词过滤器
 */
 void set_filter(ast::expression const& filter) { _filter = filter; }

 /*! \brief 启用/禁用将字符串转换为分类类型
 *  \param val true 为启用，false 为禁用
 */
 void enable_convert_strings_to_categories(bool val) { _convert_strings_to_categories = val; }

 /*! \brief 启用/禁用使用 PANDAS 元数据
 *  \param val true 为启用，false 为禁用
 */
 void enable_use_pandas_metadata(bool val) { _use_pandas_metadata = val; }

 /*! \brief 启用/禁用使用 ARROW schema
 *  \param val true 为启用，false 为禁用
 */
 void enable_use_arrow_schema(bool val) { _use_arrow_schema = val; }

 /*! \brief 启用/禁用允许从不匹配的 Parquet 文件中读取匹配的选中列。
 *  \param val true 为启用，false 为禁用
 */
 void enable_allow_mismatched_pq_schemas(bool val) { _allow_mismatched_pq_schemas = val; }

 /*! \brief 设置自定义列 schema
 *  \param val 自定义列 schema
 */
 void set_column_schema(std::vector<reader_column_schema> val)
 {
 _reader_column_schema = std::move(val);
 }

 /*! \brief 设置要跳过的起始行数
 *  \param val 行数
 */
 void set_skip_rows(int64_t val);

 /*! \brief 设置要读取的行数
 *  \param val 行数
 */
 void set_num_rows(size_type val);

 /*! \brief 设置时间戳类型
 *  \param type 时间戳类型
 */
 void set_timestamp_type(data_type type) { _timestamp_type = type; }
};

/*! \brief Parquet 读取器选项 builder */
class parquet_reader_options_builder {
 parquet_reader_options options;

 public
 /*! \brief 默认构造函数 */
 parquet_reader_options_builder() = default;

 /*! \brief 带源的构造函数
 *  \param src 数据源
 */
 explicit parquet_reader_options_builder(source_info src) : options{std::move(src)} {}

 /*! \brief 设置要读取的列
 *  \param col_names 列名列表
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& columns(std::vector<std::string> col_names)
 {
 options._columns = std::move(col_names);
 return *this;
 }

 /*! \brief 设置要读取的单个行组
 *  \param row_groups 行组索引列表
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& row_groups(std::vector<std::vector<size_type>> row_groups)
 {
 options.set_row_groups(std::move(row_groups));
 return *this;
 }

 /*! \brief 设置应用于输出行的过滤器
 *  \param filter 作为 AST 的谓词过滤器
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& filter(ast::expression const& filter)
 {
 options.set_filter(filter);
 return *this;
 }

 /*! \brief 启用/禁用将字符串转换为分类类型
 *  \param val true 为启用，false 为禁用
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& convert_strings_to_categories(bool val)
 {
 options._convert_strings_to_categories = val;
 return *this;
 }

 /*! \brief 启用/禁用使用 PANDAS 元数据
 *  \param val true 为启用，false 为禁用
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& use_pandas_metadata(bool val)
 {
 options._use_pandas_metadata = val;
 return *this;
 }

 /*! \brief 启用/禁用使用 ARROW schema
 *  \param val true 为启用，false 为禁用
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& use_arrow_schema(bool val)
 {
 options._use_arrow_schema = val;
 return *this;
 }

 /*! \brief 启用/禁用允许从不匹配的 Parquet 文件中读取匹配的选中列。
 *  \param val true 为启用，false 为禁用
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& allow_mismatched_pq_schemas(bool val)
 {
 options._allow_mismatched_pq_schemas = val;
 return *this;
 }

 /*! \brief 设置自定义列 schema
 *  \param val 自定义列 schema
 *  \return This builder object
 */
 parquet_reader_options_builder& set_column_schema(std::vector<reader_column_schema> val)
 {
 options._reader_column_schema = std::move(val);
 return *this;
 }

 /*! \brief 设置要跳过的起始行数
 *  \param val 行数
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& skip_rows(int64_t val)
 {
 options.set_skip_rows(val);
 return *this;
 }

 /*! \brief 设置要读取的行数
 *  \param val 行数
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& num_rows(size_type val)
 {
 options.set_num_rows(val);
 return *this;
 }

 /*! \brief 设置时间戳类型
 *  \param type 时间戳类型
 *  \return 此 builder 对象
 */
 parquet_reader_options_builder& timestamp_type(data_type type)
 {
 options._timestamp_type = type;
 return *this;
 }

 /*! \brief 转换为 options 对象
 *  \return parquet_reader_options 对象
 */
 operator parquet_reader_options&&() { return std::move(options); }

 /*! \brief 构建 options 对象
 *  \return parquet_reader_options 对象
 */
 parquet_reader_options&& build() { return std::move(options); }
};

 /*! \brief 读取 Parquet 数据集并返回带元数据的表。
 *
 *  \param options 控制读取行为的设置。
 *  \param stream 用于设备内存操作和内核启动的 CUDA 流。
 *  \param mr 用于分配返回的表及其列的设备内存资源。
 *  \return 表和元数据
 */
table_with_metadata read_parquet(
 parquet_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 分块读取 Parquet 数据集。
 */
class chunked_parquet_reader {
 public
 /*! \brief 默认构造函数 */
 chunked_parquet_reader();

 /*! \brief 带分块读取限制的构造函数。
 *
 *  \param chunk_read_limit 每块要读取的最大字节数。
 *  \param options 控制读取行为的设置。
 *  \param stream 用于设备内存操作和内核启动的 CUDA 流。
 *  \param mr 用于分配返回的表及其列的设备内存资源。
 */
 chunked_parquet_reader(
 std::size_t chunk_read_limit,
 parquet_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 带分块读取限制和 pass 读取限制的构造函数。
 *
 *  \param chunk_read_limit 每块要读取的最大字节数。
 *  \param pass_read_limit 每列每 pass 要读取的最大字节数。如果在 pass 内达到 chunk_read_limit，将使用 chunk_read_limit。否则，将使用 pass_read_limit。
 *  \param options 控制读取行为的设置。
 *  \param stream 用于设备内存操作和内核启动的 CUDA 流。
 *  \param mr 用于分配返回的表及其列的设备内存资源。
 */
 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::device_async_resource_ref mr = cudf::get_current_device_resource_ref());

 /*! \brief 析构函数 */
 ~chunked_parquet_reader();

 /*! \brief 如果还有更多块可读，则返回 true。
 *  \return 如果还有更多块可读，则为 true。
 */
 [[nodiscard]] bool has_next() const;

 /*! \brief 从 Parquet 数据集中读取下一块。
 *  \return 带元数据的表
 */
 [[nodiscard]] table_with_metadata read_chunk() const;

 private
 std::unique_ptr<cudf::io::parquet::detail::chunked_reader> reader;
};
 // group 结束
/*! \brief 用于排序的列的描述 */
struct sorting_column {
 int column_idx{};
 bool is_descending{false};
 bool is_nulls_first{true};
};

/*! \brief Parquet writer 选项 */
class parquet_writer_options_base {
 // 指定用于 writer 输出的 sink
 sink_info _sink;
 // 指定要使用的压缩格式
 compression_type _compression = compression_type::SNAPPY;
 // 指定输出文件中的统计信息级别
 statistics_freq _stats_level = statistics_freq::STATISTICS_ROWGROUP;
 // 可选的关联元数据
 std::optional<table_input_metadata> _metadata;
 // 可选的页脚 key_value_metadata
 std::vector<std::map<std::string, std::string>> _user_data;
 // Parquet writer 可以写入 INT96 或 TIMESTAMP_MICROS。默认为 TIMESTAMP_MICROS。
 // 如果为 true，则覆盖 _metadata 中的任何按列设置。
 bool _write_timestamps_as_int96 = false;
 // Parquet writer 可以将时间戳写入为 UTC
 // 默认为 true，因为 libcudf 时间戳隐式为 UTC
 bool _write_timestamps_as_UTC = true;
 BuilderT& sorting_columns(std::vector<sorting_column> sorting_columns);

 operator OptionsT&&();

 OptionsT&& build();
};

class parquet_writer_options_builder;

class parquet_writer_options : public parquet_writer_options_base {
 // Sets of columns to output
 table_view _table;
 // Partitions described as {start_row, num_rows} pairs
 std::vector<partition_info> _partitions;
 // Column chunks file paths to be set in the raw output metadata. One per output file
 std::vector<std::string> _column_chunks_file_paths;

 friend parquet_writer_options_builder;

 explicit parquet_writer_options(sink_info const& sink, table_view table);

 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]] table_view get_table() const { return _table; }

 [[nodiscard]] std::vector<partition_info> const& get_partitions() const { return _partitions; }

 [[nodiscard]] std::vector<std::string> const& get_column_chunks_file_paths() const
 {
 return _column_chunks_file_paths;
 }

 void set_partitions(std::vector<partition_info> partitions);

 void set_column_chunks_file_paths(std::vector<std::string> file_paths);
};

class parquet_writer_options_builder
 : public parquet_writer_options_builder_base<parquet_writer_options_builder,
 parquet_writer_options> {
 public
 explicit parquet_writer_options_builder() = default;

 explicit parquet_writer_options_builder(sink_info const& sink, table_view const& table);

 parquet_writer_options_builder& partitions(std::vector<partition_info> partitions);

 parquet_writer_options_builder& column_chunks_file_paths(std::vector<std::string> file_paths);
};

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 : public parquet_writer_options_base {
 explicit chunked_parquet_writer_options(sink_info sink);

 friend chunked_parquet_writer_options_builder;

 public
 chunked_parquet_writer_options() = default;

 static chunked_parquet_writer_options_builder builder(sink_info const& sink);
};

class chunked_parquet_writer_options_builder
 : public parquet_writer_options_builder_base<chunked_parquet_writer_options_builder,
 chunked_parquet_writer_options> {
 public
 chunked_parquet_writer_options_builder() = default;

 chunked_parquet_writer_options_builder(sink_info const& sink);
};

class parquet_chunked_writer {
 public
 parquet_chunked_writer();

 parquet_chunked_writer(chunked_parquet_writer_options const& options,
 rmm::cuda_stream_view stream = cudf::get_default_stream());
 ~parquet_chunked_writer();

 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 io
} // namespace CUDF_EXPORT cudf