diff --git a/composer.json b/composer.json
index 9b8ffe3..e5c89f4 100644
--- a/composer.json
+++ b/composer.json
@@ -37,7 +37,7 @@
         "phpunit/phpunit": "^9.5",
         "roave/infection-static-analysis-plugin": "^1.16",
         "spatie/phpunit-watcher": "^1.23",
-        "vimeo/psalm": "^4.18"
+        "vimeo/psalm": "^4.30|^5.14"
     },
     "autoload": {
         "psr-4": {
diff --git a/psalm.xml b/psalm.xml
index 2df6349..4ba0a00 100644
--- a/psalm.xml
+++ b/psalm.xml
@@ -1,14 +1,20 @@
 <?xml version="1.0"?>
 <psalm
     errorLevel="2"
+    findUnusedBaselineEntry="true"
+    findUnusedCode="false"
     xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
     xmlns="https://getpsalm.org/schema/config"
     xsi:schemaLocation="https://getpsalm.org/schema/config vendor/vimeo/psalm/config.xsd"
 >
     <projectFiles>
         <directory name="src" />
+        <ignoreFiles>
+            <directory name="vendor" />
+        </ignoreFiles>
     </projectFiles>
     <issueHandlers>
+        <MixedAssignment errorLevel="suppress" />
         <PropertyNotSetInConstructor>
             <errorLevel type="suppress">
                 <file name="src/Filter/*.php" />
diff --git a/src/.gitkeep b/src/.gitkeep
deleted file mode 100644
index e69de29..0000000
diff --git a/src/AbstractQueryDataReader.php b/src/AbstractQueryDataReader.php
index 8d6b3b3..80f0295 100644
--- a/src/AbstractQueryDataReader.php
+++ b/src/AbstractQueryDataReader.php
@@ -191,6 +191,7 @@ protected function applyHaving(QueryInterface $query): QueryInterface
 
     /**
      * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
      */
     public function withOffset(int $offset): static
     {
@@ -203,6 +204,7 @@ public function withOffset(int $offset): static
 
     /**
      * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
      */
     public function withLimit(int $limit): static
     {
@@ -217,6 +219,10 @@ public function withLimit(int $limit): static
         return $new;
     }
 
+    /**
+     * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
+     */
     public function withCountParam(?string $countParam): static
     {
         if ($this->countParam === $countParam) {
@@ -232,6 +238,7 @@ public function withCountParam(?string $countParam): static
 
     /**
      * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
      */
     public function withSort(?Sort $sort): static
     {
@@ -244,6 +251,7 @@ public function withSort(?Sort $sort): static
 
     /**
      * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
      */
     public function withFilter(FilterInterface $filter): static
     {
@@ -254,6 +262,10 @@ public function withFilter(FilterInterface $filter): static
         return $new;
     }
 
+    /**
+     * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
+     */
     public function withHaving(?FilterInterface $having): static
     {
         $new = clone $this;
@@ -263,6 +275,10 @@ public function withHaving(?FilterInterface $having): static
         return $new;
     }
 
+    /**
+     * @psalm-mutation-free
+     * @psalm-return static<TKey, TValue>
+     */
     public function withBatchSize(?int $batchSize): static
     {
         if ($batchSize !== null && $batchSize < 1) {
@@ -278,7 +294,9 @@ public function withBatchSize(?int $batchSize): static
     /**
      * @param FilterHandlerInterface ...$filterHandlers
      * @return $this
-     * @psalm-suppress ArgumentTypeCoercion
+     *
+     * @psalm-suppress ArgumentTypeCoercion    *
+     * @psalm-return static<TKey, TValue>
      */
     public function withFilterHandlers(FilterHandlerInterface ...$filterHandlers): static
     {
diff --git a/src/QueryDataReader.php b/src/QueryDataReader.php
index 661451c..68db498 100644
--- a/src/QueryDataReader.php
+++ b/src/QueryDataReader.php
@@ -5,7 +5,12 @@
 namespace Yiisoft\Data\Db;
 
 /**
- * Base class for QueryDataReaderInterface
+ * Base class for `QueryDataReaderInterface`
+ *
+ * @template TKey as array-key
+ * @template TValue as array|object
+ *
+ * @extends AbstractQueryDataReader<TKey, TValue>
  */
 final class QueryDataReader extends AbstractQueryDataReader
 {