diff --git a/framework/yii/elasticsearch/ActiveQuery.php b/framework/yii/elasticsearch/ActiveQuery.php
index d2b32f2..c734b34 100644
--- a/framework/yii/elasticsearch/ActiveQuery.php
+++ b/framework/yii/elasticsearch/ActiveQuery.php
@@ -89,6 +89,12 @@ class ActiveQuery extends Query implements ActiveQueryInterface
 			return [];
 		}
 		$models = $this->createModels($result['hits']);
+		if ($this->asArray) {
+			foreach($models as $key => $model) {
+				$models[$key] = $model['_source'];
+				$models[$key]['primaryKey'] = $model['_id'];
+			}
+		}
 		if (!empty($this->with)) {
 			$this->findWith($this->with, $models);
 		}
@@ -107,11 +113,13 @@ class ActiveQuery extends Query implements ActiveQueryInterface
 	{
 		$command = $this->createCommand($db);
 		$result = $command->queryOne();
-		if ($result['total'] == 0) {
+		if ($result['total'] == 0 || empty($result['hits'])) {
 			return null;
 		}
 		if ($this->asArray) {
-			$model = reset($result['hits']);
+			$first = reset($result['hits']);
+			$model = $first['_source'];
+			$model['primaryKey'] = $first['_id'];
 		} else {
 			/** @var ActiveRecord $class */
 			$class = $this->modelClass;
diff --git a/framework/yii/elasticsearch/Query.php b/framework/yii/elasticsearch/Query.php
index d3e7ad0..6621219 100644
--- a/framework/yii/elasticsearch/Query.php
+++ b/framework/yii/elasticsearch/Query.php
@@ -117,16 +117,20 @@ class Query extends Component implements QueryInterface
 
 	/**
 	 * Returns the number of records.
-	 * @param string $q the COUNT expression. Defaults to '*'.
-	 * Make sure you properly quote column names in the expression.
+	 * @param string $q the COUNT expression. This parameter is ignored by this implementation.
 	 * @param Connection $db the database connection used to generate the SQL statement.
 	 * If this parameter is not given (or null), the `db` application component will be used.
 	 * @return integer number of records
 	 */
 	public function count($q = '*', $db = null)
 	{
-		$this->select = ["COUNT($q)"];
-		return $this->createCommand($db)->queryScalar();
+		$count = $this->createCommand($db)->queryCount()['total'];
+		if ($this->limit === null && $this->offset === null) {
+			return $count;
+		} elseif ($this->offset !== null) {
+			$count = $this->offset < $count ? $count - $this->offset : 0;
+		}
+		return $this->limit === null ? $count : ($this->limit > $count ? $count : $this->limit);
 	}
 
 
diff --git a/tests/unit/data/ar/elasticsearch/ActiveRecord.php b/tests/unit/data/ar/elasticsearch/ActiveRecord.php
index 3309004..6c4dff6 100644
--- a/tests/unit/data/ar/elasticsearch/ActiveRecord.php
+++ b/tests/unit/data/ar/elasticsearch/ActiveRecord.php
@@ -17,6 +17,9 @@ class ActiveRecord extends \yii\elasticsearch\ActiveRecord
 {
 	public static $db;
 
+	/**
+	 * @return \yii\elasticsearch\Connection
+	 */
 	public static function getDb()
 	{
 		return self::$db;
diff --git a/tests/unit/framework/elasticsearch/ActiveRecordTest.php b/tests/unit/framework/elasticsearch/ActiveRecordTest.php
index 89383ff..ce5c7ce 100644
--- a/tests/unit/framework/elasticsearch/ActiveRecordTest.php
+++ b/tests/unit/framework/elasticsearch/ActiveRecordTest.php
@@ -92,16 +92,18 @@ class ActiveRecordTest extends ElasticSearchTestCase
 //		$orderItem->setAttributes(array('order_id' => 3, 'item_id' => 2, 'quantity' => 1, 'subtotal' => 40.0), false);
 //		$orderItem->save(false);
 
-		for($n = 0; $n < 20; $n++) {
-			$r = $db->http()->post('_count')->send();
-			$c = Json::decode($r->getBody(true));
-			if ($c['count'] != 11) {
-				usleep(100000);
-			} else {
-				return;
-			}
-		}
-		throw new \Exception('Unable to initialize elasticsearch data.');
+		Customer::getDb()->createCommand()->flushIndex();
+
+//		for($n = 0; $n < 20; $n++) {
+//			$r = $db->http()->post('_count')->send();
+//			$c = Json::decode($r->getBody(true));
+//			if ($c['count'] != 11) {
+//				usleep(100000);
+//			} else {
+//				return;
+//			}
+//		}
+//		throw new \Exception('Unable to initialize elasticsearch data.');
 	}
 
 	public function testFind()
@@ -127,22 +129,24 @@ class ActiveRecordTest extends ElasticSearchTestCase
 		$this->assertNull($customer);
 
 		// query scalar
-		$customerName = Customer::find()->where(array('id' => 2))->scalar('name');
-		$this->assertEquals('user2', $customerName);
+		$customerName = Customer::find()->where(array('status' => 2))->scalar('name');
+		$this->assertEquals('user3', $customerName);
 
 		// find by column values
-		$customer = Customer::find(array('id' => 2, 'name' => 'user2'));
+		$customer = Customer::find(array('name' => 'user2'));
 		$this->assertTrue($customer instanceof Customer);
 		$this->assertEquals('user2', $customer->name);
-		$customer = Customer::find(array('id' => 2, 'name' => 'user1'));
+		$customer = Customer::find(array('name' => 'user1', 'id' => 2));
 		$this->assertNull($customer);
-		$customer = Customer::find(array('id' => 5));
+		$customer = Customer::find(array('primaryKey' => 5));
+		$this->assertNull($customer);
+		$customer = Customer::find(array('name' => 'user5'));
 		$this->assertNull($customer);
 
 		// find by attributes
 		$customer = Customer::find()->where(array('name' => 'user2'))->one();
 		$this->assertTrue($customer instanceof Customer);
-		$this->assertEquals(2, $customer->id);
+		$this->assertEquals('user2', $customer->name);
 
 		// find count, sum, average, min, max, scalar
 		$this->assertEquals(3, Customer::find()->count());
@@ -156,13 +160,13 @@ class ActiveRecordTest extends ElasticSearchTestCase
 //		$this->assertEquals(2, Customer::find()->active()->count());
 
 		// asArray
-		$customer = Customer::find()->where(array('id' => 2))->asArray()->one();
+		$customer = Customer::find()->where(array('name' => 'user2'))->asArray()->one();
 		$this->assertEquals(array(
-			'id' => '2',
 			'email' => 'user2@example.com',
 			'name' => 'user2',
 			'address' => 'address2',
 			'status' => '1',
+			'primaryKey' => 2,
 		), $customer);
 
 		// indexBy
@@ -174,13 +178,13 @@ class ActiveRecordTest extends ElasticSearchTestCase
 
 		// indexBy callable
 		$customers = Customer::find()->indexBy(function ($customer) {
-			return $customer->id . '-' . $customer->name;
+			return $customer->status . '-' . $customer->name;
 //		})->orderBy('id')->all();
 		})->all();
 		$this->assertEquals(3, count($customers));
 		$this->assertTrue($customers['1-user1'] instanceof Customer);
-		$this->assertTrue($customers['2-user2'] instanceof Customer);
-		$this->assertTrue($customers['3-user3'] instanceof Customer);
+		$this->assertTrue($customers['1-user2'] instanceof Customer);
+		$this->assertTrue($customers['2-user3'] instanceof Customer);
 	}
 
 	public function testFindCount()