avro.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 "types.hpp"
 
#include <cudf/table/table_view.hpp>
#include <cudf/types.hpp>
 
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <memory>
#include <string>
#include <vector>
 
namespace cudf {
namespace io {
class avro_reader_options_builder;
 
class avro_reader_options {
  source_info _source;
 
  // Names of column to read; empty is all
  std::vector<std::string> _columns;
 
  // Rows to skip from the start;
  size_type _skip_rows = 0;
  // Rows to read; -1 is all
  size_type _num_rows = -1;
 
  explicit avro_reader_options(source_info src) : _source{std::move(src)} {}
 
  friend avro_reader_options_builder;
 
 public:
  avro_reader_options() = default;
 
  [[nodiscard]] source_info const& get_source() const { return _source; }
 
  [[nodiscard]] std::vector<std::string> get_columns() const { return _columns; }
 
  [[nodiscard]] size_type get_skip_rows() const { return _skip_rows; }
 
  [[nodiscard]] size_type get_num_rows() const { return _num_rows; }
 
  void set_columns(std::vector<std::string> col_names) { _columns = std::move(col_names); }
 
  void set_skip_rows(size_type val) { _skip_rows = val; }
 
  void set_num_rows(size_type val) { _num_rows = val; }
 
  static avro_reader_options_builder builder(source_info src);
};
 
class avro_reader_options_builder {
  avro_reader_options options;
 
 public:
  avro_reader_options_builder() = default;
 
  explicit avro_reader_options_builder(source_info src) : options{std::move(src)} {}
 
  avro_reader_options_builder& columns(std::vector<std::string> col_names)
  {
    options._columns = std::move(col_names);
    return *this;
  }
 
  avro_reader_options_builder& skip_rows(size_type val)
  {
    options._skip_rows = val;
    return *this;
  }
 
  avro_reader_options_builder& num_rows(size_type val)
  {
    options._num_rows = val;
    return *this;
  }
 
  operator avro_reader_options&&() { return std::move(options); }
 
  avro_reader_options&& build() { return std::move(options); }
};
 
table_with_metadata read_avro(
  avro_reader_options const& options,
  rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource());
  // end of group
}  // namespace io
}  // namespace cudf