Improve select usage with Statement

src/Statement.php

@@ -51,6 +51,21 @@ public static function create(callable $where = null, int $offset = 0, int $limi
         return $stmt->offset($offset)->limit($limit);
     }
 
+    /**
+     * Sets the Iterator element columns.
+     */
+    public function select(string|int ...$columns): self
+    {
+        if ($columns === $this->select) {
+            return $this;
+        }
+
+        $clone = clone $this;
+        $clone->select = $columns;
+
+        return $clone;
+    }
+
     /**
      * Sets the Iterator filter method.
      */
@@ -138,7 +153,7 @@ public function process(TabularDataReader $tabular_data, array $header = []): Ta
         /** @var Iterator<array-key, array<array-key, string|null>> $iterator */
         $iterator = new LimitIterator($iterator, $this->offset, $this->limit);
 
-        return $this->applySelect(new ResultSet($iterator, $header));
+        return $this->applySelect($iterator, $header);
     }
 
     /**
@@ -175,50 +190,36 @@ protected function buildOrderBy(Iterator $iterator): Iterator
         return $it;
     }
 
-    public function select(string|int ...$columns): self
-    {
-        if ($columns === $this->select) {
-            return $this;
-        }
-
-        $clone = clone $this;
-        $clone->select = $columns;
-
-        return $clone;
-    }
-
-    protected function applySelect(TabularDataReader $tabularDataReader): TabularDataReader
+    protected function applySelect(Iterator $records, array $recordsHeader): TabularDataReader
     {
         if ([] === $this->select) {
-            return $tabularDataReader;
+            return new ResultSet($records, $recordsHeader);
         }
 
-        $header = [];
-        $documentHeader = $tabularDataReader->getHeader();
-        $hasNoHeader = [] === $documentHeader;
-        foreach ($this->select as $field) {
+        $hasHeader = [] !== $recordsHeader;
+        /** @var array<string> $header */
+        $header = array_reduce($this->select, function (array $header, string|int $field) use ($recordsHeader, $hasHeader): array {
             if (is_string($field)) {
-                if ($hasNoHeader) {
-                    throw new InvalidArgument(__METHOD__.' can only use named column if the tabular data has a non-empty header.');
-                }
-
-                $index = array_search($field, $documentHeader, true);
+                $index = array_search($field, $recordsHeader, true);
                 if (false === $index) {
                     throw InvalidArgument::dueToInvalidColumnIndex($field, 'offset', __METHOD__);
                 }
 
                 $header[$index] = $field;
-                continue;
+
+                return $header;
             }
 
-            if (!$hasNoHeader && !array_key_exists($field, $documentHeader)) {
+            if ($hasHeader && !array_key_exists($field, $recordsHeader)) {
                 throw InvalidArgument::dueToInvalidColumnIndex($field, 'offset', __METHOD__);
             }
 
-            $header[$field] = $documentHeader[$field] ?? $field;
-        }
+            $header[$field] = $recordsHeader[$field] ?? $field;
 
-        $records = new MapIterator($tabularDataReader, function (array $record) use ($header): array {
+            return $header;
+        }, []);
+
+        $records = new MapIterator($records, function (array $record) use ($header): array {
             $assocRecord = [];
             $row = array_values($record);
             foreach ($header as $offset => $headerName) {
@@ -228,9 +229,6 @@ protected function applySelect(TabularDataReader $tabularDataReader): TabularDat
             return $assocRecord;
         });
 
-        /** @var array<string>. $finalHeader */
-        $finalHeader = $hasNoHeader ? [] : $header;
-
-        return new ResultSet($records, $finalHeader);
+        return new ResultSet($records, $hasHeader ? $header : []);
     }
 }