Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
3 changes: 3 additions & 0 deletions cpp/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -115,6 +115,7 @@ cc_library(
"@abseil-cpp//absl/status",
"@abseil-cpp//absl/status:statusor",
"@abseil-cpp//absl/strings",
"@abseil-cpp//absl/strings:cord",
"@abseil-cpp//absl/synchronization",
"@abseil-cpp//absl/types:span",
"@protobuf//:protobuf_lite",
Expand Down Expand Up @@ -235,6 +236,8 @@ cc_test(
":test_utils",
":thread_pool",
"@abseil-cpp//absl/strings",
"@abseil-cpp//absl/strings:cord",
"@abseil-cpp//absl/strings:cord_test_helpers",
"@googletest//:gtest_main",
"@riegeli//riegeli/base:initializer",
"@riegeli//riegeli/bytes:string_reader",
Expand Down
13 changes: 12 additions & 1 deletion cpp/array_record_writer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ limitations under the License.
#include "absl/log/check.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
Expand All @@ -39,8 +40,8 @@ limitations under the License.
#include "cpp/common.h"
#include "cpp/layout.pb.h"
#include "cpp/sequenced_chunk_writer.h"
#include "cpp/tri_state_ptr.h"
#include "cpp/thread_pool.h"
#include "cpp/tri_state_ptr.h"
#include "google/protobuf/message_lite.h"
#include "riegeli/base/object.h"
#include "riegeli/base/options_parser.h"
Expand Down Expand Up @@ -418,6 +419,16 @@ bool ArrayRecordWriterBase::WriteRecord(absl::string_view record) {
return WriteRecordImpl(std::move(record));
}

bool ArrayRecordWriterBase::WriteRecord(const absl::Cord& record) {
if (auto flat = record.TryFlat(); flat.has_value()) {
return WriteRecord(*flat);
}

std::string cord_string;
absl::AppendCordToString(record, &cord_string);
return WriteRecord(cord_string);
}

bool ArrayRecordWriterBase::WriteRecord(const void* data, size_t num_bytes) {
auto view = absl::string_view(reinterpret_cast<const char*>(data), num_bytes);
return WriteRecordImpl(std::move(view));
Expand Down
4 changes: 3 additions & 1 deletion cpp/array_record_writer.h
Original file line number Diff line number Diff line change
Expand Up @@ -66,12 +66,13 @@ limitations under the License.
#include <utility>

#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "cpp/common.h"
#include "cpp/sequenced_chunk_writer.h"
#include "cpp/tri_state_ptr.h"
#include "cpp/thread_pool.h"
#include "cpp/tri_state_ptr.h"
#include "riegeli/base/initializer.h"
#include "riegeli/base/object.h"
#include "riegeli/bytes/writer.h"
Expand Down Expand Up @@ -304,6 +305,7 @@ class ArrayRecordWriterBase : public riegeli::Object {
// Write records of various types.
bool WriteRecord(const google::protobuf::MessageLite& record);
bool WriteRecord(absl::string_view record);
bool WriteRecord(const absl::Cord& record);
bool WriteRecord(const void* data, size_t num_bytes);
template <typename T>
bool WriteRecord(absl::Span<const T> record) {
Expand Down
35 changes: 35 additions & 0 deletions cpp/array_record_writer_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ limitations under the License.
#include <vector>

#include "gtest/gtest.h"
#include "absl/strings/cord.h"
#include "absl/strings/cord_test_helpers.h"
#include "absl/strings/string_view.h"
#include "cpp/common.h"
#include "cpp/layout.pb.h"
Expand Down Expand Up @@ -116,6 +118,39 @@ TEST_P(ArrayRecordWriterTest, MoveTest) {
}
}

TEST_P(ArrayRecordWriterTest, CordTest) {
std::string encoded;
ARThreadPool* pool = nullptr;
if (std::get<3>(GetParam())) {
pool = ArrayRecordGlobalPool();
}
auto options = GetOptions();
options.set_group_size(2);
auto writer = ArrayRecordWriter(
riegeli::Maker<riegeli::StringWriter>(&encoded), options, pool);

absl::Cord flat_cord("test");
// Empty string should not crash the writer.
absl::Cord empty_cord("");
absl::Cord fragmented_cord = absl::MakeFragmentedCord({"aaa ", "", "c"});

EXPECT_TRUE(writer.WriteRecord(flat_cord));
EXPECT_TRUE(writer.WriteRecord(empty_cord));
EXPECT_TRUE(writer.WriteRecord(fragmented_cord));
ASSERT_TRUE(writer.Close());

// Empty string should not crash the reader.
std::vector<std::string> expected_strings{"test", "", "aaa c"};

auto reader =
riegeli::RecordReader(riegeli::Maker<riegeli::StringReader>(encoded));
for (const auto& expected : expected_strings) {
std::string result;
reader.ReadRecord(result);
EXPECT_EQ(result, expected);
}
}

TEST_P(ArrayRecordWriterTest, RandomDatasetTest) {
std::mt19937 bitgen;
constexpr uint32_t kGroupSize = 100;
Expand Down
Loading