blob: 0963e79b0dd0dc4250da3343295b0d04bfa6c218 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#include "tensor_store_saver.h"
#include "tensor_store.h"
#include <vespa/searchlib/attribute/iattributesavetarget.h>
#include <vespa/searchlib/util/bufferwriter.h>
#include <vespa/vespalib/objects/nbostream.h>
using vespalib::GenerationHandler;
namespace search::tensor {
TensorStoreSaver::
TensorStoreSaver(GenerationHandler::Guard &&guard,
const attribute::AttributeHeader &header,
RefCopyVector &&refs,
const TensorStore &tensorStore)
: AttributeSaver(std::move(guard), header),
_refs(std::move(refs)),
_tensorStore(tensorStore)
{
}
TensorStoreSaver::~TensorStoreSaver() = default;
bool
TensorStoreSaver::onSave(IAttributeSaveTarget &saveTarget)
{
auto datWriter = saveTarget.datWriter().allocBufferWriter();
const uint32_t docIdLimit(_refs.size());
vespalib::nbostream stream;
for (uint32_t lid = 0; lid < docIdLimit; ++lid) {
if (_tensorStore.encode_stored_tensor(_refs[lid], stream)) {
uint32_t sz = stream.size();
datWriter->write(&sz, sizeof(sz));
datWriter->write(stream.peek(), stream.size());
stream.clear();
} else {
uint32_t sz = 0;
datWriter->write(&sz, sizeof(sz));
}
}
datWriter->flush();
return true;
}
} // namespace search::tensor
|
