| 968 | } |
| 969 | |
| 970 | Status BundleReader::GetValue(const BundleEntryProto& entry, Tensor* val) { |
| 971 | Tensor* ret = val; |
| 972 | const TensorShape stored_shape(TensorShape(entry.shape())); |
| 973 | if (val->NumElements() == 0) { |
| 974 | ret = new Tensor(entry.dtype(), stored_shape); |
| 975 | } |
| 976 | |
| 977 | // Validates the "size" field. |
| 978 | if (entry.dtype() != DT_STRING && entry.dtype() != DT_VARIANT) { |
| 979 | if (entry.size() != ret->TotalBytes()) { |
| 980 | return errors::DataLoss("Invalid size in bundle entry: key ", key(), |
| 981 | "; stored size ", entry.size(), |
| 982 | "; expected size ", ret->TotalBytes()); |
| 983 | } |
| 984 | } else if (entry.dtype() == DT_STRING) { |
| 985 | // Relaxes the check for string tensors as follows: |
| 986 | // entry.size() == bytes(varint lengths) + bytes(data) |
| 987 | // >= NumElems + bytes(data), since size bytes(varint) >= 1. |
| 988 | // TotalBytes() == sizeof(tstring) * NumElems + bytes(data) |
| 989 | // Since we don't know bytes(varint lengths), we just check an inequality. |
| 990 | const size_t lower_bound = ret->NumElements() + ret->TotalBytes() - |
| 991 | sizeof(tstring) * ret->NumElements(); |
| 992 | if (entry.size() < lower_bound) { |
| 993 | return errors::DataLoss("Invalid size in bundle entry: key ", key(), |
| 994 | "; stored size ", entry.size(), |
| 995 | "; expected size is at least ", lower_bound); |
| 996 | } |
| 997 | } |
| 998 | |
| 999 | // Open the data file if it has not been opened. |
| 1000 | io::InputBuffer* buffered_file = data_[entry.shard_id()]; |
| 1001 | if (buffered_file == nullptr) { |
| 1002 | std::unique_ptr<RandomAccessFile> file = nullptr; |
| 1003 | TF_RETURN_IF_ERROR(env_->NewRandomAccessFile( |
| 1004 | DataFilename(prefix_, entry.shard_id(), num_shards_), &file)); |
| 1005 | buffered_file = new io::InputBuffer(file.release(), kBufferSize); |
| 1006 | // The InputBuffer and RandomAccessFile objects are both released in dtor. |
| 1007 | data_[entry.shard_id()] = buffered_file; |
| 1008 | } |
| 1009 | CHECK(buffered_file != nullptr); |
| 1010 | |
| 1011 | TF_RETURN_IF_ERROR(buffered_file->Seek(entry.offset())); |
| 1012 | uint32 actual_crc32c = 0; |
| 1013 | |
| 1014 | if (DataTypeCanUseMemcpy(entry.dtype())) { |
| 1015 | char* backing_buffer = const_cast<char*>((ret->tensor_data().data())); |
| 1016 | size_t unused_bytes_read; |
| 1017 | if (entry.size() > kBufferSize) { |
| 1018 | StringPiece sp; |
| 1019 | TF_RETURN_IF_ERROR(buffered_file->file()->Read( |
| 1020 | entry.offset(), entry.size(), &sp, backing_buffer)); |
| 1021 | if (sp.data() != backing_buffer) { |
| 1022 | memmove(backing_buffer, sp.data(), entry.size()); |
| 1023 | } |
| 1024 | } else { |
| 1025 | TF_RETURN_IF_ERROR(buffered_file->ReadNBytes(entry.size(), backing_buffer, |
| 1026 | &unused_bytes_read)); |
| 1027 | } |
nothing calls this directly
no test coverage detected