Enable Cord reading and writing for InputStream and OutputStreams
PiperOrigin-RevId: 343014311 Change-Id: If1b7fee3cb1432e35952104984e64e54a1aa2f72
This commit is contained in:
parent
8d06472606
commit
516ae286f6
tensorflow/core
@ -263,7 +263,7 @@ uint32 Extend(uint32 crc, const char *buf, size_t size) {
|
||||
return l ^ 0xffffffffu;
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
uint32 Extend(uint32 crc, const absl::Cord &cord) {
|
||||
for (absl::string_view fragment : cord.Chunks()) {
|
||||
crc = Extend(crc, fragment.data(), fragment.size());
|
||||
|
@ -32,11 +32,14 @@ namespace crc32c {
|
||||
// crc32c of a stream of data.
|
||||
extern uint32 Extend(uint32 init_crc, const char* data, size_t n);
|
||||
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
extern uint32 Extend(uint32 init_crc, const absl::Cord& cord);
|
||||
#endif
|
||||
|
||||
// Return the crc32c of data[0,n-1]
|
||||
inline uint32 Value(const char* data, size_t n) { return Extend(0, data, n); }
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
extern uint32 Extend(uint32 init_crc, const absl::Cord& cord);
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
inline uint32 Value(const absl::Cord& cord) { return Extend(0, cord); }
|
||||
#endif
|
||||
|
||||
|
@ -37,7 +37,7 @@ class InputStreamInterface {
|
||||
// * OUT_OF_RANGE - not enough bytes remaining before end of file.
|
||||
virtual Status ReadNBytes(int64 bytes_to_read, tstring* result) = 0;
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
// Reads the next bytes_to_read from the file. Typical return codes:
|
||||
// * OK - in case of success.
|
||||
// * OUT_OF_RANGE - not enough bytes remaining before end of file.
|
||||
|
@ -49,7 +49,7 @@ Status RandomAccessInputStream::ReadNBytes(int64 bytes_to_read,
|
||||
return s;
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status RandomAccessInputStream::ReadNBytes(int64 bytes_to_read,
|
||||
absl::Cord* result) {
|
||||
if (bytes_to_read < 0) {
|
||||
|
@ -35,7 +35,7 @@ class RandomAccessInputStream : public InputStreamInterface {
|
||||
|
||||
Status ReadNBytes(int64 bytes_to_read, tstring* result) override;
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status ReadNBytes(int64 bytes_to_read, absl::Cord* result) override;
|
||||
#endif
|
||||
|
||||
|
@ -115,7 +115,7 @@ Status RecordWriter::WriteRecord(StringPiece data) {
|
||||
return dest_->Append(StringPiece(footer, sizeof(footer)));
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status RecordWriter::WriteRecord(const absl::Cord& data) {
|
||||
if (dest_ == nullptr) {
|
||||
return Status(::tensorflow::error::FAILED_PRECONDITION,
|
||||
|
@ -68,8 +68,8 @@ class RecordWriter {
|
||||
// Create a writer that will append data to "*dest".
|
||||
// "*dest" must be initially empty.
|
||||
// "*dest" must remain live while this Writer is in use.
|
||||
RecordWriter(WritableFile* dest,
|
||||
const RecordWriterOptions& options = RecordWriterOptions());
|
||||
explicit RecordWriter(WritableFile* dest, const RecordWriterOptions& options =
|
||||
RecordWriterOptions());
|
||||
|
||||
// Calls Close() and logs if an error occurs.
|
||||
//
|
||||
@ -79,7 +79,7 @@ class RecordWriter {
|
||||
|
||||
Status WriteRecord(StringPiece data);
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status WriteRecord(const absl::Cord& data);
|
||||
#endif
|
||||
|
||||
@ -98,12 +98,13 @@ class RecordWriter {
|
||||
// "header[0,kHeaderSize-1]". The record-header is based on data[0, n-1].
|
||||
inline static void PopulateHeader(char* header, const char* data, size_t n);
|
||||
|
||||
inline static void PopulateHeader(char* header, const absl::Cord& data);
|
||||
|
||||
// Utility method to populate TFRecord footers. Populates record-footer in
|
||||
// "footer[0,kFooterSize-1]". The record-footer is based on data[0, n-1].
|
||||
inline static void PopulateFooter(char* footer, const char* data, size_t n);
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
inline static void PopulateHeader(char* header, const absl::Cord& data);
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
inline static void PopulateFooter(char* footer, const absl::Cord& data);
|
||||
#endif
|
||||
|
||||
@ -115,7 +116,7 @@ class RecordWriter {
|
||||
return crc32c::Mask(crc32c::Value(data, n));
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
inline static uint32 MaskedCrc(const absl::Cord& data) {
|
||||
return crc32c::Mask(crc32c::Value(data));
|
||||
}
|
||||
@ -134,7 +135,7 @@ void RecordWriter::PopulateFooter(char* footer, const char* data, size_t n) {
|
||||
core::EncodeFixed32(footer, MaskedCrc(data, n));
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
void RecordWriter::PopulateHeader(char* header, const absl::Cord& data) {
|
||||
core::EncodeFixed64(header + 0, data.size());
|
||||
core::EncodeFixed32(header + sizeof(uint64),
|
||||
|
@ -62,7 +62,7 @@ class StringDest : public WritableFile {
|
||||
contents_->append(slice.data(), slice.size());
|
||||
return Status::OK();
|
||||
}
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status Append(const absl::Cord& data) override {
|
||||
contents_->append(std::string(data));
|
||||
return Status::OK();
|
||||
@ -136,7 +136,7 @@ class RecordioTest : public ::testing::Test {
|
||||
TF_ASSERT_OK(writer_->WriteRecord(StringPiece(msg)));
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
void Write(const absl::Cord& msg) {
|
||||
ASSERT_TRUE(!reading_) << "Write() after starting to read";
|
||||
TF_ASSERT_OK(writer_->WriteRecord(msg));
|
||||
@ -204,7 +204,7 @@ TEST_F(RecordioTest, ReadWrite) {
|
||||
ASSERT_EQ("EOF", Read()); // Make sure reads at eof work
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
TEST_F(RecordioTest, ReadWriteCords) {
|
||||
Write(absl::Cord("foo"));
|
||||
Write(absl::Cord("bar"));
|
||||
|
@ -68,7 +68,7 @@ Status SnappyInputStream::ReadNBytes(int64 bytes_to_read, tstring* result) {
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status SnappyInputStream::ReadNBytes(int64 bytes_to_read, absl::Cord* result) {
|
||||
// TODO(frankchn): Optimize this instead of bouncing through the buffer.
|
||||
tstring buf;
|
||||
|
@ -46,7 +46,7 @@ class SnappyInputStream : public InputStreamInterface {
|
||||
// others: If reading from stream failed.
|
||||
Status ReadNBytes(int64 bytes_to_read, tstring* result) override;
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status ReadNBytes(int64 bytes_to_read, absl::Cord* result) override;
|
||||
#endif
|
||||
|
||||
|
@ -40,7 +40,7 @@ SnappyOutputBuffer::~SnappyOutputBuffer() {
|
||||
|
||||
Status SnappyOutputBuffer::Append(StringPiece data) { return Write(data); }
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status SnappyOutputBuffer::Append(const absl::Cord& cord) {
|
||||
for (absl::string_view fragment : cord.Chunks()) {
|
||||
TF_RETURN_IF_ERROR(Append(fragment));
|
||||
|
@ -65,7 +65,7 @@ class SnappyOutputBuffer : public WritableFile {
|
||||
// later time. To immediately write contents to file call `Flush()`.
|
||||
Status Append(StringPiece data) override;
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status Append(const absl::Cord& cord) override;
|
||||
#endif
|
||||
|
||||
|
@ -190,7 +190,7 @@ Status ZlibOutputBuffer::Append(StringPiece data) {
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status ZlibOutputBuffer::Append(const absl::Cord& cord) {
|
||||
for (absl::string_view fragment : cord.Chunks()) {
|
||||
TF_RETURN_IF_ERROR(Append(fragment));
|
||||
|
@ -65,7 +65,7 @@ class ZlibOutputBuffer : public WritableFile {
|
||||
// To immediately write contents to file call `Flush()`.
|
||||
Status Append(StringPiece data) override;
|
||||
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
Status Append(const absl::Cord& cord) override;
|
||||
#endif
|
||||
|
||||
|
@ -751,8 +751,7 @@ class RandomAccessFile {
|
||||
virtual tensorflow::Status Read(uint64 offset, size_t n, StringPiece* result,
|
||||
char* scratch) const = 0;
|
||||
|
||||
// TODO(ebrevdo): Remove this ifdef when absl is updated.
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
/// \brief Read up to `n` bytes from the file starting at `offset`.
|
||||
virtual tensorflow::Status Read(uint64 offset, size_t n,
|
||||
absl::Cord* cord) const {
|
||||
@ -778,8 +777,7 @@ class WritableFile {
|
||||
/// \brief Append 'data' to the file.
|
||||
virtual tensorflow::Status Append(StringPiece data) = 0;
|
||||
|
||||
// TODO(ebrevdo): Remove this ifdef when absl is updated.
|
||||
#if defined(PLATFORM_GOOGLE)
|
||||
#if defined(TF_CORD_SUPPORT)
|
||||
// \brief Append 'data' to the file.
|
||||
virtual tensorflow::Status Append(const absl::Cord& cord) {
|
||||
return errors::Unimplemented("Append(absl::Cord) is not implemented");
|
||||
|
Loading…
Reference in New Issue
Block a user