Keep all rows for native vstack merges with list inputs (#1320)

src/python/txtai/workflow/task/base.py

@@ -368,7 +368,14 @@ class Task:
             return outputs
 
         if self.merge == "vstack":
-            return self.vstack(outputs)
+            results = self.vstack(outputs)
+
+            if isinstance(results, np.ndarray) or torch.is_tensor(results):
+                size = len(self.action)
+                results = [OneToMany(list(results[x : x + size])) for x in range(0, len(results), size)]
+
+            return results
+
         if self.merge == "concat":
             return self.concat(outputs)
 

test/python/testworkflow.py

@@ -391,6 +391,17 @@ class TestWorkflow(unittest.TestCase):
         results = list(workflow(np.array([2, 4, 6])))
         self.assertTrue(np.array_equal(np.array(results), np.array([[4, 16, 36], [8, 64, 216]])))
 
+    def testNumpyListWorkflow(self):
+        """
+        Test a vstack merge of numpy and torch action outputs with list inputs
+        """
+
+        # pylint: disable=E1101
+        for function in [lambda x, y: np.power(np.array(x), y), lambda x, y: torch.pow(torch.tensor(x), y)]:
+            task = Task([lambda x, f=function: f(x, 2), lambda x, f=function: f(x, 3)], merge="vstack")
+            results = list(Workflow([task])([2, 4]))
+            self.assertEqual([int(x) for x in results], [4, 8, 16, 64])
+
     def testRetrieveWorkflow(self):
         """
         Test a retrieve task