Concatenate method for DataPipeline class

fairseq2n/python/src/fairseq2n/bindings/data/data_pipeline.cc

@@ -358,8 +358,35 @@ def_data_pipeline(py::module_ &data_module)
                 return data_pipeline::count(start, std::move(key));
             },
             py::arg("start") = 0,
-            py::arg("key") = std::nullopt);
+            py::arg("key") = std::nullopt)
+        .def_static(
+            "cat",
+            [](
+                std::vector<std::reference_wrapper<data_pipeline>> &refs1,
+                std::vector<std::reference_wrapper<data_pipeline>> &refs2
+            )
+            {
+                std::vector<data_pipeline> pipeline1{};
+                std::vector<data_pipeline> pipeline2{};
+
+                pipeline1.reserve(refs1.size());
+                pipeline2.reserve(refs2.size());
+
+                std::transform(
+                    refs1.begin(), refs1.end(), std::back_inserter(pipeline1), [](auto &r) {
+                        return std::move(r.get());
+                    });
 
+                std::transform(
+                    refs2.begin(), refs2.end(), std::back_inserter(pipeline2), [](auto &r) {
+                        return std::move(r.get());
+                    });
+
+                return data_pipeline::cat(std::move(pipeline1), std::move(pipeline2));
+            },
+            py::arg("pipeline1"),
+            py::arg("pipeline2"));
+
     // DataPipelineIterator
     py::class_<data_pipeline_iterator>(m, "_DataPipelineIterator")
         .def(

fairseq2n/src/fairseq2n/CMakeLists.txt

@@ -17,6 +17,7 @@ target_sources(fairseq2n
         data/bucket_by_length_data_source.cc
         data/bucket_data_source.cc
         data/byte_stream.cc
+        data/cat_data_source.cc
         data/collater.cc
         data/constant_data_source.cc
         data/count_data_source.cc

fairseq2n/src/fairseq2n/data/cat_data_source.cc

@@ -0,0 +1,76 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include "fairseq2n/data/cat_data_source.h"
+#include <vector>
+
+namespace fairseq2n::detail {
+
+cat_data_source::cat_data_source(
+    std::vector<data_pipeline> &&pipeline1,
+    std::vector<data_pipeline> &&pipeline2)
+    : pipeline1_{std::move(pipeline1)}
+    , pipeline2_{std::move(pipeline2)}
+{}
+
+std::optional<data>
+cat_data_source::next()
+{
+    std::optional<data> d;
+    for (auto &p : pipeline1_) {
+        d = p.next();
+        if (d) {
+            return d;
+        }
+    }
+    for (auto &p : pipeline2_) {
+        d = p.next();
+        if (d) {
+            return d;
+        }
+    }
+    return {};
+}
+
+void cat_data_source::reset()
+{
+    for (auto &pipeline : pipeline1_)
+        pipeline.reset();
+    for (auto &pipeline : pipeline2_)
+        pipeline.reset();
+}  
+
+void cat_data_source::record_position(tape &t) const
+{
+    for (auto &pipeline : pipeline1_)
+        pipeline.record_position(t);
+    for (auto &pipeline : pipeline2_)
+        pipeline.record_position(t);
+}
+
+void cat_data_source::reload_position(tape &t)
+{
+    for (auto &pipeline : pipeline1_)
+        pipeline.reload_position(t);
+    for (auto &pipeline : pipeline2_)
+        pipeline.reload_position(t);
+}
+
+std::vector<data_pipeline> cat_data_source::concatenate(
+    std::vector<data_pipeline> &&pipeline1,
+    std::vector<data_pipeline> &&pipeline2)
+{
+    std::vector<data_pipeline> result;
+    result.reserve(pipeline1.size() + pipeline2.size());
+    for (auto &&pipeline : pipeline1) {
+        result.push_back(std::move(pipeline));
+    }
+    for (auto &&pipeline : pipeline2) {
+        result.push_back(std::move(pipeline));
+    }
+    return result;
+}
+} // namespace fairseq2n::detail

fairseq2n/src/fairseq2n/data/cat_data_source.h

@@ -0,0 +1,47 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+// All rights reserved.
+//
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "fairseq2n/data/data_pipeline.h"
+#include "fairseq2n/data/data_source.h"
+
+namespace fairseq2n::detail {
+
+class cat_data_source final : public data_source {
+public:
+    explicit
+    cat_data_source(
+        std::vector<data_pipeline> &&pipeline1,
+        std::vector<data_pipeline> &&pipeline2);
+
+    std::optional<data>
+    next() override;
+
+    void
+    reset() override;
+
+    void
+    record_position(tape &t) const override;
+
+    void
+    reload_position(tape &t) override;
+
+    static std::vector<data_pipeline>
+    concatenate(
+        std::vector<data_pipeline> &&pipeline1,
+        std::vector<data_pipeline> &&pipeline2);
+
+private:
+    std::vector<data_pipeline> pipeline1_;
+    std::vector<data_pipeline> pipeline2_;
+};
+
+}  // namespace fairseq2n::detail

fairseq2n/src/fairseq2n/data/data_pipeline.cc

@@ -11,25 +11,27 @@
 #include <system_error>
 #include <utility>
 
+#include "data_pipeline.h"
 #include "fairseq2n/data/bucket_by_length_data_source.h"
 #include "fairseq2n/data/bucket_data_source.h"
+#include "fairseq2n/data/cat_data_source.h"
 #include "fairseq2n/data/constant_data_source.h"
 #include "fairseq2n/data/count_data_source.h"
+#include "fairseq2n/data/detail/file_system.h"
 #include "fairseq2n/data/filter_data_source.h"
 #include "fairseq2n/data/list_data_source.h"
 #include "fairseq2n/data/map_data_source.h"
 #include "fairseq2n/data/prefetch_data_source.h"
-#include "fairseq2n/data/take_data_source.h"
 #include "fairseq2n/data/round_robin_data_source.h"
 #include "fairseq2n/data/sample_data_source.h"
 #include "fairseq2n/data/shard_data_source.h"
 #include "fairseq2n/data/shuffle_data_source.h"
 #include "fairseq2n/data/skip_data_source.h"
+#include "fairseq2n/data/take_data_source.h"
 #include "fairseq2n/data/tape.h"
 #include "fairseq2n/data/yield_from_data_source.h"
 #include "fairseq2n/data/zip_data_source.h"
 #include "fairseq2n/data/zip_file_data_source.h"
-#include "fairseq2n/data/detail/file_system.h"
 #include "fairseq2n/detail/exception.h"
 
 using namespace fairseq2n::detail;
@@ -283,6 +285,42 @@ data_pipeline::count(std::int64_t start, std::optional<std::string> key)
     return data_pipeline_builder{std::move(factory)};
 }
 
+data_pipeline_builder
+data_pipeline::cat(
+    std::vector<data_pipeline> pipeline1,
+    std::vector<data_pipeline> pipeline2)
+{
+    bool is_broken = false;
+
+    is_broken = std::any_of(
+        pipeline1.begin(), pipeline1.end(), [](const data_pipeline &pipeline)
+        {
+            return pipeline.is_broken();
+        });
+
+    is_broken = is_broken || std::any_of(
+        pipeline2.begin(), pipeline2.end(), [](const data_pipeline &pipeline)
+        {
+            return pipeline.is_broken();
+        });
+
+    if (is_broken)
+        throw_<std::invalid_argument>(
+            "At least one of the specified data pipelines is broken and cannot be used in cat.");
+
+    auto tmp = std::make_shared<std::vector<data_pipeline>>(std::move(pipeline1));
+    auto tmp2 = std::make_shared<std::vector<data_pipeline>>(std::move(pipeline2));
+
+    auto factory = [
+        tmp,
+        tmp2]() mutable
+    {
+        return std::make_unique<cat_data_source>(std::move(*tmp), std::move(*tmp2));
+    };
+
+    return data_pipeline_builder{std::move(factory)};
+}
+
 data_pipeline_builder
 data_pipeline_builder::bucket(std::size_t bucket_size, bool drop_remainder) &&
 {

fairseq2n/src/fairseq2n/data/data_pipeline.h

@@ -97,6 +97,10 @@ class FAIRSEQ2_API data_pipeline {
     static data_pipeline_builder
     count(std::int64_t start = 0, std::optional<std::string> key = {});
 
+    static data_pipeline_builder
+    cat(std::vector<data_pipeline> pipeline1,
+        std::vector<data_pipeline> pipeline2);
+
 private:
     data_source_factory factory_{};
     std::unique_ptr<data_source> source_{};

src/fairseq2/data/data_pipeline.py

@@ -1,4 +1,4 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -134,6 +134,14 @@
         def count(start: int = 0, key: Optional[str] = None) -> "DataPipelineBuilder":
             ...
 
+        @staticmethod
+        def cat(pipeline1: Sequence["DataPipeline"], pipeline2: Sequence["DataPipeline"]) -> "DataPipelineBuilder":
+            """Concatenate examples from ``pipelines``.
+
+            :param pipelines:
+                The data pipelines to concatenate.
+            """
+
     class DataPipelineBuilder:
         """API to create DataPipeline"""
 

tests/unit/data/data_pipeline/test_cat.py

@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import pytest
+
+from fairseq2.data import DataPipeline, DataPipelineError, read_sequence
+from fairseq2.data.text import read_text
+
+class TestCatOp:
+    def test_op_works(self) -> None:
+        pipeline1 = read_sequence([1, 2, 3, 4]).and_return()
+        pipeline2 = read_sequence([5, 6, 7, 8]).and_return()
+
+        pipeline = DataPipeline.cat([pipeline1], [pipeline2]).and_return()
+
+        for _ in range(2):
+            assert list(pipeline) == [1, 2, 3, 4, 5, 6, 7, 8]
+            pipeline.reset()
+
+    def test_op_works_when_pipelines_are_empty(self) -> None:
+        pipeline1 = read_sequence([]).and_return()
+        pipeline2 = read_sequence([]).and_return()
+
+        pipeline = DataPipeline.cat([pipeline1], [pipeline2]).and_return()
+
+        for _ in range(2):
+            assert list(pipeline) == []
+            pipeline.reset()
+
+    def test_op_works_when_pipelines_have_different_lengths(self) -> None:
+        pipeline1 = read_sequence([1, 2, 3]).and_return()
+        pipeline2 = read_sequence([4, 5]).and_return()
+
+        pipeline = DataPipeline.cat([pipeline1], [pipeline2]).and_return()
+
+        for _ in range(2):
+            assert list(pipeline) == [1, 2, 3, 4, 5]
+            pipeline.reset()
+
+    def test_op_raises_error_when_one_of_the_pipelines_is_broken(self) -> None:
+        # Force a non-recoverable error.
+        pipeline1 = read_text(pathname=" &^#").and_return()
+        pipeline2 = read_text(pathname=" &^#").and_return()
+
+        # Break the first pipeline.
+        try:
+            next(iter(pipeline1))
+        except DataPipelineError:
+            assert pipeline1.is_broken
+
+        with pytest.raises(
+            ValueError,
+            match=r"^At least one of the specified data pipelines is broken and cannot be used in cat\.$",
+        ):
+            DataPipeline.cat([pipeline1], [pipeline2]).and_return()
+
+    def test_op_saves_and_restores_its_state(self) -> None:
+        pipeline1 = read_sequence([1, 2, 3, 4]).and_return()
+        pipeline2 = read_sequence([5, 6, 7, 8]).and_return()
+
+        pipeline = DataPipeline.cat([pipeline1], [pipeline2]).and_return()
+
+        d = None
+
+        it = iter(pipeline)
+
+        # Move the the second example.
+        for _ in range(2):
+            d = next(it)
+
+        assert d == 2
+
+        state_dict = pipeline.state_dict()
+
+        # Read one more example before we roll back.
+        d = next(it)
+
+        assert d == 3
+
+        # Expected to roll back to the second example.
+        pipeline.load_state_dict(state_dict)
+
+        # Move to EOD.
+        for _ in range(6):
+            d = next(it)
+
+        assert d == 8
+
+        state_dict = pipeline.state_dict()
+
+        pipeline.reset()
+
+        # Expected to be EOD.
+        pipeline.load_state_dict(state_dict)
+
+        with pytest.raises(StopIteration):
+            next(iter(pipeline))