diff --git a/src/InnerResultIterator.php b/src/InnerResultIterator.php
index 937e65e1..cb3d44ba 100644
--- a/src/InnerResultIterator.php
+++ b/src/InnerResultIterator.php
@@ -9,6 +9,7 @@
 use Psr\Log\LoggerInterface;
 use Psr\Log\NullLogger;
 use TheCodingMachine\TDBM\QueryFactory\SmartEagerLoad\PartialQueryFactory;
+use TheCodingMachine\TDBM\QueryFactory\SmartEagerLoad\Query\PartialQuery;
 use TheCodingMachine\TDBM\QueryFactory\SmartEagerLoad\StorageNode;
 use TheCodingMachine\TDBM\Utils\DbalUtils;
 
@@ -68,11 +69,13 @@ class InnerResultIterator implements \Iterator, \Countable, \ArrayAccess
      */
     private $logger;
     /**
-     * @var PartialQueryFactory|null
+     * @var PartialQuery|null
      */
-    private $partialQueryFactory;
-
-    protected $count = null;
+    private $partialQuery;
+    /**
+     * @var int|null
+     */
+    protected $count;
 
     private function __construct()
     {
@@ -96,7 +99,11 @@ public static function createInnerResultIterator(string $magicSql, array $parame
         $iterator->magicQuery = $magicQuery;
         $iterator->databasePlatform = $iterator->tdbmService->getConnection()->getDatabasePlatform();
         $iterator->logger = $logger;
-        $iterator->partialQueryFactory = $partialQueryFactory;
+        $partialQuery = null;
+        if ($iterator instanceof StorageNode && $partialQueryFactory !== null) {
+            $iterator->partialQuery = $partialQueryFactory->getPartialQuery($iterator, $magicQuery, $parameters);
+        }
+
         return $iterator;
     }
 
@@ -208,11 +215,6 @@ public function next()
                 $beansData[$columnDescriptor['tableGroup']][$columnDescriptor['table']][$columnDescriptor['column']] = $value;
             }
 
-            $partialQuery = null;
-            if ($this instanceof StorageNode && $this->partialQueryFactory !== null) {
-                $partialQuery = $this->partialQueryFactory->getPartialQuery($this, $this->magicQuery, $this->parameters);
-            }
-
             $reflectionClassCache = [];
             $firstBean = true;
             foreach ($beansData as $beanData) {
@@ -250,7 +252,7 @@ public function next()
                     // Let's bypass the constructor when creating the bean!
                     /** @var AbstractTDBMObject $bean */
                     $bean = $reflectionClassCache[$actualClassName]->newInstanceWithoutConstructor();
-                    $bean->_constructFromData($beanData, $this->tdbmService, $partialQuery);
+                    $bean->_constructFromData($beanData, $this->tdbmService, $this->partialQuery);
                 }
 
                 // The first bean is the one containing the main table.