Added Rescore Query

changes.txt

@@ -1,5 +1,9 @@
 CHANGES
 
+2013-08-14
+- Added Rescore
+- Set up structure to be easily extended for more types of rescores
+
 2013-08-13
 - Add a getQuery method on the FilteredQuery object
 

lib/Elastica/Query.php

@@ -8,6 +8,7 @@
 use Elastica\Query\AbstractQuery;
 use Elastica\Query\MatchAll;
 use Elastica\Query\QueryString;
+use Elastica\Rescore\AbstractRescore;
 
 /**
  * Elastica query object
@@ -100,6 +101,18 @@ public function setQuery(AbstractQuery $query)
         return $this->setParam('query', $query->toArray());
     }
 
+    /**
+     * Sets the rescore
+     *
+     * @param  \Elastica\Rescore\AbstractRescore $rescore Rescore object
+     * @return \Elastica\Query               Query object
+     */
+    public function setRescore(AbstractRescore $rescore)
+    {
+        $data = $rescore->toArray();
+        return $this->setParam('rescore', $data['rescore']);
+    }
+
     /**
      * Gets the query array
      *

lib/Elastica/Rescore/QueryRescore.php

@@ -0,0 +1,77 @@
+<?php
+
+namespace Elastica\Rescore;
+
+use Elastica\Query as BaseQuery;
+
+/**
+ * Rescore query
+ *
+ * @category Xodoa
+ * @package Elastica
+ * @author Jason Hu <mjhu91@gmail.com>
+ * @link http://www.elasticsearch.org/guide/reference/api/search/rescore/
+ */
+class QueryRescore extends AbstractRescore
+{
+
+    /**
+     * Gets the query param
+     *
+     * @return array
+     */
+    protected function getQueryParam()
+    {
+        if (!$this->hasParam('query'))
+        {
+            $this->setParam('query', array());
+        }
+
+        return $this->getParam('query');
+    }
+
+    /**
+     * Sets query object
+     *
+     * @param  string|\Elastica\Query|\Elastica\Query\AbstractQuery $query
+     * @return \Elastica\Rescore
+     */
+    public function setQuery($rescoreQuery)
+    {
+        $query = BaseQuery::create($rescoreQuery);
+        $data = $query->toArray();
+        $query = $this->getQueryParam();
+
+        $query['rescore_query'] = $data['query'];
+
+        return $this->setParam('query', $query);
+    }
+
+    /**
+     * Sets query_weight
+     *
+     * @param float $weight
+     * @return \Elastica\Rescore
+     */
+    public function setQueryWeight($weight)
+    {
+        $query = $this->getQueryParam();
+        $query['query_weight'] = $weight;
+
+        return $this->setParam('query', $query);
+    }
+
+    /**
+     * Sets rescore_query_weight
+     *
+     * @param float $size
+     * @return \Elastica\Rescore
+     */
+    public function setRescoreQueryWeight($weight)
+    {
+        $query = $this->getQueryParam();
+        $query['rescore_query_weight'] = $weight;
+
+        return $this->setParam('query', $query);
+    }
+}

lib/Elastica/Search.php

@@ -3,6 +3,7 @@
 namespace Elastica;
 
 use Elastica\Exception\InvalidException;
