Skip to content

Commit 2e7ba80

Browse files
drymancopybara-github
dryman
authored and
copybara-github
committed
Add new writer option to reduce IOPs for small groups.
Also change the default group size to 1 as this is the most common usecase. PiperOrigin-RevId: 606995301
1 parent ee89b74 commit 2e7ba80

8 files changed

Lines changed: 113 additions & 38 deletions

‎cpp/array_record_reader_test.cc‎

Lines changed: 43 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,16 @@ class ArrayRecordReaderTest
8282

8383
TEST_P(ArrayRecordReaderTest, MoveTest) {
8484
std::string encoded;
85-
auto writer_options = GetWriterOptions().set_group_size(2);
85+
auto writer_options = GetWriterOptions();
86+
int32_t group_size;
87+
if (optimize_for_random_access()) {
88+
group_size = 1;
89+
writer_options.set_groups_awaiting_flush(256);
90+
} else {
91+
group_size = 3;
92+
writer_options.set_groups_awaiting_flush(1);
93+
}
94+
writer_options.set_group_size(group_size);
8695
auto writer = ArrayRecordWriter(
8796
riegeli::Maker<riegeli::StringWriter>(&encoded), writer_options, nullptr);
8897

@@ -113,7 +122,7 @@ TEST_P(ArrayRecordReaderTest, MoveTest) {
113122
})
114123
.ok());
115124

116-
EXPECT_EQ(reader_before_move.RecordGroupSize(), 2);
125+
EXPECT_EQ(reader_before_move.RecordGroupSize(), group_size);
117126

118127
ArrayRecordReader reader = std::move(reader_before_move);
119128
// Once a reader is moved, it is closed.
@@ -159,7 +168,7 @@ TEST_P(ArrayRecordReaderTest, MoveTest) {
159168
EXPECT_FALSE(reader.ReadRecord(&record_view));
160169
EXPECT_TRUE(reader.ok());
161170

162-
EXPECT_EQ(reader.RecordGroupSize(), 2);
171+
EXPECT_EQ(reader.RecordGroupSize(), group_size);
163172

164173
ASSERT_TRUE(reader.Close());
165174
}
@@ -172,11 +181,21 @@ TEST_P(ArrayRecordReaderTest, RandomDatasetTest) {
172181
size_t len = dist(bitgen);
173182
records[i] = MTRandomBytes(bitgen, len);
174183
}
184+
auto writer_options = GetWriterOptions();
185+
int32_t group_size;
186+
if (optimize_for_random_access()) {
187+
group_size = 1;
188+
writer_options.set_groups_awaiting_flush(1024);
189+
} else {
190+
group_size = 1024;
191+
writer_options.set_groups_awaiting_flush(1);
192+
}
193+
writer_options.set_group_size(group_size);
175194

176195
std::string encoded;
177196
auto writer =
178197
ArrayRecordWriter(riegeli::Maker<riegeli::StringWriter>(&encoded),
179-
GetWriterOptions(), get_pool());
198+
writer_options, get_pool());
180199
for (auto i : Seq(kDatasetSize)) {
181200
EXPECT_TRUE(writer.WriteRecord(records[i]));
182201
}
@@ -193,58 +212,65 @@ TEST_P(ArrayRecordReaderTest, RandomDatasetTest) {
193212
reader_opt, use_thread_pool() ? get_pool() : nullptr);
194213
ASSERT_TRUE(reader.status().ok());
195214
EXPECT_EQ(reader.NumRecords(), kDatasetSize);
196-
uint64_t group_size =
197-
std::min(ArrayRecordWriterBase::Options::kDefaultGroupSize, kDatasetSize);
198215
EXPECT_EQ(reader.RecordGroupSize(), group_size);
199216

