Merge branch 'main' into acl-facade

contrib/charts/dragonfly/ci/priorityclassname-values.golden.yaml

@@ -0,0 +1,103 @@
+---
+# Source: dragonfly/templates/serviceaccount.yaml
+apiVersion: v1
+kind: ServiceAccount
+metadata:
+  name: test-dragonfly
+  namespace: default
+  labels:
+    app.kubernetes.io/name: dragonfly
+    app.kubernetes.io/instance: test
+    app.kubernetes.io/version: "v1.12.1"
+    app.kubernetes.io/managed-by: Helm
+---
+# Source: dragonfly/templates/service.yaml
+apiVersion: v1
+kind: Service
+metadata:
+  name: test-dragonfly
+  namespace: default
+  labels:
+    app.kubernetes.io/name: dragonfly
+    app.kubernetes.io/instance: test
+    app.kubernetes.io/version: "v1.12.1"
+    app.kubernetes.io/managed-by: Helm
+spec:
+  type: ClusterIP
+  ports:
+    - port: 6379
+      targetPort: dragonfly
+      protocol: TCP
+      name: dragonfly
+  selector:
+    app.kubernetes.io/name: dragonfly
+    app.kubernetes.io/instance: test
+---
+# Source: dragonfly/templates/deployment.yaml
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: test-dragonfly
+  namespace: default
+  labels:
+    app.kubernetes.io/name: dragonfly
+    app.kubernetes.io/instance: test
+    app.kubernetes.io/version: "v1.12.1"
+    app.kubernetes.io/managed-by: Helm
+spec:
+  replicas: 1
+  selector:
+    matchLabels:
+      app.kubernetes.io/name: dragonfly
+      app.kubernetes.io/instance: test
+  template:
+    metadata:
+      annotations:
+      labels:
+        app.kubernetes.io/name: dragonfly
+        app.kubernetes.io/instance: test
+    spec:
+      priorityClassName: high-priority
+      serviceAccountName: test-dragonfly
+      containers:
+        - name: dragonfly
+          image: "docker.dragonflydb.io/dragonflydb/dragonfly:v1.12.1"
+          imagePullPolicy: IfNotPresent
+          ports:
+            - name: dragonfly
+              containerPort: 6379
+              protocol: TCP
+          livenessProbe:
+            exec:
+              command:
+              - /bin/sh
+              - /usr/local/bin/healthcheck.sh
+            failureThreshold: 3
+            initialDelaySeconds: 10
+            periodSeconds: 10
+            successThreshold: 1
+            timeoutSeconds: 5
+          readinessProbe:
+            exec:
+              command:
+              - /bin/sh
+              - /usr/local/bin/healthcheck.sh
+            failureThreshold: 3
+            initialDelaySeconds: 10
+            periodSeconds: 10
+            successThreshold: 1
+            timeoutSeconds: 5
+          args:
+            - "--alsologtostderr"
+          resources:
+            limits: {}
+            requests: {}
+---
+# Source: dragonfly/templates/extra-manifests.yaml
+apiVersion: scheduling.k8s.io/v1
+description: This priority class should be used only for tests.
+globalDefault: false
+kind: PriorityClass
+metadata:
+  name: high-priority
+value: 1000000

contrib/charts/dragonfly/ci/priorityclassname-values.yaml

@@ -0,0 +1,10 @@
+priorityClassName: "high-priority"
+
+extraObjects:
+  - apiVersion: scheduling.k8s.io/v1
+    kind: PriorityClass
+    metadata:
+      name: high-priority
+    value: 1000000
+    globalDefault: false
+    description: "This priority class should be used only for tests."

contrib/charts/dragonfly/templates/_pod.tpl

@@ -16,6 +16,9 @@ volumeMounts:
 {{- end }}
 
 {{- define "dragonfly.pod" -}}
+{{- if ne .Values.priorityClassName "" }}
+priorityClassName: {{ .Values.priorityClassName }}
+{{- end }}
 {{- with .Values.tolerations }}
 tolerations:
   {{- toYaml . | trim | nindent 2 -}}

contrib/charts/dragonfly/values.yaml

@@ -193,6 +193,9 @@ resources:
   #   cpu: 100m
   #   memory: 128Mi
 
+# -- Priority class name for pod assignment
+priorityClassName: ""
+
 # -- Node labels for pod assignment
 nodeSelector: {}
 

src/core/search/search_test.cc

@@ -513,6 +513,54 @@ TEST_P(KnnTest, Simple2D) {
   }
 }
 
