datasource.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 <cudf/io/types.hpp>
#include <cudf/utilities/error.hpp>
#include <cudf/utilities/export.hpp>
#include <cudf/utilities/span.hpp>
 
#include <rmm/cuda_stream_view.hpp>
 
#include <future>
#include <memory>
 
namespace CUDF_EXPORT cudf {
namespace io {
 
class datasource {
 public:
  template <typename Container>
  class owning_buffer;  // forward declaration
  class buffer {
   public:
    [[nodiscard]] virtual size_t size() const = 0;
 
    [[nodiscard]] virtual uint8_t const* data() const = 0;
 
    virtual ~buffer() {}
 
    template <typename Container>
    static std::unique_ptr<buffer> create(Container&& data_owner)
    {
      return std::make_unique<owning_buffer<Container>>(std::forward<Container>(data_owner));
    }
  };
 
  static std::unique_ptr<datasource> create(std::string const& filepath,
                                            size_t offset            = 0,
                                            size_t max_size_estimate = 0);
 
  static std::unique_ptr<datasource> create(host_buffer const& buffer);
 
  static std::unique_ptr<datasource> create(cudf::host_span<std::byte const> buffer);
 
  static std::unique_ptr<datasource> create(cudf::device_span<std::byte const> buffer);
 
  static std::unique_ptr<datasource> create(datasource* source);
 
  template <typename T>
  static std::vector<std::unique_ptr<datasource>> create(std::vector<T> const& args)
  {
    std::vector<std::unique_ptr<datasource>> sources;
    sources.reserve(args.size());
    std::transform(args.cbegin(), args.cend(), std::back_inserter(sources), [](auto const& arg) {
      return datasource::create(arg);
    });
    return sources;
  }
 
  virtual ~datasource(){};
 
  virtual std::unique_ptr<datasource::buffer> host_read(size_t offset, size_t size) = 0;
 
  virtual size_t host_read(size_t offset, size_t size, uint8_t* dst) = 0;
 
  [[nodiscard]] virtual bool supports_device_read() const { return false; }
 
  [[nodiscard]] virtual bool is_device_read_preferred(size_t size) const
  {
    return supports_device_read();
  }
 
  virtual std::unique_ptr<datasource::buffer> device_read(size_t offset,
                                                          size_t size,
                                                          rmm::cuda_stream_view stream)
  {
    CUDF_FAIL("datasource classes that support device_read must override it.");
  }
 
  virtual size_t device_read(size_t offset, size_t size, uint8_t* dst, rmm::cuda_stream_view stream)
  {
    CUDF_FAIL("datasource classes that support device_read must override it.");
  }
 
  virtual std::future<size_t> device_read_async(size_t offset,
                                                size_t size,
                                                uint8_t* dst,
                                                rmm::cuda_stream_view stream)
  {
    CUDF_FAIL("datasource classes that support device_read_async must override it.");
  }
 
  [[nodiscard]] virtual size_t size() const = 0;
 
  [[nodiscard]] virtual bool is_empty() const { return size() == 0; }
 
  class non_owning_buffer : public buffer {
   public:
    non_owning_buffer() {}
 
    non_owning_buffer(uint8_t const* data, size_t size) : _data(data), _size(size) {}
 
    [[nodiscard]] size_t size() const override { return _size; }
 
    [[nodiscard]] uint8_t const* data() const override { return _data; }
 
   private:
    uint8_t const* _data{nullptr};
    size_t _size{0};
  };
 
  template <typename Container>
  class owning_buffer : public buffer {
   public:
    // Require that the argument passed to the constructor be an rvalue (Container&& being an rvalue
    // reference).
    static_assert(std::is_rvalue_reference_v<Container&&>,
                  "The container argument passed to the constructor must be an rvalue.");
 
    owning_buffer(Container&& moved_data_owner)
      : _data(std::move(moved_data_owner)), _data_ptr(_data.data()), _size(_data.size())
    {
    }
 
    owning_buffer(Container&& moved_data_owner, uint8_t const* data_ptr, size_t size)
      : _data(std::move(moved_data_owner)), _data_ptr(data_ptr), _size(size)
    {
    }
 
    [[nodiscard]] size_t size() const override { return _size; }
 
    [[nodiscard]] uint8_t const* data() const override
    {
      return static_cast<uint8_t const*>(_data_ptr);
    }
 
   private:
    Container _data;
    void const* _data_ptr;
    size_t _size;
  };
};
  // end of group
}  // namespace io
}  // namespace CUDF_EXPORT cudf