avro.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 "types.hpp"

#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
#include <cudf/utilities/memory_resource.hpp>

#include <memory>
#include <string>
#include <vector>

namespace CUDF_EXPORT cudf {
namespace io {
/// @brief 用于构建 read_avro() 选项的构建器。
class avro_reader_options_builder;
/// @brief 用于 read_avro() 的设置。
class avro_reader_options {
 source_info _source;

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

 // 从开始跳过的行数；
 size_type _skip_rows = 0;
 // 要读取的行数；-1 表示所有行
 size_type _num_rows = -1;
 /// @brief 构造函数。
 explicit avro_reader_options(source_info src) : _source{std::move(src)} {}

 friend avro_reader_options_builder;

 /// @brief 默认构造函数。
 avro_reader_options() = default;
 /// @brief 返回源信息。
 [[nodiscard]] source_info const& get_source() const { return _source; }
 /// @brief 返回要读取的列名。
 [[nodiscard]] std::vector<std::string> get_columns() const { return _columns; }
 /// @brief 返回从开始跳过的行数。
 [[nodiscard]] size_type get_skip_rows() const { return _skip_rows; }
 /// @brief 返回要读取的行数。
 [[nodiscard]] size_type get_num_rows() const { return _num_rows; }
 /// @brief 设置要读取的列名。
 void set_columns(std::vector<std::string> col_names) { _columns = std::move(col_names); }
 /// @brief 设置要跳过的行数。
 void set_skip_rows(size_type val) { _skip_rows = val; }
 /// @brief 设置要读取的行数。
 void set_num_rows(size_type val) { _num_rows = val; }
 /// @brief 创建将构建 avro_reader_options 的 avro_reader_options_builder。
 static avro_reader_options_builder builder(source_info src);
};
/// @brief 用于构建 read_avro() 选项的构建器。
class avro_reader_options_builder {
 avro_reader_options options;

 /// @brief 默认构造函数。
 avro_reader_options_builder() = default;
 /// @brief 从源信息构造函数。
 explicit avro_reader_options_builder(source_info src) : options{std::move(src)} {}
 /// @brief 设置要读取的列名。
 avro_reader_options_builder& columns(std::vector<std::string> col_names)
 {
 options._columns = std::move(col_names);
 return *this;
 }
 /// @brief 设置要跳过的行数。
 avro_reader_options_builder& skip_rows(size_type val)
 {
 options._skip_rows = val;
 return *this;
 }
 /// @brief 设置要读取的行数。
 avro_reader_options_builder& num_rows(size_type val)
 {
 options._num_rows = val;
 return *this;
 }
 /// @brief 构建完成后移动 avro_reader_options 成员。
 operator avro_reader_options&&() { return std::move(options); }
 /// @brief 构建完成后移动 avro_reader_options 成员。
 avro_reader_options&& build() { return std::move(options); }
};
/// @brief 将 Avro 数据集读取到一组列中。
table_with_metadata read_avro(
 avro_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());
 // end of group
} // namespace io