+TEST_P(KnnTest, Cosine) {
+  // Four arrows, closest cosing distance will be closes by angle
+  // 0 🡢 1 🡣 2 🡠 3 🡡
+  const pair<float, float> kTestCoords[] = {{1, 0}, {0, -1}, {-1, 0}, {0, 1}};
+
+  auto schema = MakeSimpleSchema({{"pos", SchemaField::VECTOR}});
+  schema.fields["pos"].special_params =
+      SchemaField::VectorParams{GetParam(), 2, VectorSimilarity::COSINE};
+  FieldIndices indices{schema, PMR_NS::get_default_resource()};
+
+  for (size_t i = 0; i < ABSL_ARRAYSIZE(kTestCoords); i++) {
+    string coords = ToBytes({kTestCoords[i].first, kTestCoords[i].second});
+    MockedDocument doc{Map{{"pos", coords}}};
+    indices.Add(i, &doc);
+  }
+
+  SearchAlgorithm algo{};
+  QueryParams params;
+
+  // Point down
+  {
+    params["vec"] = ToBytes({-0.1, -10});
+    algo.Init("* =>[KNN 1 @pos $vec]", &params);
+    EXPECT_THAT(algo.Search(&indices).ids, testing::UnorderedElementsAre(1));
+  }
+
+  // Point left
+  {
+    params["vec"] = ToBytes({-0.1, -0.01});
+    algo.Init("* =>[KNN 1 @pos $vec]", &params);
+    EXPECT_THAT(algo.Search(&indices).ids, testing::UnorderedElementsAre(2));
+  }
+
+  // Point up
+  {
+    params["vec"] = ToBytes({0, 5});
+    algo.Init("* =>[KNN 1 @pos $vec]", &params);
+    EXPECT_THAT(algo.Search(&indices).ids, testing::UnorderedElementsAre(3));
+  }
+
+  // Point right
+  {
+    params["vec"] = ToBytes({0.2, 0.05});
+    algo.Init("* =>[KNN 1 @pos $vec]", &params);
+    EXPECT_THAT(algo.Search(&indices).ids, testing::UnorderedElementsAre(0));
+  }
+}
+
 TEST_P(KnnTest, AutoResize) {
   // Make sure index resizes automatically even with a small initial capacity
   const size_t kInitialCapacity = 5;

src/core/search/vector_utils.cc

@@ -24,6 +24,7 @@ __attribute__((optimize("fast-math"))) float L2Distance(const float* u, const fl
   return sqrt(sum);
 }
 
+// TODO: Normalize vectors ahead if cosine distance is used
 __attribute__((optimize("fast-math"))) float CosineDistance(const float* u, const float* v,
                                                             size_t dims) {
   float sum_uv = 0, sum_uu = 0, sum_vv = 0;
@@ -34,7 +35,7 @@ __attribute__((optimize("fast-math"))) float CosineDistance(const float* u, cons
   }
 
   if (float denom = sum_uu * sum_vv; denom != 0.0f)
-    return sum_uv / sqrt(denom);
+    return 1 - sum_uv / sqrt(denom);
   return 0.0f;
 }
 

src/facade/cmd_arg_parser.cc

@@ -34,18 +34,30 @@ CmdArgParser::CheckProxy::operator bool() const {
   return true;
 }
 
-template <typename T> T CmdArgParser::NextProxy::Int() {
+template <typename T> T CmdArgParser::Num(size_t idx) {
+  auto arg = SafeSV(idx);
   T out;
-  if (absl::SimpleAtoi(operator std::string_view(), &out))
-    return out;
-  parser_->Report(INVALID_INT, idx_);
-  return T{0};
+  if constexpr (std::is_same_v<T, float>) {
+    if (absl::SimpleAtof(arg, &out))
+      return out;
+  } else if constexpr (std::is_same_v<T, double>) {
+    if (absl::SimpleAtod(arg, &out))
+      return out;
+  } else if constexpr (std::is_integral_v<T>) {
+    if (absl::SimpleAtoi(arg, &out))
+      return out;
+  }
+
+  Report(INVALID_INT, idx);
+  return {};
 }
 
-template uint64_t CmdArgParser::NextProxy::Int<uint64_t>();
-template int64_t CmdArgParser::NextProxy::Int<int64_t>();
-template uint32_t CmdArgParser::NextProxy::Int<uint32_t>();
-template int32_t CmdArgParser::NextProxy::Int<int32_t>();
+template float CmdArgParser::Num<float>(size_t);
+template double CmdArgParser::Num<double>(size_t);
+template uint64_t CmdArgParser::Num<uint64_t>(size_t);
+template int64_t CmdArgParser::Num<int64_t>(size_t);
+template uint32_t CmdArgParser::Num<uint32_t>(size_t);
+template int32_t CmdArgParser::Num<int32_t>(size_t);
 
 ErrorReply CmdArgParser::ErrorInfo::MakeReply() const {
   switch (type) {