From 406b35cb0db82c1a879e65c9d4deeaa1b7514f64 Mon Sep 17 00:00:00 2001
From: Norbert Orzechowicz <1921950+norberttech@users.noreply.github.com>
Date: Sun, 28 Jan 2024 15:41:01 +0100
Subject: [PATCH] Added to_branch loader (#937)

* Added to_branch loader

* Added Closure interface to BranchingLoader
---
 src/core/etl/src/Flow/ETL/DSL/functions.php   |  5 ++
 .../src/Flow/ETL/Loader/BranchingLoader.php   | 42 +++++++++++++
 .../Integration/DataFrame/BranchingTest.php   | 60 +++++++++++++++++++
 3 files changed, 107 insertions(+)
 create mode 100644 src/core/etl/src/Flow/ETL/Loader/BranchingLoader.php
 create mode 100644 src/core/etl/tests/Flow/ETL/Tests/Integration/DataFrame/BranchingTest.php

diff --git a/src/core/etl/src/Flow/ETL/DSL/functions.php b/src/core/etl/src/Flow/ETL/DSL/functions.php
index ea37f6324..f3b8a9cc4 100644
--- a/src/core/etl/src/Flow/ETL/DSL/functions.php
+++ b/src/core/etl/src/Flow/ETL/DSL/functions.php
@@ -248,6 +248,11 @@ function to_transformation(Transformer $transformer, Loader $loader) : Transform
     return new TransformerLoader($transformer, $loader);
 }
 
+function to_branch(ScalarFunction $condition, Loader $loader) : Loader
+{
+    return new Loader\BranchingLoader($condition, $loader);
+}
+
 /**
  * @param array<mixed> $data
  */
diff --git a/src/core/etl/src/Flow/ETL/Loader/BranchingLoader.php b/src/core/etl/src/Flow/ETL/Loader/BranchingLoader.php
new file mode 100644
index 000000000..0e1166dac
--- /dev/null
+++ b/src/core/etl/src/Flow/ETL/Loader/BranchingLoader.php
@@ -0,0 +1,42 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Flow\ETL\Loader;
+
+use Flow\ETL\FlowContext;
+use Flow\ETL\Function\ScalarFunction;
+use Flow\ETL\Loader;
+use Flow\ETL\Rows;
+use Flow\ETL\Transformer\ScalarFunctionFilterTransformer;
+
+final class BranchingLoader implements Closure, Loader, OverridingLoader
+{
+    public function __construct(
+        private readonly ScalarFunction $condition,
+        private readonly Loader $loader
+    ) {
+    }
+
+    public function closure(FlowContext $context) : void
+    {
+        if ($this->loader instanceof Closure) {
+            $this->loader->closure($context);
+        }
+    }
+
+    public function load(Rows $rows, FlowContext $context) : void
+    {
+        $this->loader->load(
+            (new ScalarFunctionFilterTransformer($this->condition))->transform($rows, $context),
+            $context
+        );
+    }
+
+    public function loaders() : array
+    {
+        return [
+            $this->loader,
+        ];
+    }
+}
diff --git a/src/core/etl/tests/Flow/ETL/Tests/Integration/DataFrame/BranchingTest.php b/src/core/etl/tests/Flow/ETL/Tests/Integration/DataFrame/BranchingTest.php
new file mode 100644
index 000000000..e406fe03c
--- /dev/null
+++ b/src/core/etl/tests/Flow/ETL/Tests/Integration/DataFrame/BranchingTest.php
@@ -0,0 +1,60 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Flow\ETL\Tests\Integration\DataFrame;
+
+use function Flow\ETL\DSL\df;
+use function Flow\ETL\DSL\from_array;
+use function Flow\ETL\DSL\lit;
+use function Flow\ETL\DSL\ref;
+use function Flow\ETL\DSL\to_branch;
+use function Flow\ETL\DSL\to_memory;
+use Flow\ETL\Memory\ArrayMemory;
+use Flow\ETL\Tests\Integration\IntegrationTestCase;
+
+final class BranchingTest extends IntegrationTestCase
+{
+    public function test_branching() : void
+    {
+        df()
+            ->read(from_array([
+                ['id' => 1, 'group' => 'A'],
+                ['id' => 2, 'group' => 'B'],
+                ['id' => 3, 'group' => 'A'],
+                ['id' => 4, 'group' => 'B'],
+                ['id' => 5, 'group' => 'A'],
+                ['id' => 6, 'group' => 'C'],
+            ]))
+            ->write(
+                to_branch(
+                    ref('group')->equals(lit('A')),
+                    to_memory($memoryA = new ArrayMemory()),
+                )
+            )
+            ->write(
+                to_branch(
+                    ref('group')->isIn(lit(['B', 'C'])),
+                    to_memory($memoryBC = new ArrayMemory()),
+                )
+            )
+            ->run();
+
+        $this->assertSame(
+            [
+                ['id' => 1, 'group' => 'A'],
+                ['id' => 3, 'group' => 'A'],
+                ['id' => 5, 'group' => 'A'],
+            ],
+            $memoryA->dump(),
+        );
+        $this->assertSame(
+            [
+                ['id' => 2, 'group' => 'B'],
+                ['id' => 4, 'group' => 'B'],
+                ['id' => 6, 'group' => 'C'],
+            ],
+            $memoryBC->dump(),
+        );
+    }
+}