From 98c965cc3bd65c69e39298b65555848978fb75a7 Mon Sep 17 00:00:00 2001 From: Edson <113824357+Zhuoxi2000@users.noreply.github.com> Date: Fri, 2 Oct 2026 14:30:31 -0400 Subject: [PATCH] Keep all rows for native vstack merges with list inputs (#1320) --- src/python/txtai/workflow/task/base.py | 9 ++++++++- test/python/testworkflow.py | 11 +++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/src/python/txtai/workflow/task/base.py b/src/python/txtai/workflow/task/base.py index b00d226b..3731db66 100644 --- a/src/python/txtai/workflow/task/base.py +++ b/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) diff --git a/test/python/testworkflow.py b/test/python/testworkflow.py index b240371c..baf38a83 100644 --- a/test/python/testworkflow.py +++ b/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