+use Elastica\Rescore\AbstractRescore;
 
 /**
  * Elastica search object
@@ -162,6 +163,24 @@ public function setQuery($query)
         return $this;
     }
 
+    /**
+     * Sets the rescore
+     *
+     * @param  \Elastica\Rescore\AbstractRescore $rescore Rescore object
+     * @return \Elastica\Query               Query object
+     */
+    public function setRescore(AbstractRescore $rescore)
+    {
+        if (!isset($this->_query))
+        {
+            throw new InvalidException('Please set query before setting rescore');
+        }
+        $data = $rescore->toArray();
+        $this->_query->setParam('rescore', $data['rescore']);
+
+        return $this;
+    }
+
     /**
      * @param  string          $key
      * @param  mixed           $value

test/lib/Elastica/Test/QueryTest.php

@@ -8,6 +8,9 @@
 use Elastica\Query\Term;
 use Elastica\Query\Text;
 use Elastica\Query;
+use Elastica\Query\Match;
+use Elastica\Rescore\QueryRescore;
+use Elastica\Client;
 use Elastica\Test\Base as BaseTest;
 
 class QueryTest extends BaseTest
@@ -174,4 +177,57 @@ public function testGetQuery()
 
         $this->assertEquals($termQuery->toArray(), $query->getQuery());
     }
+
+    public function testSetRescore()
+    {
+        $client = new Client(array('connections' => array(array('host' => 'localhost', 'port' => 9200))));
+        $index = $client->getIndex('elastica_test1');
+        $index->create(array(), true);
+
+        $mainQuery = new Match();
+        $mainQuery = $mainQuery->setFieldQuery('test1', 'foo');
+        $queryRescore = new QueryRescore();
+        $rescoreQuery = new Term();
+        $rescoreQuery = $rescoreQuery->setTerm('test2', 'bar', 2);
+        $queryRescore->setQuery($rescoreQuery);
+        $queryRescore->setWindowSize(50);
+        $queryRescore->setQueryWeight(.7);
+        $queryRescore->setRescoreQueryWeight(1.2);
+
+        $query = Query::create($mainQuery);
+        $query = $query->setRescore($queryRescore);
+
+        $response = $index->search($query)->getResponse();
+        $data = $response->getData();
+        $shards = $data['_shards'];
+        $failed = $shards['failed'];
+        $expectedData = array(
+            'query' => array(
+                'match' => array(
+                    'test1' => array(
+                        'query' => 'foo',
+                    ),
+                ),
+            ),
+            'rescore' => array(
+                'window_size' => 50,
+                'query' => array(
+                    'rescore_query' => array(
+                        'term' => array(
+                            'test2' => array(
+                                'value' => 'bar',
+                                'boost' => 2,
+                            ),
+                        ),
+                    ),
+                    'query_weight' => 0.7,
+                    'rescore_query_weight' => 1.2
+                ),
+            ),
+        );
+
+        $this->assertEquals($expectedData, $query->toArray());
+        $this->assertEquals(false, $response->hasError());
+        $this->assertEquals(0, $failed);
+    }
 }

test/lib/Elastica/Test/Rescore/QueryRescoreTest.php

@@ -0,0 +1,99 @@
+<?php
+
+namespace Elastica\Test;
+
+use Elastica\Rescore\QueryRescore;
+use Elastica\Query\Term;
+use Elastica\Query\Match;
+use Elastica\Client;
+use Elastica\Query;
+use Elastica\Test\Base as BaseTest;
+
+class QueryRescoreTest extends BaseTest
+{
+    public function testToArray()
+    {
+        $query = new QueryRescore();
+        $rescoreQuery = new Term();
+        $rescoreQuery = $rescoreQuery->setTerm('test2', 'bar', 2);
+        $query->setQuery($rescoreQuery);
+
+        $data = $query->toArray();
+        $expected = array(
+            'rescore' => array(
+                'query' => array(
+                    'rescore_query' => array(
+                        'term' => array(
+                            'test2' => array(
+                                'value' => 'bar',
+                                'boost' => 2,
+                            ),
+                        ),
+                    ),
+                ),
+            ),
+        );
+
+        $this->assertEquals($expected, $data);
+    }
+
+    public function testSetSize()
+    {
+        $query = new QueryRescore();
+        $rescoreQuery = new Term();
+        $rescoreQuery = $rescoreQuery->setTerm('test2', 'bar', 2);
+        $query->setQuery($rescoreQuery);
+        $query->setWindowSize(50);
+
+        $data = $query->toArray();
+        $expected = array(
+            'rescore' => array(
+                'window_size' => 50,
+                'query' => array(
+                    'rescore_query' => array(
+                        'term' => array(
+                            'test2' => array(
+                                'value' => 'bar',
+                                'boost' => 2,
+                            ),
+                        ),
+                    ),
+                ),
+            ),
+        );
+
+        $this->assertEquals($expected, $data);
+    }
+
+    public function testSetWeights()
+    {
+        $query = new QueryRescore();
+        $rescoreQuery = new Term();
+        $rescoreQuery = $rescoreQuery->setTerm('test2', 'bar', 2);
+        $query->setQuery($rescoreQuery);
+        $query->setWindowSize(50);
+        $query->setQueryWeight(.7);
+        $query->setRescoreQueryWeight(1.2);
+
+        $data = $query->toArray();
+        $expected = array(
+            'rescore' => array(
+                'window_size' => 50,
+                'query' => array(
+                    'rescore_query' => array(
+                        'term' => array(
+                            'test2' => array(
+                                'value' => 'bar',
+                                'boost' => 2,
+                            ),
+                        ),
+                    ),
+                    'query_weight' => 0.7,
+                    'rescore_query_weight' => 1.2
+                ),
+            ),
+        );
+
+        $this->assertEquals($expected, $data);
+    }
+}

test/lib/Elastica/Test/SearchTest.php

@@ -12,6 +12,9 @@
 use Elastica\Query;
 use Elastica\Script;
 use Elastica\Search;
+use Elastica\Query\Match;
+use Elastica\Query\Term;
+use Elastica\Rescore\QueryRescore;
 use Elastica\Test\Base as BaseTest;
 use Elastica\Type;
 
@@ -484,4 +487,59 @@ public function testCount() {
 		$this->assertInstanceOf('\Elastica\ResultSet', $result2);
 		$this->assertEquals(1, $result2->getTotalHits());
 	}
+
+    public function testSetRescore()
+    {
+        $client = new Client(array('connections' => array(array('host' => 'localhost', 'port' => 9200))));
+        $index = $client->getIndex('elastica_test1');
+        $index->create(array(), true);
+        $search = new Search($client);
+        $search = $search->addIndex($index);
+
+        $mainQuery = new Match();
+        $mainQuery = $mainQuery->setFieldQuery('test1', 'foo');
+        $queryRescore = new QueryRescore();
+        $rescoreQuery = new Term();
+        $rescoreQuery = $rescoreQuery->setTerm('test2', 'bar', 2);
+        $queryRescore->setQuery($rescoreQuery);
+        $queryRescore->setWindowSize(50);
+        $queryRescore->setQueryWeight(.7);
+        $queryRescore->setRescoreQueryWeight(1.2);
+
+        $search = $search->setQuery($mainQuery);
+        $search = $search->setRescore($queryRescore);
+
+        $response = $search->search()->getResponse();
+        $data = $response->getData();
+        $shards = $data['_shards'];
+        $failed = $shards['failed'];
+        $expectedData = array(
+            'query' => array(
+                'match' => array(
+                    'test1' => array(
+                        'query' => 'foo',
+                    ),
+                ),
+            ),
+            'rescore' => array(
+                'window_size' => 50,
+                'query' => array(
+                    'rescore_query' => array(
+                        'term' => array(
+                            'test2' => array(
+                                'value' => 'bar',
+                                'boost' => 2,
+                            ),
+                        ),
+                    ),
+                    'query_weight' => 0.7,
+                    'rescore_query_weight' => 1.2
+                ),
+            ),
+        );
+
+        $this->assertEquals($expectedData, $search->getQuery()->toArray());
+        $this->assertEquals(false, $response->hasError());
+        $this->assertEquals(0, $failed);
+    }
 }