NVIDIA / NVIDIA/cudf

[BUG] Limit size of buffer read by batched multi-source JSON lines reader to be at most `INT_MAX` bytes

Open
#17,058 0 comments 0 reactions 0 assignees View on GitHub
bug cuIO libcudf
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Describe the bug**
With the implementation of the reallocate-and-retry logic when the initial buffer size estimate fails for byte range reading (PR #16687), the total buffer size read per batch can exceed 2GiB.

**Steps/Code to reproduce bug**
```
TEST_F(JsonLargeReaderTest, MultiBatchFailing)
{
std::string json_string = R"(
{ "a": { "y" : 6}, "b" : [1, 2, 3], "c": 11 }
{ "a": { "y" : 6}, "b" : [4, 5 ], "c": 12 }
{ "a": { "y" : 6}, "b" : [6 ], "c": 13 }
{ "a": { "y" : 6}, "b" : [7 ], "c": 14 }
)";

std::size_t const batch_size_upper_bound = std::numeric_limits::max() / 16;
// set smaller batch_size to reduce file size and execution time
setenv("LIBCUDF_JSON_BATCH_SIZE", std::to_string(batch_size_upper_bound).c_str(), 1);

constexpr std::size_t expected_file_size = 1.0 * static_cast(batch_size_upper_bound);
std::size_t log_repetitions =
static_cast(std::floor(std::log2(expected_file_size / json_string.size())));

json_string.reserve(json_string.size() * (1UL << log_repetitions));
for (std::size_t i = 0; i < log_repetitions; i++) {
json_string += json_string;
}

std::string really_long_string = R"(haha)";
log_repetitions += 4;
really_long_string.reserve(really_long_string.size() * (1UL << log_repetitions));
for (std::size_t i = 0; i < log_repetitions; i++) {
really_long_string += really_long_string;
}
std::string last_line = "{ \"lol\" : \"" + really_long_string + "\" }";
json_string += last_line;

constexpr int num_sources = 1;
std::vector> hostbufs(
num_sources,
cudf::host_span(reinterpret_cast(json_string.data()),
json_string.size()));

// Initialize parsing options (reading json lines)
cudf::io::json_reader_options json_lines_options =
cudf::io::json_reader_options::builder(
cudf::io::source_info{
cudf::host_span>(hostbufs.data(), hostbufs.size())})
.lines(true)
.compression(cudf::io::compression_type::NONE)
.recovery_mode(cudf::io::json_recovery_mode_t::FAIL);

// Read full test data via existing, nested JSON lines reader
cudf::io::table_with_metadata current_reader_table = cudf::io::read_json(json_lines_options);

std::vector> datasources;
for (auto& hb : hostbufs) {
datasources.emplace_back(cudf::io::datasource::create(hb));
}
// Test for different chunk sizes
std::vector chunk_sizes{batch_size_upper_bound / 4,
batch_size_upper_bound / 2,
batch_size_upper_bound,
static_cast(batch_size_upper_bound * 2)};

for (auto chunk_size : chunk_sizes) {
auto const tables =
split_byte_range_reading(datasources,
json_lines_options,
chunk_size,
cudf::get_default_stream(),
cudf::get_current_device_resource_ref());

auto table_views = std::vector(tables.size());
std::transform(tables.begin(), tables.end(), table_views.begin(), [](auto& table) {
return table.tbl->view();
});
auto result = cudf::concatenate(table_views);

// Verify that the data read via chunked reader matches the data read via nested JSON reader
// cannot use EQUAL due to concatenate removing null mask
CUDF_TEST_EXPECT_TABLES_EQUIVALENT(current_reader_table.tbl->view(), result->view());
}

// go back to normal batch_size
unsetenv("LIBCUDF_JSON_BATCH_SIZE");
}
```
diff file:
```
diff --git a/cpp/src/io/json/read_json.cu b/cpp/src/io/json/read_json.cu
index 50f0e5cf33..9aeeb2a0f1 100644
--- a/cpp/src/io/json/read_json.cu
+++ b/cpp/src/io/json/read_json.cu
@@ -204,6 +204,7 @@ datasource::owning_buffer get_record_range_raw_input(
}
}

+ CUDF_EXPECTS(static_cast(next_delim_pos - first_delim_pos - shift_for_nonzero_offset) < get_batch_size_upper_bound(), "oh no!");
return datasource::owning_buffer(
std::move(buffer),
reinterpret_cast(buffer.data()) + first_delim_pos + shift_for_nonzero_offset,

```

**Expected behavior**
The size of the buffer returned by `get_record_range_raw_input` should not exceed 2GiB.
Proposed solution: stop at the line end before the 2GB limit and adjust the batch offsets for the remaining batches.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.