Skip to content

Do not zero-fill the kvikio_ofstream staging buffer #2749

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
rapids-bot merged 2 commits into from
Oct 8, 2026
Merged

Do not zero-fill the kvikio_ofstream staging buffer #2749

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a82873d
Do not zero-fill the kvikio_ofstream staging buffer
bdice Oct 5, 2026
a7557dc
Merge branch 'main' into kvikio-ofstream-no-zero-fill
bdice Oct 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 15 additions & 7 deletions cpp/src/util/file_io.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,8 @@
#include <cstring>
#include <exception>
#include <limits>
#include <memory>
#include <sys/stat.h>
#include <vector>

namespace cuvs::util {
namespace {
Expand Down Expand Up @@ -262,10 +262,14 @@ void kvikio_file_reader::read_device(void* data, size_t size) { impl_->read_devi
class kvikio_ofstream::sbuf : public std::streambuf {
public:
sbuf(const std::string& path, size_t cap)
try : path_(path), handle_(path, "w"), buffer_(std::max<size_t>(cap, kNumpyDataAlignment)) {
RAFT_EXPECTS(buffer_.size() <= static_cast<size_t>(std::numeric_limits<int>::max()),
try
: path_(path),
handle_(path, "w"),
buffer_size_(std::max<size_t>(cap, kNumpyDataAlignment)),
buffer_(std::make_unique_for_overwrite<char[]>(buffer_size_)) {
RAFT_EXPECTS(buffer_size_ <= static_cast<size_t>(std::numeric_limits<int>::max()),
"kvikio_ofstream buffer size must fit in std::streambuf::pbump");
setp(buffer_.data(), buffer_.data() + buffer_.size());
setp(buffer_.get(), buffer_.get() + buffer_size_);
} catch (const std::exception& e) {
RAFT_FAIL("Cannot open file %s for writing: %s", path.c_str(), e.what());
}
Expand Down Expand Up @@ -338,7 +342,7 @@ class kvikio_ofstream::sbuf : public std::streambuf {
while (remaining > 0) {
// If the caller hands us a large contiguous chunk, flush pending staged bytes and pass the
// chunk straight to kvikio. This avoids std::streambuf's byte-at-a-time fallback path.
if (remaining >= buffer_.size()) {
if (remaining >= buffer_size_) {
flush_buffer();
write_at_current_offset(handle_, current, remaining);
return requested;
Expand Down Expand Up @@ -385,7 +389,7 @@ class kvikio_ofstream::sbuf : public std::streambuf {
const size_t n = static_cast<size_t>(pptr() - pbase());
if (n > 0) {
write_at_current_offset(handle_, pbase(), n);
setp(buffer_.data(), buffer_.data() + buffer_.size());
setp(buffer_.get(), buffer_.get() + buffer_size_);
}
}

Expand All @@ -400,7 +404,11 @@ class kvikio_ofstream::sbuf : public std::streambuf {
std::string path_;
kvikio::FileHandle handle_;
std::unique_ptr<kvikio::FileHandle> device_handle_;
std::vector<char> buffer_;
// Staging buffer, deliberately left uninitialized: only [pbase(), pptr()) is ever passed to
// kvikio, and those bytes are always written first. Value-initializing it would fault in and
// zero every page of the (32 MiB by default) buffer on each open, even for tiny files.
size_t buffer_size_;
std::unique_ptr<char[]> buffer_;
size_t offset_ = 0;
bool closed_ = false;
};
Expand Down
Loading