blob: 0de4491cfcc401de8c1f0e8679ba9e96c969f6f7 (
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
49
50
51
52
53
54
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#include "direct_tensor_saver.h"
#include "direct_tensor_store.h"
#include <vespa/eval/eval/value_codec.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 {
DirectTensorAttributeSaver::
DirectTensorAttributeSaver(GenerationHandler::Guard &&guard,
const attribute::AttributeHeader &header,
RefCopyVector &&refs,
const DirectTensorStore &tensorStore)
: AttributeSaver(std::move(guard), header),
_refs(std::move(refs)),
_tensorStore(tensorStore)
{
}
DirectTensorAttributeSaver::~DirectTensorAttributeSaver()
{
}
bool
DirectTensorAttributeSaver::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) {
const vespalib::eval::Value *tensor = _tensorStore.get_tensor_ptr(_refs[lid]);
if (tensor) {
stream.clear();
encode_value(*tensor, stream);
uint32_t sz = stream.size();
datWriter->write(&sz, sizeof(sz));
datWriter->write(stream.peek(), stream.size());
} else {
uint32_t sz = 0;
datWriter->write(&sz, sizeof(sz));
}
}
datWriter->flush();
return true;
}
} // namespace search::tensor
|
