Knn functional tests

elasticsearch version changed as 8.4.0 for github ci/cd
handcraftedinthealps · Oct 8, 2023 · 96d4f5a · 96d4f5a
1 parent 48b8e3d
commit 96d4f5a
Show file tree

Hide file tree

Showing 4 changed files with 109 additions and 15 deletions.
diff --git a/.github/workflows/test-application.yaml b/.github/workflows/test-application.yaml
@@ -50,7 +50,7 @@ jobs:
 
         services:
             elasticsearch:
-                image: elasticsearch:8.0.0
+                image: elasticsearch:8.4.0
                 ports:
                     - 9200:9200
                 env:

diff --git a/src/SearchEndpoint/KnnEndpoint.php b/src/SearchEndpoint/KnnEndpoint.php
@@ -42,14 +42,22 @@ public function normalize(
         $format = null,
         array $context = []
     ): array|string|int|float|bool {
-        $output = [];
-        if (count($this->getAll()) > 0) {
+        $knns = $this->getAll();
+        if (count($knns) === 1) {
             /** @var Knn $knn */
-            foreach ($this->getAll() as $knn) {
+            $knn = array_values($knns)[0];
+            return $knn->toArray();
+        }
+
+        if (count($knns) > 1) {
+            $output = [];
+            /** @var Knn $knn */
+            foreach ($knns as $knn) {
                 $output[] = $knn->toArray();
             }
+            return $output;
         }
 
-        return $output;
+        return [];
     }
 }
diff --git a/tests/Functional/AbstractElasticsearchTestCase.php b/tests/Functional/AbstractElasticsearchTestCase.php
@@ -21,7 +21,7 @@ abstract class AbstractElasticsearchTestCase extends TestCase
     /**
      * Test index name in the elasticsearch.
      */
-    const INDEX_NAME = 'elasticsaerch-dsl-test';
+    const INDEX_NAME = 'elasticsearch-dsl-test';
 
     /**
      * @var Client
@@ -35,17 +35,18 @@ protected function setUp(): void
     {
         parent::setUp();
 
-        $this->client = ClientBuilder::create()->build();
+        $this->client = ClientBuilder::create()
+            ->build();
         $this->deleteIndex();
 
-        $this->client->indices()->create(
-            array_filter(
-                [
-                    'index' => self::INDEX_NAME,
-                    'mapping' => $this->getMapping()
-                ]
-            )
-        );
+        $params = [
+            'index' => self::INDEX_NAME
+        ];
+        if ($this->getMapping()) {
+            $params['body']['mappings'] = $this->getMapping();
+        }
+
+        $response = $this->client->indices()->create($params);
 
         $bulkBody = [];
 

diff --git a/tests/Functional/Knn/KnnTest.php b/tests/Functional/Knn/KnnTest.php
@@ -0,0 +1,85 @@
+<?php
+/**
+ * @since     Feb 2022
+ * @author    Haydar KULEKCI <haydarkulekci@gmail.com>
+ */
+
+namespace ONGR\ElasticsearchDSL\Tests\Functional\Knn;
+
+use ONGR\ElasticsearchDSL\Aggregation\Bucketing\DateHistogramAggregation;
+use ONGR\ElasticsearchDSL\Knn\Knn;
+use ONGR\ElasticsearchDSL\Query\TermLevel\TermQuery;
+use ONGR\ElasticsearchDSL\Search;
+use ONGR\ElasticsearchDSL\Tests\Functional\AbstractElasticsearchTestCase;
+
+class KnnTest extends AbstractElasticsearchTestCase
+{
+    /**
+     * {@inheritdoc}
+     */
+    protected function getDataArray(): array
+    {
+        return [
+            'knn_data' => [
+                'doc_1' => [
+                    'label' => 1,
+                    'vector_field' => [1, 2, 3],
+                ],
+                'doc_2' => [
+                    'label' => 1,
+                    'vector_field' => [1, 2, 4],
+                ],
+                'doc_3' => [
+                    'label' => 2,
+                    'vector_field' => [1, 2, 30],
+                ],
+            ]
+        ];
+    }
+
+    protected function getMapping(): array
+    {
+        return [
+            'properties' => [
+                'label' => [
+                    'type' => 'long'
+                ],
+                'vector_field' => [
+                    'type' => 'dense_vector',
+                    'similarity' => 'cosine',
+                    'index' => true,
+                    'dims' => 3
+                ]
+            ]
+        ];
+    }
+
+    /**
+     * Match all test
+     */
+    public function testKnnSearch(): void
+    {
+        $knn = new Knn('vector_field', [1, 2, 3], 1, 1);
+
+        $search = new Search();
+        $search->addKnn($knn);
+        $results = $this->executeSearch($search, true);
+        $this->assertCount(1, $results['hits']['hits']);
+        $this->assertEquals('doc_1', $results['hits']['hits'][0]['_id']);
+    }
+
+    /**
+     * Match all test
+     */
+    public function testKnnSearchWithFilter(): void
+    {
+        $knn = new Knn('vector_field', [1, 2, 3], 1, 1);
+        $knn->setFilter(new TermQuery('label', 2));
+
+        $search = new Search();
+        $search->addKnn($knn);
+        $results = $this->executeSearch($search, true);
+        $this->assertCount(1, $results['hits']['hits']);
+        $this->assertEquals('doc_3', $results['hits']['hits'][0]['_id']);
+    }
+}