| 109 | } |
| 110 | |
| 111 | Status RecordReader::GetMetadata(Metadata* md) { |
| 112 | if (!md) { |
| 113 | return errors::InvalidArgument( |
| 114 | "Metadata object call to GetMetadata() was null"); |
| 115 | } |
| 116 | |
| 117 | // Compute the metadata of the TFRecord file if not cached. |
| 118 | if (!cached_metadata_) { |
| 119 | TF_RETURN_IF_ERROR(input_stream_->Reset()); |
| 120 | |
| 121 | int64 data_size = 0; |
| 122 | int64 entries = 0; |
| 123 | |
| 124 | // Within the loop, we always increment offset positively, so this |
| 125 | // loop should be guaranteed to either return after reaching EOF |
| 126 | // or encountering an error. |
| 127 | uint64 offset = 0; |
| 128 | string record; |
| 129 | while (true) { |
| 130 | // Read header, containing size of data. |
| 131 | Status s = ReadChecksummed(offset, sizeof(uint64), &record); |
| 132 | if (!s.ok()) { |
| 133 | if (errors::IsOutOfRange(s)) { |
| 134 | // We should reach out of range when the record file is complete. |
| 135 | break; |
| 136 | } |
| 137 | return s; |
| 138 | } |
| 139 | |
| 140 | // Read the length of the data. |
| 141 | const uint64 length = core::DecodeFixed64(record.data()); |
| 142 | |
| 143 | // Skip reading the actual data since we just want the number |
| 144 | // of records and the size of the data. |
| 145 | TF_RETURN_IF_ERROR(input_stream_->SkipNBytes(length + kFooterSize)); |
| 146 | offset += kHeaderSize + length + kFooterSize; |
| 147 | |
| 148 | // Increment running stats. |
| 149 | data_size += length; |
| 150 | ++entries; |
| 151 | } |
| 152 | |
| 153 | cached_metadata_.reset(new Metadata()); |
| 154 | cached_metadata_->stats.entries = entries; |
| 155 | cached_metadata_->stats.data_size = data_size; |
| 156 | cached_metadata_->stats.file_size = |
| 157 | data_size + (kHeaderSize + kFooterSize) * entries; |
| 158 | } |
| 159 | |
| 160 | md->stats = cached_metadata_->stats; |
| 161 | return Status::OK(); |
| 162 | } |
| 163 | |
| 164 | Status RecordReader::ReadRecord(uint64* offset, string* record) { |
| 165 | // Position the input stream. |