feat: Add ID-based model chunking (#10363)

Author: memleakd

system/BaseModel.php

@@ -600,6 +600,19 @@ abstract public function countAllResults(bool $reset = true, bool $test = false)
      */
     abstract public function chunk(int $size, Closure $userFunc);
 
+    /**
+     * Loops over records in batches ordered by the primary key.
+     * This method works only with DB calls.
+     *
+     * @param Closure(array<string, string>|object): mixed $userFunc
+     *
+     * @return void
+     *
+     * @throws DataException
+     * @throws InvalidArgumentException if $size is not a positive integer or the current query cannot be chunked by ID
+     */
+    abstract public function chunkById(int $size, Closure $userFunc);
+
     /**
      * Loops over records in batches, allowing you to operate on each chunk at a time.
      * This method works only with DB calls.
@@ -616,6 +629,22 @@ abstract public function chunk(int $size, Closure $userFunc);
      */
     abstract public function chunkRows(int $size, Closure $userFunc);
 
+    /**
+     * Loops over records in batches ordered by the primary key, allowing you to operate on each chunk at a time.
+     * This method works only with DB calls.
+     *
+     * This method calls the `$userFunc` with the chunk, instead of a single record as in `chunkById()`.
+     * This allows you to operate on multiple records at once, which can be more efficient for certain operations.
+     *
+     * @param Closure(list<array<string, string>>|list<object>): mixed $userFunc
+     *
+     * @return void
+     *
+     * @throws DataException
+     * @throws InvalidArgumentException if $size is not a positive integer or the current query cannot be chunked by ID
+     */
+    abstract public function chunkRowsById(int $size, Closure $userFunc);
+
     /**
      * Fetches the row of database.
      *

system/Database/BaseBuilder.php

@@ -2061,6 +2061,16 @@ public function offset(int $offset)
         return $this;
     }
 
+    /**
+     * Checks if the current query has a LIMIT, OFFSET or UNION clause.
+     *
+     * @internal This method is for internal Model use only.
+     */
+    public function hasLimitOffsetOrUnion(): bool
+    {
+        return $this->QBLimit !== false || $this->QBOffset !== false || $this->QBUnion !== [];
+    }
+
     /**
      * Generates a platform-specific LIMIT clause.
      */

system/Model.php

@@ -165,6 +165,7 @@ class Model extends BaseModel
         'getCompiledInsert',
         'getCompiledSelect',
         'getCompiledUpdate',
+        'hasLimitOffsetOrUnion',
     ];
 
     public function __construct(?ConnectionInterface $db = null, ?ValidationInterface $validation = null)
@@ -599,9 +600,7 @@ private function prepareSoftDeleteQuery(bool $reset): void
      */
     private function iterateChunks(int $size): Generator
     {
-        if ($size <= 0) {
-            throw new InvalidArgumentException('$size must be a positive integer.');
-        }
+        $this->assertValidChunkSize($size);
 
         $total  = $this->builder()->countAllResults(false);
         $offset = 0;
@@ -626,6 +625,89 @@ private function iterateChunks(int $size): Generator
         }
     }
 
+    /**
+     * Iterates over the result set in chunks of the specified size ordered by the primary key.
+     *
+     * @param int $size The number of records to retrieve in each chunk.
+     *
+     * @return Generator<list<array<string, string>>|list<object>>
+     */
+    private function iterateChunksById(int $size): Generator
+    {
+        $this->assertValidChunkSize($size);
+
+        if ($this->primaryKey === '') {
+            throw new InvalidArgumentException('ID-based chunking requires a primary key.');
+        }
+
+        $builder             = clone $this->builder();
+        $qualifiedPrimaryKey = $this->table . '.' . $this->primaryKey;
+        $lastPrimaryKey      = null;
+        $hasLastPrimaryKey   = false;
+
+        if ($builder->QBOrderBy !== []) {
+            throw new InvalidArgumentException('ID-based chunking cannot be used with orderBy().');
+        }
+
+        if ($builder->QBGroupBy !== []) {
+            throw new InvalidArgumentException('ID-based chunking cannot be used with groupBy().');
+        }
+
+        if ($builder->hasLimitOffsetOrUnion()) {
+            throw new InvalidArgumentException('ID-based chunking cannot be used with limit(), offset() or union().');
+        }
+
+        while (true) {
+            $chunkBuilder = clone $builder;
+
+            if ($this->tempUseSoftDeletes) {
+                $chunkBuilder->where($this->table . '.' . $this->deletedField, null);
+            }
+
+            if ($hasLastPrimaryKey) {
+                $chunkBuilder->where($qualifiedPrimaryKey . ' >', $lastPrimaryKey);
+            }
+
+            $rows = $chunkBuilder
+                ->orderBy($qualifiedPrimaryKey, 'ASC')
+                ->get($size);
+
+            if (! $rows) {
+                throw DataException::forEmptyDataset('chunkById');
+            }
+
+            $rows = $rows->getResult($this->tempReturnType);
+
+            if ($rows === []) {
+                return;
+            }
+
+            $lastPrimaryKey = $this->getIdValue($rows[array_key_last($rows)]);
+
+            if ($lastPrimaryKey === null) {
+                throw new InvalidArgumentException('The primary key must be selected for ID-based chunking.');
+            }
+
+            $hasLastPrimaryKey = true;
+
+            yield $rows;
+
+            if (count($rows) < $size) {
+                return;
+            }
+        }
+    }
+
+    /**
+     * Asserts the chunk size is valid.
+     */
+    private function assertValidChunkSize(int $size): void
+    {
+        if ($size <= 0) {
+            throw new InvalidArgumentException('$size must be a positive integer.');
+        }
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -640,6 +722,20 @@ public function chunk(int $size, Closure $userFunc)
         }
     }
 
+    /**
+     * {@inheritDoc}
+     */
+    public function chunkById(int $size, Closure $userFunc)
+    {
+        foreach ($this->iterateChunksById($size) as $rows) {
+            foreach ($rows as $row) {
+                if ($userFunc($row) === false) {
+                    return;
+                }
+            }
+        }
+    }
+
     /**
      * {@inheritDoc}
      */
@@ -652,6 +748,18 @@ public function chunkRows(int $size, Closure $userFunc): void
         }
     }
 
+    /**
+     * {@inheritDoc}
+     */
+    public function chunkRowsById(int $size, Closure $userFunc): void
+    {
+        foreach ($this->iterateChunksById($size) as $rows) {
+            if ($userFunc($rows) === false) {
+                return;
+            }
+        }
+    }
+
     /**
      * Provides a shared instance of the Query Builder.
      *