Finished batch query feature.

docs/guide/query-builder.md

@@ -359,7 +359,7 @@ foreach ($query->batch(10) as $users) {
 The method [[yii\db\Query::batch()]] returns an [[yii\db\BatchQueryResult]] object which implements
 the `Iterator` interface and thus can be used in the `foreach` construct. For each iterator,
 it returns an array of query result. The size of the array is determined by the so-called batch
-size, which is the first parameter (defaults to 10) to the method.
+size, which is the first parameter (defaults to 100) to the method.
 
 Compared to the `$query->all()` call, the above code only loads 10 rows of data at a time into the memory.
 If you process the data and then discard it right away, the batch query can help keep the memory usage under a limit.
@@ -385,3 +385,21 @@ foreach ($query->batch(1) as $user) {
 	// $user represents a row from the user table
 }
 ```
+
+If you specify the query result to be indexed by some column via [[yii\db\Query::indexBy()]], the batch query
+will still keep the proper index. For example,
+
+```php
+use yii\db\Query;
+
+$query = (new Query)
+	->from('tbl_user')
+	->indexBy('username');
+
+foreach ($query->batch(10) as $users) {
+	// $users is indexed by the "username" column
+}
+
+foreach ($query->each() as $username => $user) {
+}
+```

framework/CHANGELOG.md

@@ -109,7 +109,7 @@ Yii Framework 2 Change Log
 - Enh #2240: Improved `yii\web\AssetManager::publish()`, `yii\web\AssetManager::getPublishedPath()` and `yii\web\AssetManager::getPublishedUrl()` to support aliases (vova07)
 - Enh #2325: Adding support for the `X-HTTP-Method-Override` header in `yii\web\Request::getMethod()` (pawzar)
 - Enh #2364: Take into account current error reporting level in error handler (gureedo)
-- Enh #2409: Added support for fetching data from database in batches (nineinchnick, qiangxue)
+- Enh #2387: Added support for fetching data from database in batches (nineinchnick, qiangxue)
 - Enh #2417: Added possibility to set `dataType` for `$.ajax` call in yii.activeForm.js (Borales)
 - Enh: Added support for using arrays as option values for console commands (qiangxue)
 - Enh: Added `favicon.ico` and `robots.txt` to default application templates (samdark)

framework/db/ActiveQuery.php

@@ -67,6 +67,9 @@ class ActiveQuery extends Query implements ActiveQueryInterface
 		return parent::all($db);
 	}
 
+	/**
+	 * @inheritdoc
+	 */
 	public function prepareResult($rows)
 	{
 		if (empty($rows)) {

framework/db/BatchQueryResult.php

@@ -10,39 +10,61 @@ namespace yii\db;
 use yii\base\Object;
 
 /**
- * BatchQueryResult represents the query result from which you can retrieve the data in batches.
+ * BatchQueryResult represents a batch query from which you can retrieve data in batches.
+ *
+ * You usually do not instantiate BatchQueryResult directly. Instead, you obtain it by
+ * calling [[Query::batch()]] or [[Query::each()]]. Because BatchQueryResult implements the `Iterator` interface,
+ * you can iterate it to obtain a batch of data in each iteration. For example,
+ *
+ * ```php
+ * $query = (new Query)->from('tbl_user');
+ * foreach ($query->batch() as $i => $users) {
+ *     // $users represents the rows in the $i-th batch
+ * }
+ * ```
  *
- * BatchQueryResult is mainly used with [[Query::batch()]].
  * @author Qiang Xue <qiang.xue@gmail.com>
  * @since 2.0
  */
 class BatchQueryResult extends Object implements \Iterator
 {
 	/**
-	 * @var Connection
+	 * @var Connection the DB connection to be used when performing batch query.
+	 * If null, the "db" application component will be used.
 	 */
 	public $db;
 	/**
-	 * @var Query
+	 * @var Query the query object associated with this batch query.
+	 * Do not modify this property directly unless after [[reset()]] is called explicitly.
 	 */
 	public $query;
 	/**
-	 * @var integer
+	 * @var DataReader the data reader associated with this batch query.
+	 * Do not modify this property directly unless after [[reset()]] is called explicitly.
 	 */
-	public $batchSize = 10;
+	public $dataReader;
 	/**
-	 * @var DataReader
+	 * @var integer the number of rows to be returned in each batch.
 	 */
-	public $dataReader;
+	public $batchSize = 100;
 
 	private $_data;
+	private $_key;
 	private $_index = -1;
 
+	/**
+	 * Destructor.
+	 */
 	public function __destruct()
 	{
+		// make sure cursor is closed
 		$this->reset();
 	}
 
+	/**
+	 * Resets the batch query.
+	 * This method will clean up the existing batch query so that a new batch query can be performed.
+	 */
 	public function reset()
 	{
 		if ($this->dataReader !== null) {
@@ -64,19 +86,19 @@ class BatchQueryResult extends Object implements \Iterator
 	}
 
 	/**
-	 * Returns the index of the current row.
+	 * Returns the index of the current dataset.
 	 * This method is required by the interface Iterator.
 	 * @return integer the index of the current row.
 	 */
 	public function key()
 	{
-		return $this->_index;
+		return $this->batchSize == 1 ? $this->_key : $this->_index;
 	}
 
 	/**
-	 * Returns the current row.
+	 * Returns the current dataset.
 	 * This method is required by the interface Iterator.
-	 * @return mixed the current row.
+	 * @return mixed the current dataset.
 	 */
 	public function current()
 	{
@@ -84,7 +106,7 @@ class BatchQueryResult extends Object implements \Iterator
 	}
 
 	/**
-	 * Moves the internal pointer to the next row.
+	 * Moves the internal pointer to the next dataset.
 	 * This method is required by the interface Iterator.
 	 */
 	public function next()
@@ -106,15 +128,17 @@ class BatchQueryResult extends Object implements \Iterator
 		} else {
 			$this->_data = $this->query->prepareResult($rows);
 			if ($this->batchSize == 1) {
-				$this->_data = reset($this->_data);
+				$row = reset($this->_data);
+				$this->_key = key($this->_data);
+				$this->_data = $row;
 			}
 		}
 	}
 
 	/**
-	 * Returns whether there is a row of data at current position.
+	 * Returns whether there is a valid dataset at the current position.
 	 * This method is required by the interface Iterator.
-	 * @return boolean whether there is a row of data at current position.
+	 * @return boolean whether there is a valid dataset at the current position.
 	 */
 	public function valid()
 	{

framework/db/Query.php

@@ -123,7 +123,28 @@ class Query extends Component implements QueryInterface
 		return $db->createCommand($sql, $params);
 	}
 
-	public function batch($size = 10, $db = null)
+	/**
+	 * Starts a batch query.
+	 *
+	 * A batch query supports fetching data in batches, which can keep the memory usage under a limit.
+	 * This method will return a [[BatchQueryResult]] object which implements the `Iterator` interface
+	 * and can be traversed to retrieve the data in batches.
+	 *
+	 * For example,
+	 *
+	 * ```php
+	 * $query = (new Query)->from('tbl_user');
+	 * foreach ($query->batch() as $rows) {
+	 *     // $rows is an array of 10 or fewer rows from tbl_user
+	 * }
+	 * ```
+	 *
+	 * @param integer $size the number of records to be fetched in each batch.
+	 * @param Connection $db the database connection. If not set, the "db" application component will be used.
+	 * @return BatchQueryResult the batch query result. It implements the `Iterator` interface
+	 * and can be traversed to retrieve the data in batches.
+	 */
+	public function batch($size = 100, $db = null)
 	{
 		return Yii::createObject([
 			'class' => BatchQueryResult::className(),
@@ -133,6 +154,13 @@ class Query extends Component implements QueryInterface
 		]);
 	}
 
+	/**
+	 * Starts a batch query and retrieves data row by row.
+	 * This method is a shortcut to [[batch()]] with batch size fixed to be 1.
+	 * @param Connection $db the database connection. If not set, the "db" application component will be used.
+	 * @return BatchQueryResult the batch query result. It implements the `Iterator` interface
+	 * and can be traversed to retrieve the data in batches.
+	 */
 	public function each($db = null)
 	{
 		return $this->batch(1, $db);
@@ -150,6 +178,13 @@ class Query extends Component implements QueryInterface
 		return $this->prepareResult($rows);
 	}
 
+	/**
+	 * Converts the raw query results into the format as specified by this query.
+	 * This method is internally used to convert the data fetched from database
+	 * into the format as required by this query.
+	 * @param array $rows the raw query result from database
+	 * @return array the converted query result
+	 */
 	public function prepareResult($rows)
 	{
 		if ($this->indexBy === null) {

tests/unit/framework/db/BatchQueryResultTest.php

@@ -84,6 +84,7 @@ class BatchQueryResultTest extends DatabaseTestCase
 		$this->assertEquals('user2', $allRows[1]['name']);
 		$this->assertEquals('user3', $allRows[2]['name']);
 
+		// each
 		$query = new Query();
 		$query->from('tbl_customer')->orderBy('id');
 		$allRows = [];
@@ -94,6 +95,18 @@ class BatchQueryResultTest extends DatabaseTestCase
 		$this->assertEquals('user1', $allRows[0]['name']);
 		$this->assertEquals('user2', $allRows[1]['name']);
 		$this->assertEquals('user3', $allRows[2]['name']);
+
+		// each with key
+		$query = new Query();
+		$query->from('tbl_customer')->orderBy('id')->indexBy('name');
+		$allRows = [];
+		foreach ($query->each($db) as $key => $row) {
+			$allRows[$key] = $row;
+		}
+		$this->assertEquals(3, count($allRows));
+		$this->assertEquals('address1', $allRows['user1']['address']);
+		$this->assertEquals('address2', $allRows['user2']['address']);
+		$this->assertEquals('address3', $allRows['user3']['address']);
 	}
 
 	public function testActiveQuery()