diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp index cf2ba369023..d2026473b6c 100644 --- a/cpp/src/io/utilities/datasource.cpp +++ b/cpp/src/io/utilities/datasource.cpp @@ -18,7 +18,6 @@ #include "io/utilities/config_utils.hpp" #include -#include #include #include #include @@ -27,7 +26,6 @@ #include -#include #include #include #include @@ -338,6 +336,33 @@ class device_buffer_source final : public datasource { cudf::device_span _d_buffer; ///< A non-owning view of the existing device data }; +// zero-copy host buffer source +class host_buffer_source final : public datasource { + public: + explicit host_buffer_source(cudf::host_span h_buffer) : _h_buffer{h_buffer} {} + + size_t host_read(size_t offset, size_t size, uint8_t* dst) override + { + auto const count = std::min(size, this->size() - offset); + std::memcpy(dst, _h_buffer.data() + offset, count); + return count; + } + + std::unique_ptr host_read(size_t offset, size_t size) override + { + auto const count = std::min(size, this->size() - offset); + return std::make_unique( + reinterpret_cast(_h_buffer.data() + offset), count); + } + + [[nodiscard]] bool supports_device_read() const override { return false; } + + [[nodiscard]] size_t size() const override { return _h_buffer.size(); } + + private: + cudf::host_span _h_buffer; ///< A non-owning view of the existing host data +}; + /** * @brief Wrapper class for user implemented data sources * @@ -424,9 +449,7 @@ std::unique_ptr datasource::create(host_buffer const& buffer) std::unique_ptr datasource::create(cudf::host_span buffer) { - // Use Arrow IO buffer class for zero-copy reads of host memory - return std::make_unique(std::make_shared( - reinterpret_cast(buffer.data()), buffer.size())); + return std::make_unique(buffer); } std::unique_ptr datasource::create(cudf::device_span buffer)