diff --git a/src/QueryBuilder/AbstractDMLQueryBuilder.php b/src/QueryBuilder/AbstractDMLQueryBuilder.php
index 0db417372..645ea7b5f 100644
--- a/src/QueryBuilder/AbstractDMLQueryBuilder.php
+++ b/src/QueryBuilder/AbstractDMLQueryBuilder.php
@@ -142,6 +142,8 @@ public function upsert(
      * @param Traversable $rows The rows to be batch inserted into the table.
      *
      * @return array|Iterator The prepared rows.
+     *
+     * @psalm-return Iterator|array<iterable<array-key, mixed>>
      */
     protected function prepareTraversable(Traversable $rows): Iterator|array
     {
@@ -153,6 +155,7 @@ protected function prepareTraversable(Traversable $rows): Iterator|array
             return [];
         }
 
+        /** @var Iterator $rows */
         return $rows;
     }
 
@@ -210,6 +213,8 @@ protected function prepareBatchInsertValues(string $table, iterable $rows, array
      * @param string[] $columns The column names.
      *
      * @return string[] The column names.
+     *
+     * @psalm-param Iterator|non-empty-array<iterable<array-key, mixed>> $rows
      */
     protected function extractColumnNames(array|Iterator $rows, array $columns): array
     {
diff --git a/src/QueryBuilder/DMLQueryBuilderInterface.php b/src/QueryBuilder/DMLQueryBuilderInterface.php
index 24aa4c856..e00113a57 100644
--- a/src/QueryBuilder/DMLQueryBuilderInterface.php
+++ b/src/QueryBuilder/DMLQueryBuilderInterface.php
@@ -45,7 +45,7 @@ interface DMLQueryBuilderInterface
      * @return string The batch INSERT SQL statement.
      *
      * @psalm-param string[] $columns
-     * @psalm-param iterable<array-key, array<array-key, mixed>> $rows
+     * @psalm-param iterable<iterable<array-key, mixed>> $rows
      * @psalm-param ParamsType $params
      *
      * Note: