Close Excel extractor when limit is reached (#1659)

* Close Excel extractor when the limit is reached

* Prevent memory leak in ExcelExtractor by not copying each row in loop

Author: stloyd

src/adapter/etl-adapter-excel/src/Flow/ETL/Adapter/Excel/ExcelExtractor.php

@@ -4,20 +4,20 @@
 
 namespace Flow\ETL\Adapter\Excel;
 
-use function Flow\ETL\DSL\array_to_rows;
+use function Flow\ETL\DSL\{array_to_rows};
 use Flow\ETL\{Adapter\Excel\Sheet\SheetNameAssertion,
     Adapter\Excel\Sheet\SheetsManager,
     Exception\InvalidArgumentException,
     Extractor,
-    FlowContext,
-    Loader\Closure};
+    FlowContext
+};
 use Flow\ETL\Extractor\{FileExtractor, Limitable, LimitableExtractor, PathFiltering, Signal};
 use Flow\Filesystem\{Path, SourceStream};
 use OpenSpout\Common\Entity\{Cell, Row};
 use OpenSpout\Reader\ODS\Reader as OdsReader;
 use OpenSpout\Reader\XLSX\Reader as XlsxReader;
 
-final class ExcelExtractor implements Closure, Extractor, FileExtractor, LimitableExtractor
+final class ExcelExtractor implements Extractor, FileExtractor, LimitableExtractor
 {
     use Limitable;
     use PathFiltering;
@@ -43,11 +43,6 @@ public function __construct(private readonly Path $path)
         $this->resetLimit();
     }
 
-    public function closure(FlowContext $context) : void
-    {
-        $this->reader()->close();
-    }
-
     public function extract(FlowContext $context) : \Generator
     {
         $headers = [];
@@ -61,13 +56,17 @@ public function extract(FlowContext $context) : \Generator
 
         foreach ($context->streams()->list($this->path, $this->filter()) as $stream) {
             foreach ($this->extractRows($stream, $headers, $offset) as $row) {
-                $signal = yield array_to_rows($row, $context->entryFactory());
+                $signal = yield array_to_rows($row, $context->entryFactory(), $stream->path()->partitions());
                 $this->incrementReturnedRows();
 
                 if ($signal === Signal::STOP || $this->reachedLimit()) {
+                    $stream->close();
+
                     return;
                 }
             }
+
+            $stream->close();
         }
     }
 
@@ -136,7 +135,7 @@ private function createRowsFromCells(Row $row, int $previousRowDataCount = 0) :
         return $rowData;
     }
 
-    private function extractRows(SourceStream $stream, array $headers, int $offset) : array
+    private function extractRows(SourceStream $stream, array $headers, int $offset) : \Generator
     {
         try {
             $this->reader()->open($stream->path()->path());
@@ -148,9 +147,9 @@ private function extractRows(SourceStream $stream, array $headers, int $offset)
 
             $sheet = $this->sheetName ? $manager->get($this->sheetName) : $manager->first();
 
-            foreach ($sheet->getRowIterator() as $rowIndex => $row) {
+            foreach ($sheet->getRowIterator() as $rowIndex => $sheetRow) {
                 if (1 === $rowIndex && $this->withHeader) {
-                    $headers = $this->createRowsFromCells($row);
+                    $headers = $this->createRowsFromCells($sheetRow);
 
                     continue;
                 }
@@ -161,17 +160,17 @@ private function extractRows(SourceStream $stream, array $headers, int $offset)
                 }
 
                 // ODS format reader skips empty cells when reading rows
-                $rowData = $this->createRowsFromCells($row, $previousRowDataCount);
-                $previousRowDataCount = \count($rowData);
+                $row = $this->createRowsFromCells($sheetRow, $previousRowDataCount);
+                $previousRowDataCount = \count($row);
 
                 if ($this->withHeader) {
-                    $rowData = \array_combine($headers, $rowData);
+                    yield \array_combine($headers, $row);
+                } else {
+                    yield $row;
                 }
-
-                $rows[] = $rowData;
             }
 
-            return $rows;
+            $this->reader()->close();
         } catch (\Throwable $e) {
             throw new InvalidArgumentException('Failed to open file: ' . $e->getMessage(), previous: $e);
         }

src/adapter/etl-adapter-excel/tests/Flow/ETL/Adapter/Excel/Tests/Fixtures/partitioned/group=1/file_01.xlsx

@@ -8,8 +8,9 @@
 use function Flow\ETL\DSL\{config, df, flow_context};
 use Flow\ETL\Adapter\Excel\ExcelReader;
 use Flow\ETL\Exception\InvalidArgumentException;
-use Flow\ETL\Row;
+use Flow\ETL\{Extractor\Signal, Row, Rows};
 use Flow\ETL\Tests\FlowTestCase;
+use Flow\Filesystem\{Partition, Path};
 use PHPUnit\Framework\Attributes\DataProvider;
 
 final class ExcelExtractorTest extends FlowTestCase
@@ -199,4 +200,33 @@ public function test_extract_with_wrongly_selected_reader() : void
 
         iterator_to_array($extractor->extract(flow_context(config())));
     }
+
+    public function test_loading_data_from_all_partitions() : void
+    {
+        df()
+            ->read(from_excel(__DIR__ . '/../Fixtures/partitioned/group=*/*.xlsx'))
+            ->run(function (Rows $rows) : void {
+                $this->assertSame(
+                    ['group'],
+                    \array_map(
+                        fn (Partition $p) => $p->name,
+                        $rows->partitions()->toArray()
+                    )
+                );
+            });
+    }
+
+    public function test_signal_stop() : void
+    {
+        $generator = from_excel(Path::realpath(__DIR__ . '/../Fixtures/fixture.xlsx'))
+            ->extract(flow_context(config()));
+
+        self::assertTrue($generator->valid());
+        $generator->next();
+        self::assertTrue($generator->valid());
+        $generator->next();
+        self::assertTrue($generator->valid());
+        $generator->send(Signal::STOP);
+        self::assertFalse($generator->valid());
+    }
 }