200-
std::vector<bool> read_all_records(kDatasetSize, false);
217+
std::vector<int32_t> read_all_records(kDatasetSize, 0);
201218
ASSERT_TRUE(reader
202219
.ParallelReadRecords(
203220
[&](uint64_t record_index,
204221
absl::string_view result_view) -> absl::Status {
205222
EXPECT_EQ(result_view, records[record_index]);
206223
EXPECT_FALSE(read_all_records[record_index]);
207-
read_all_records[record_index] = true;
224+
read_all_records[record_index] = 1;
208225
return absl::OkStatus();
209226
})
210227
.ok());
211-
for (bool record_was_read : read_all_records) {
212-
EXPECT_TRUE(record_was_read);
228+
uint32_t records_read = 0;
229+
for (auto record_was_read : read_all_records) {
230+
if (record_was_read) {
231+
records_read++;
232+
}
213233
}
234+
EXPECT_EQ(records_read, kDatasetSize);
235+
EXPECT_TRUE(reader.SeekRecord(0));
214236

215237
std::vector<uint64_t> indices = {0, 3, 5, 7, 101, 2000};
216-
std::vector<bool> read_indexed_records(indices.size(), false);
238+
std::vector<int32_t> read_indexed_records(indices.size(), 0);
217239
ASSERT_TRUE(reader
218240
.ParallelReadRecordsWithIndices(
219241
indices,
220242
[&](uint64_t indices_idx,
221243
absl::string_view result_view) -> absl::Status {
222244
EXPECT_EQ(result_view, records[indices[indices_idx]]);
223245
EXPECT_FALSE(read_indexed_records[indices_idx]);
224-
read_indexed_records[indices_idx] = true;
246+
read_indexed_records[indices_idx] = 1;
225247
return absl::OkStatus();
226248
})
227249
.ok());
228-
for (bool record_was_read : read_indexed_records) {
250+
for (auto record_was_read : read_indexed_records) {
229251
EXPECT_TRUE(record_was_read);
230252
}
231253

232254
uint64_t begin = 10, end = 101;
233-
std::vector<bool> read_range_records(end - begin, false);
255+
std::vector<int32_t> read_range_records(end - begin, 0);
234256
ASSERT_TRUE(reader
235257
.ParallelReadRecordsInRange(
236258
begin, end,
237259
[&](uint64_t record_index,
238260
absl::string_view result_view) -> absl::Status {
239261
EXPECT_EQ(result_view, records[record_index]);
240262
EXPECT_FALSE(read_range_records[record_index - begin]);
241-
read_range_records[record_index - begin] = true;
263+
read_range_records[record_index - begin] = 1;
242264
return absl::OkStatus();
243265
})
244266
.ok());
245-
for (bool record_was_read : read_range_records) {
246-
EXPECT_TRUE(record_was_read);
267+
records_read = 0;
268+
for (auto record_was_read : read_range_records) {
269+
if (record_was_read) {
270+
records_read++;
271+
}
247272
}
273+
EXPECT_EQ(records_read, end - begin);
248274

249275
// Test sequential read
250276
absl::string_view result_view;

‎cpp/array_record_writer.cc‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,9 @@ ArrayRecordWriterBase::Options::FromString(absl::string_view text) {
131131
// Group
132132
options_parser.AddOption(
133133
"group_size", ValueParser::Int(1, INT32_MAX, &options.group_size_));
134+
options_parser.AddOption(
135+
"groups_awaiting_flush",
136+
ValueParser::Int(0, INT32_MAX, &options.groups_awaiting_flush_));
134137
int32_t max_parallelism = 0;
135138
options_parser.AddOption(
136139
"max_parallelism",
@@ -196,6 +199,7 @@ ArrayRecordWriterBase::Options::FromString(absl::string_view text) {
196199
std::string ArrayRecordWriterBase::Options::ToString() const {
197200
std::string option;
198201
absl::StrAppend(&option, "group_size:", this->group_size_,
202+
",groups_awaiting_flush:", this->groups_awaiting_flush_,
199203
",transpose:", this->transpose_ ? "true" : "false",
200204
",pad_to_block_boundary:",
201205
this->pad_to_block_boundary_ ? "true" : "false");
@@ -362,6 +366,7 @@ void ArrayRecordWriterBase::Initialize() {
362366
}
363367
// Add callback only after we serialize header and metadata.
364368
writer->set_submit_chunk_callback(submit_chunk_callback_.get());
369+
// writer->set_chunks_awaiting_flush(options_.groups_awaiting_flush());
365370
}
366371

367372
void ArrayRecordWriterBase::Done() {

‎cpp/array_record_writer.h‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,6 +94,7 @@ class ArrayRecordWriterBase : public riegeli::Object {
9494
// options ::= option? ("," option?)*
9595
// option ::=
9696
// "group_size" ":" group_size |
97+
// "groups_awaiting_flush" ":" groups_awaiting_flush |
9798
// "max_parallelism" ":" max_parallelism |
9899
// "saturation_delay_ms" : saturation_delay_ms |
99100
// "uncompressed" |
@@ -105,7 +106,11 @@ class ArrayRecordWriterBase : public riegeli::Object {
105106
// "window_log" : window_log |
106107
// "pad_to_block_boundary" (":" ("true" | "false"))?
107108
// group_size ::= positive integer which specifies number of records to be
108-
// grouped into a chunk before compression. (default 65536)
109+
// grouped into a chunk before compression. (default 1)
110+
// groups_awaiting_flush ::= positive integer which specify the number of
111+
// groups stored in a write buffer before sending to the storage IO.
112+
// This option reduce the total IOPs for small group_size. (default
113+
// 1024)
109114
// saturation_delay_ms ::= positive integer which specifies a delay in
110115
// milliseconds when the parallel writing queue is saturated.
111116
// max_parallelism ::= `auto` or positive integers which specifies
@@ -123,13 +128,23 @@ class ArrayRecordWriterBase : public riegeli::Object {
123128
//
124129
// The larger the value, the denser the file, at the cost of more expansive
125130
// random accessing.
126-
static constexpr uint32_t kDefaultGroupSize = 65536;
131+
// TODO(fchern) group_size=1 can trigger writer bugs..
132+
static constexpr uint32_t kDefaultGroupSize = 1;
127133
Options& set_group_size(uint32_t group_size) {
128134
group_size_ = group_size;
129135
return *this;
130136
}
131137
uint32_t group_size() const { return group_size_; }
132138

139+
// Set the number of gruops pending in the write buffer before sending to
140+
// the storage IO. This option reduces the total IOPs for small group size.
141+
static constexpr uint32_t kDefaultGroupsAwaitingFlush = 1024;
142+
Options& set_groups_awaiting_flush(uint32_t groups_awaiting_flush) {
143+
groups_awaiting_flush_ = groups_awaiting_flush;
144+
return *this;
145+
}
146+
uint32_t groups_awaiting_flush() const { return groups_awaiting_flush_; }
147+
133148
// Specifies max number of concurrent chunk encoders allowed. Default to the
134149
// thread pool size.
135150
Options& set_max_parallelism(std::optional<uint32_t> max_parallelism) {
@@ -290,6 +305,7 @@ class ArrayRecordWriterBase : public riegeli::Object {
290305

291306
private:
292307
int32_t group_size_ = kDefaultGroupSize;
308+
int32_t groups_awaiting_flush_ = kDefaultGroupsAwaitingFlush;
293309
riegeli::CompressorOptions compressor_options_;
294310
std::optional<riegeli::RecordsMetadata> metadata_;
295311
bool pad_to_block_boundary_ = false;

‎cpp/array_record_writer_test.cc‎

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,8 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
209209
auto option = ArrayRecordWriterBase::Options::FromString("").value();
210210
EXPECT_EQ(option.group_size(),
211211
ArrayRecordWriterBase::Options::kDefaultGroupSize);
212+
EXPECT_EQ(option.groups_awaiting_flush(),
213+
ArrayRecordWriterBase::Options::kDefaultGroupsAwaitingFlush);
212214
EXPECT_FALSE(option.transpose());
213215
EXPECT_EQ(option.max_parallelism(), std::nullopt);
214216
EXPECT_EQ(option.compressor_options().compression_type(),
@@ -218,7 +220,8 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
218220
EXPECT_FALSE(option.pad_to_block_boundary());
219221

220222
EXPECT_EQ(option.ToString(),
221-
"group_size:65536,"
223+
"group_size:1,"
224+
"groups_awaiting_flush:1024,"
222225
"transpose:false,"
223226
"pad_to_block_boundary:false,"
224227
"zstd:3,"
@@ -230,6 +233,8 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
230233
auto option = ArrayRecordWriterBase::Options::FromString("default").value();
231234
EXPECT_EQ(option.group_size(),
232235
ArrayRecordWriterBase::Options::kDefaultGroupSize);
236+
EXPECT_EQ(option.groups_awaiting_flush(),
237+
ArrayRecordWriterBase::Options::kDefaultGroupsAwaitingFlush);
233238
EXPECT_FALSE(option.transpose());
234239
EXPECT_EQ(option.max_parallelism(), std::nullopt);
235240
EXPECT_EQ(option.compressor_options().compression_type(),
@@ -239,7 +244,8 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
239244
EXPECT_FALSE(option.pad_to_block_boundary());
240245

241246
EXPECT_EQ(option.ToString(),
242-
"group_size:65536,"
247+
"group_size:1,"
248+
"groups_awaiting_flush:1024,"
243249
"transpose:false,"
244250
"pad_to_block_boundary:false,"
245251
"zstd:3,"
@@ -248,10 +254,12 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
248254
ArrayRecordWriterBase::Options::FromString(option.ToString()).ok());
249255
}
250256
{
251-
auto option = ArrayRecordWriterBase::Options::FromString(
252-
"group_size:32,transpose,window_log:20")
253-
.value();
257+
auto option =
258+
ArrayRecordWriterBase::Options::FromString(
259+
"group_size:32,groups_awaiting_flush:256,transpose,window_log:20")
260+
.value();
254261
EXPECT_EQ(option.group_size(), 32);
262+
EXPECT_EQ(option.groups_awaiting_flush(), 256);
255263
EXPECT_TRUE(option.transpose());
256264
EXPECT_EQ(option.max_parallelism(), std::nullopt);
257265
EXPECT_EQ(option.compressor_options().compression_type(),
@@ -261,6 +269,7 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
261269

262270
EXPECT_EQ(option.ToString(),
263271
"group_size:32,"
272+
"groups_awaiting_flush:256,"
264273
"transpose:true,"
265274
"pad_to_block_boundary:false,"
266275
"transpose_bucket_size:256,"
@@ -274,6 +283,8 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
274283
"brotli:6,group_size:32,transpose,window_log:25")
275284
.value();
276285
EXPECT_EQ(option.group_size(), 32);
286+
EXPECT_EQ(option.groups_awaiting_flush(),
287+
ArrayRecordWriterBase::Options::kDefaultGroupsAwaitingFlush);
277288
EXPECT_TRUE(option.transpose());
278289
EXPECT_EQ(option.max_parallelism(), std::nullopt);
279290
EXPECT_EQ(option.compressor_options().compression_type(),
@@ -283,6 +294,7 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
283294

284295
EXPECT_EQ(option.ToString(),
285296
"group_size:32,"
297+
"groups_awaiting_flush:1024,"
286298
"transpose:true,"
287299
"pad_to_block_boundary:false,"
288300
"transpose_bucket_size:256,"
@@ -296,6 +308,8 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
296308
"group_size:32,transpose,zstd:5")
297309
.value();
298310
EXPECT_EQ(option.group_size(), 32);
311+
EXPECT_EQ(option.groups_awaiting_flush(),
312+
ArrayRecordWriterBase::Options::kDefaultGroupsAwaitingFlush);
299313
EXPECT_TRUE(option.transpose());
300314
EXPECT_EQ(option.max_parallelism(), std::nullopt);
301315
EXPECT_EQ(option.compressor_options().compression_type(),
@@ -306,6 +320,7 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
306320

307321
EXPECT_EQ(option.ToString(),
308322
"group_size:32,"
323+
"groups_awaiting_flush:1024,"
309324
"transpose:true,"
310325
"pad_to_block_boundary:false,"
311326
"transpose_bucket_size:256,"
@@ -320,14 +335,17 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
320335
.value();
321336
EXPECT_EQ(option.group_size(),
322337
ArrayRecordWriterBase::Options::kDefaultGroupSize);
338+
EXPECT_EQ(option.groups_awaiting_flush(),
339+
ArrayRecordWriterBase::Options::kDefaultGroupsAwaitingFlush);
323340
EXPECT_FALSE(option.transpose());
324341
EXPECT_EQ(option.max_parallelism(), std::nullopt);
325342
EXPECT_EQ(option.compressor_options().compression_type(),
326343
riegeli::CompressionType::kNone);
327344
EXPECT_TRUE(option.pad_to_block_boundary());
328345

329346
EXPECT_EQ(option.ToString(),
330-
"group_size:65536,"
347+
"group_size:1,"
348+
"groups_awaiting_flush:1024,"
331349
"transpose:false,"
332350
"pad_to_block_boundary:true,"
333351
"uncompressed");
@@ -340,14 +358,17 @@ TEST(ArrayRecordWriterOptionsTest, ParsingTest) {
340358
.value();
341359
EXPECT_EQ(option.group_size(),
342360
ArrayRecordWriterBase::Options::kDefaultGroupSize);
361+
EXPECT_EQ(option.groups_awaiting_flush(),
362+
ArrayRecordWriterBase::Options::kDefaultGroupsAwaitingFlush);
343363
EXPECT_FALSE(option.transpose());
344364
EXPECT_EQ(option.max_parallelism(), std::nullopt);
345365
EXPECT_EQ(option.compressor_options().compression_type(),
346366
riegeli::CompressionType::kSnappy);
347367
EXPECT_TRUE(option.pad_to_block_boundary());
348368

349369
EXPECT_EQ(option.ToString(),
350-
"group_size:65536,"
370+
"group_size:1,"
371+
"groups_awaiting_flush:1024,"
351372
"transpose:false,"
352373
"pad_to_block_boundary:true,"
353374
"snappy");

‎cpp/sequenced_chunk_writer.cc‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,8 @@ void SequencedChunkWriterBase::Initialize() {
149149
"Failed to create the file header"));
150150
}
151151
if (!chunk_writer->Flush(riegeli::FlushType::kFromObject)) {
152-
Fail(riegeli::Annotate(chunk_writer->status(), "Could not flush"));
152+
Fail(riegeli::Annotate(chunk_writer->status(),
153+
"Could not flush the file header."));
153154
}
154155
}
155156

@@ -159,6 +160,11 @@ void SequencedChunkWriterBase::Done() {
159160
return;
160161
}
161162
auto* chunk_writer = get_writer();
163+
// if (!chunk_writer->Flush(riegeli::FlushType::kFromObject)) {
164+
// Fail(riegeli::Annotate(chunk_writer->status(),
165+
// "Could not flush before close."));
166+
// return;
167+
// }
162168
if (!chunk_writer->Close()) {
163169
Fail(riegeli::Annotate(chunk_writer->status(),
164170
"Failed to close chunk_writer"));

‎cpp/sequenced_chunk_writer.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,6 +145,11 @@ class SequencedChunkWriterBase : public riegeli::Object {
145145
return pad_to_block_boundary_;
146146
}
147147

148+
void set_chunks_awaiting_flush(uint32_t chunks_awaiting_flush) {
149+
absl::MutexLock l(&mu_);
150+
chunks_awaiting_flush_ = chunks_awaiting_flush;
151+
}
152+
148153
// Setup a callback for each committed chunk. See CommitChunkCallback
149154
// comments for details.
150155
void set_submit_chunk_callback(SubmitChunkCallback* callback) {
@@ -182,6 +187,8 @@ class SequencedChunkWriterBase : public riegeli::Object {
182187
// Records the sequence number of submitted chunks.
183188
uint64_t submitted_chunks_ ABSL_GUARDED_BY(mu_) = 0;
184189

190+
uint64_t chunks_awaiting_flush_ ABSL_GUARDED_BY(mu_) = 0;
191+
185192
// Queue for storing the future chunks.
186193
std::queue<std::future<absl::StatusOr<riegeli::Chunk>>> queue_
187194
ABSL_GUARDED_BY(mu_);

0 commit comments

Comments
 (0)