Skip to content

Exact vector search #2519

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 19 commits into from
Mar 14, 2024
Merged

Exact vector search #2519

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
55ec947
First Knn version
azevaykin Mar 5, 2024
2e20f2c
Serialization/deserialization
azevaykin Mar 6, 2024
b636bcc
InnerProductDistance
azevaykin Mar 6, 2024
40afce6
Trivial PR fixes
azevaykin Mar 7, 2024
897ee20
EnumerateVectors
azevaykin Mar 7, 2024
d33cc3e
Remove memcpy
azevaykin Mar 7, 2024
8a24959
double -> float
azevaykin Mar 7, 2024
6a37abd
InnerProductSimilarity
azevaykin Mar 7, 2024
6be9f7b
TAutoMap
azevaykin Mar 7, 2024
385cf97
CosineSimilarity
azevaykin Mar 7, 2024
320200b
Typo
azevaykin Mar 7, 2024
1317e6e
CosineDistance
azevaykin Mar 11, 2024
5167376
EuclideanDistance
azevaykin Mar 11, 2024
7e85069
template <typename TCallback>
azevaykin Mar 12, 2024
95de228
Remove excessive passes
azevaykin Mar 12, 2024
02344d4
Review fixes
azevaykin Mar 12, 2024
9d46bed
ISerializer
azevaykin Mar 13, 2024
1006413
refactor TFloatVectorSerializer
azevaykin Mar 13, 2024
f21ee11
Remove virtual classes
azevaykin Mar 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 130 additions & 0 deletions ydb/library/yql/udfs/common/knn/knn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
#include <ydb/library/yql/public/udf/udf_helpers.h>

#include <util/generic/buffer.h>
#include <util/stream/format.h>

using namespace NYql;
using namespace NYql::NUdf;

enum EFormat : unsigned char {
DoubleVector = 1
};

TString SerializeDoubleVector(const TUnboxedValuePod x) {
const EFormat format = EFormat::DoubleVector;

TString str;
TStringOutput outStr(str);
if (const auto elements = x.GetElements()) {
const auto size = x.GetListLength();
outStr.Reserve(1 + size * sizeof(double));
outStr.Write(&format, sizeof(unsigned char));
for (ui32 i = 0; i < size; ++i) {
double element = elements[i].Get<double>();
outStr.Write(&element, sizeof(double));
}
} else {
outStr.Write(&format, sizeof(unsigned char));
const auto it = x.GetListIterator();
TUnboxedValue v;
while(it.Next(v)) {
double element = v.Get<double>();
outStr.Write(&element, sizeof(double));
}
}
return str;
}

NYql::NUdf::TUnboxedValue DeserializeDoubleVector(const IValueBuilder *valueBuilder, TStringRef str) {
if (str.Size() % sizeof(double) != 0)
return TUnboxedValue(TUnboxedValuePod());

const ui32 count = str.Size() / sizeof(double);

TUnboxedValue* items = nullptr;
auto res = valueBuilder->NewArray(count, items);

for (ui32 i = 0; i < count; ++i) {
double element;
memcpy(&element, str.Data() + i * sizeof(double), sizeof(double));
*items++ = TUnboxedValuePod{element};
}

return res.Release();
}

SIMPLE_STRICT_UDF(TToBinaryString, char*(TListType<double>)) {
return valueBuilder->NewString(SerializeDoubleVector(args[0]));
}

SIMPLE_STRICT_UDF(TFromBinaryString, TOptional<TListType<double>>(const char*)) {
TStringRef str = args[0].AsStringRef();
if (str.Size() == 0)
return TUnboxedValue();

const EFormat format = (EFormat)str.Data()[0];
str = TStringRef{str.Data() + 1, str.Size() -1};
switch (format) {
case EFormat::DoubleVector: {
return DeserializeDoubleVector(valueBuilder, str);
break;
}
default:
return TUnboxedValue();
}
}


std::optional<double> InnerProductDistance(const TUnboxedValuePod vector1, const TUnboxedValuePod vector2) {
double ret = 0;

if (vector1.HasFastListLength() && vector2.HasFastListLength()) {
const auto elements1 = vector1.GetElements();
const auto elements2 = vector2.GetElements();
const auto size1 = vector1.GetListLength();
const auto size2 = vector2.GetListLength();

if (size1 != size2)
return {};

for (ui32 i = 0; i < size1; ++i) {
double element1 = elements1[i].Get<double>();
double element2 = elements2[i].Get<double>();
ret += element1 * element2;
}
} else {
TUnboxedValue value1, value2;
const auto it1 = vector1.GetListIterator();
const auto it2 = vector2.GetListIterator();
for (; it1.Next(value1) && it2.Next(value2);) {
ret += value1.Get<double>() * value2.Get<double>();
}

// Lenght mismatch
if (it1.Next(value1) || it2.Next(value2))
return {};
}

return ret;
}

SIMPLE_STRICT_UDF(TInnerProductDistance, TOptional<double>(TOptional<TListType<double>>, TOptional<TListType<double>>)) {
Y_UNUSED(valueBuilder);

if (!args[0].HasValue() || !args[1].HasValue())
return TUnboxedValue(TUnboxedValuePod());

if (auto distance = InnerProductDistance(args[0], args[1]))
return TUnboxedValuePod{distance.value()};
else
return TUnboxedValuePod();
}

SIMPLE_MODULE(TKnnModule,
TFromBinaryString,
TToBinaryString,
TInnerProductDistance
)

REGISTER_MODULES(TKnnModule)

22 changes: 22 additions & 0 deletions ydb/library/yql/udfs/common/knn/test/canondata/result.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
{
"test.test[DeserializationError]": [
{
"uri": "file://test.test_DeserializationError_/results.txt"
}
],
"test.test[InnerProductDistance]": [
{
"uri": "file://test.test_InnerProductDistance_/results.txt"
}
],
"test.test[LazyListSerialization]": [
{
"uri": "file://test.test_LazyListSerialization_/results.txt"
}
],
"test.test[ListSerialization]": [
{
"uri": "file://test.test_ListSerialization_/results.txt"
}
]
}
66 changes: 66 additions & 0 deletions ydb/library/yql/udfs/common/knn/test/canondata/test.test_DeserializationError_/results.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
[
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"ListType";
[
"DataType";
"Double"
]
]
]
]
]
]
];
"Data" = [
[
#
]
]
}
]
};
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"ListType";
[
"DataType";
"Double"
]
]
]
]
]
]
];
"Data" = [
[
#
]
]
}
]
}
]
153 changes: 153 additions & 0 deletions ydb/library/yql/udfs/common/knn/test/canondata/test.test_InnerProductDistance_/results.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
[
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"DataType";
"Bool"
]
]
]
]
]
];
"Data" = [
[
[
%true
]
]
]
}
]
};
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"DataType";
"Bool"
]
]
]
]
]
];
"Data" = [
[
[
%true
]
]
]
}
]
};
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"DataType";
"Bool"
]
]
]
]
]
];
"Data" = [
[
[
%true
]
]
]
}
]
};
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"DataType";
"Double"
]
]
]
]
]
];
"Data" = [
[
#
]
]
}
]
};
{
"Write" = [
{
"Type" = [
"ListType";
[
"StructType";
[
[
"column0";
[
"OptionalType";
[
"DataType";
"Double"
]
]
]
]
]
];
"Data" = [
[
#
]
]
}
]
}
]
Loading