1 files changed, 54 insertions, 14 deletions

diff --git a/engine/classes/ElggBatch.php b/engine/classes/ElggBatch.php
index 5d59425d0..d810ea066 100644
--- a/engine/classes/ElggBatch.php
+++ b/engine/classes/ElggBatch.php
@@ -150,6 +150,20 @@ class ElggBatch
 	private $incrementOffset = true;
 
 	/**
+	 * Entities that could not be instantiated during a fetch
+	 *
+	 * @var stdClass[]
+	 */
+	private $incompleteEntities = array();
+
+	/**
+	 * Total number of incomplete entities fetched
+	 *
+	 * @var int
+	 */
+	private $totalIncompletes = 0;
+
+	/**
 	 * Batches operations on any elgg_get_*() or compatible function that supports
 	 * an options array.
 	 *
@@ -222,16 +236,22 @@ class ElggBatch
 	}
 
 	/**
+	 * Tell the process that an entity was incomplete during a fetch
+	 *
+	 * @param stdClass $row
+	 *
+	 * @access private
+	 */
+	public function reportIncompleteEntity(stdClass $row) {
+		$this->incompleteEntities[] = $row;
+	}
+
+	/**
 	 * Fetches the next chunk of results
 	 *
 	 * @return bool
 	 */
 	private function getNextResultsChunk() {
-		// reset memory caches after first chunk load
-		if ($this->chunkIndex > 0) {
-			global $DB_QUERY_CACHE, $ENTITY_CACHE;
-			$DB_QUERY_CACHE = $ENTITY_CACHE = array();
-		}
 
 		// always reset results.
 		$this->results = array();
@@ -265,27 +285,47 @@ class ElggBatch
 		if ($this->incrementOffset) {
 			$offset = $this->offset + $this->retrievedResults;
 		} else {
-			$offset = $this->offset;
+			$offset = $this->offset + $this->totalIncompletes;
 		}
 
 		$current_options = array(
 			'limit' => $limit,
-			'offset' => $offset
+			'offset' => $offset,
+			'__ElggBatch' => $this,
 		);
 
 		$options = array_merge($this->options, $current_options);
-		$getter = $this->getter;
 
-		if (is_string($getter)) {
-			$this->results = $getter($options);
-		} else {
-			$this->results = call_user_func_array($getter, array($options));
+		$this->incompleteEntities = array();
+		$this->results = call_user_func_array($this->getter, array($options));
+
+		$num_results = count($this->results);
+		$num_incomplete = count($this->incompleteEntities);
+
+		$this->totalIncompletes += $num_incomplete;
+
+		if ($this->incompleteEntities) {
+			// pad the front of the results with nulls representing the incompletes
+			array_splice($this->results, 0, 0, array_pad(array(), $num_incomplete, null));
+			// ...and skip past them
+			reset($this->results);
+			for ($i = 0; $i < $num_incomplete; $i++) {
+				next($this->results);
+			}
 		}
 
 		if ($this->results) {
 			$this->chunkIndex++;
-			$this->resultIndex = 0;
-			$this->retrievedResults += count($this->results);
+
+			// let the system know we've jumped past the nulls
+			$this->resultIndex = $num_incomplete;
+
+			$this->retrievedResults += ($num_results + $num_incomplete);
+			if ($num_results == 0) {
+				// This fetch was *all* incompletes! We need to fetch until we can either
+				// offer at least one row to iterate over, or give up.
+				return $this->getNextResultsChunk();
+			}
 			return true;
 		} else {
 			return false;