#include "client.h"

#include <saas/library/rtyt/lib/operation/factory.h>
#include <saas/library/rtyt/lib/io/ut_row.pb.h>
#include <robot/library/yt/static/command.h>

#include <library/cpp/testing/unittest/registar.h>
#include <mapreduce/yt/interface/operation.h>
#include <mapreduce/yt/interface/io.h>
#include <util/system/fs.h>

namespace {
    class TRenamingMapper
        : public NYT::IMapper<NYT::TTableReader<NYT::TNode>, NYT::TTableWriter<NYT::TNode>>
    {
    public:
        TRenamingMapper() = default;

        void Do(NYT::TTableReader<NYT::TNode>* input, NYT::TTableWriter<NYT::TNode>* output) override {
            for (; input->IsValid(); input->Next()) {
                NYT::TNode row = input->GetRow();
                row["b"] = row["a"];
                output->AddRow(row);
            }
        }
    };

    REGISTER_RTYT_MAPPER(TRenamingMapper);

    class TProtoMultiplyingMapper
        : public NYT::IMapper<NYT::TTableReader<NRTYT::NTesting::TRow>, NYT::TTableWriter<NRTYT::NTesting::TRow>>
    {
    public:
        TProtoMultiplyingMapper() = default;

        void Do(TReader* input, TWriter* output) override {
            NRTYT::NTesting::TRow row;
            for (; input->IsValid(); input->Next()) {
                row = input->GetRow();
                i64 result = row.GetInt32() * row.GetFixed64();
                row.SetFixed64(result);
                output->AddRow(row);
            }
        }
    };

    REGISTER_RTYT_MAPPER(TProtoMultiplyingMapper);

    class TSimpleProtoReducer
        : public NYT::IReducer<NYT::TTableReader<NRTYT::NTesting::TRow>, NYT::TTableWriter<NRTYT::NTesting::TRow>>
    {
        int Modifier = 1;
    public:
        TSimpleProtoReducer() = default;
        TSimpleProtoReducer(int mod) : Modifier(mod) {
        }

        void Do(TReader* input, TWriter* output) override {
            int result = 0;
            int key = 0;
            for (; input->IsValid(); input->Next()) {
                auto row = input->GetRow();
                result += Modifier * row.GetInt32();
                key = row.GetFixed64();
            }
            NRTYT::NTesting::TRow row;
            row.SetInt32(result);
            row.SetFixed64(key);
            output->AddRow(row);
        }

        Y_SAVELOAD_JOB(Modifier);
    };
    REGISTER_RTYT_REDUCER(TSimpleProtoReducer);

    TFsPath path = "./rtyt_test";

    void Cleanup() {
        TVector<TFsPath> children;
        if (path.Exists()) {
            path.List(children);
            for (const auto& child : children) {
                child.ForceDelete();
            }
        }
    }
}

Y_UNIT_TEST_SUITE(RTYT_SIMPLE_MAPPER) {
    Y_UNIT_TEST(SIMPLE_CREATING_TEST) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NYT::TNode;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions().Attributes(NYT::TNode()("test_attr", 1)));
        client.Create("//test1/test2", NYT::NT_TABLE, NYT::TCreateOptions().Attributes(NYT::TNode()("test_attr", 2)));
        client.Create("//test1/output", NYT::NT_TABLE);
        auto writer = client.CreateTableWriter(TString("//test1/test2"));
        writer->AddRow(NYT::TNode()("a", 1)("b", "b"), 0);
        writer->AddRow(NYT::TNode()("a", 2)("b", "b"), 0);
        NYT::TMapOperationSpec spec;
        spec.AddInput<NYT::TNode>(NYT::TRichYPath("//test1/test2")).AddOutput<NYT::TNode>(NYT::TRichYPath("//test1/output"));
        client.Map(spec, new TRenamingMapper(), NYT::TOperationOptions());
        auto reader = client.CreateTableReader<NYT::TNode>(NYT::TRichYPath("//test1/output"));
        auto row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 1);
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("b"), 1);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("b"), 2);
        UNIT_ASSERT_VALUES_EQUAL(client.Get("//test1").ChildAsInt64("test_attr"), 1);
    }

    Y_UNIT_TEST(PROTO_MAPPER_TEST) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        NRTYT::TClientBase client(".");
        TIntrusivePtr<NRTYT::TClientBase> client2 = MakeIntrusive<NRTYT::TClientBase>();
        NRTYT::TClient c(*client2);
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        client.Create("//test1/output", NYT::NT_TABLE);
        TRow sample;
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetString("abcabc");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(2);
            second.SetFixed64(2);
            writer->AddRow(second, 0);
        }

        NYT::TMapOperationSpec spec;
        spec.AddInput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/test2")).AddOutput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
        client.Map(spec, new TProtoMultiplyingMapper(), NYT::TOperationOptions());

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 4);
    }
}


Y_UNIT_TEST_SUITE(RTYT_TABLE_IO) {
    Y_UNIT_TEST(CREATE_BY_APPEND) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        TRow sample;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        {
            auto writer = client.CreateTableWriter(NYT::TRichYPath("//test1/test2").Append(true), *sample.GetDescriptor());
            TRow first;
            first.SetInt32(1);
            writer->AddRow(first, 0);

            TRow second;
            second.SetInt32(2);
            writer->AddRow(second, 0);
        }
        {
            auto reader = client.CreateTableReader<TRow>(NYT::TRichYPath("//test1/test2"));
            auto row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 1);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        }
    }
    Y_UNIT_TEST(PROTO_APPEND_MODE) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        TRow sample;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetInt32(1);
            writer->AddRow(first, 0);

            TRow second;
            second.SetInt32(2);
            writer->AddRow(second, 0);
        }
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetInt32(3);
            writer->AddRow(first, 0);

            TRow second;
            second.SetInt32(4);
            writer->AddRow(second, 0);
        }
        {
            auto reader = client.CreateTableReader<TRow>(NYT::TRichYPath("//test1/test2"));
            auto row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 3);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 4);
        }
        {
            auto writer = client.CreateTableWriter(NYT::TRichYPath(TString("//test1/test2")).Append(true), *sample.GetDescriptor());
            TRow first;
            first.SetInt32(5);
            writer->AddRow(first, 0);

            TRow second;
            second.SetInt32(6);
            writer->AddRow(second, 0);
        }
        {
            
            auto reader = client.CreateTableReader<TRow>(NYT::TRichYPath("//test1/test2"));
            auto row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 3);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 4);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 5);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 6);
        }
    }
    Y_UNIT_TEST(TEST_APPEND_MODE) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NYT::TNode;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        {
            auto writer = client.CreateTableWriter(TString("//test1/test2"));
            writer->AddRow(NYT::TNode()("a", 1)("b", "b"), 0);
            writer->AddRow(NYT::TNode()("a", 2)("b", "b"), 0);
        }
        {
            auto writer = client.CreateTableWriter(TString("//test1/test2"));
            writer->AddRow(NYT::TNode()("a", 3)("b", "b"), 0);
            writer->AddRow(NYT::TNode()("a", 4)("b", "b"), 0);
        }
        {
            auto reader = client.CreateTableReader<NYT::TNode>(NYT::TRichYPath("//test1/test2"));
            auto row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 3);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 4);
        }
        {
            auto writer = client.CreateTableWriter(NYT::TRichYPath(TString("//test1/test2")).Append(true));
            writer->AddRow(NYT::TNode()("a", 5)("b", "b"), 0);
            writer->AddRow(NYT::TNode()("a", 6)("b", "b"), 0);
        }
        {
            auto reader = client.CreateTableReader<NYT::TNode>(NYT::TRichYPath("//test1/test2"));
            auto row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 3);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 4);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 5);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.ChildAsInt64("a"), 6);
        }
    }
}

Y_UNIT_TEST_SUITE(RTYT_SORT) {
    Y_UNIT_TEST(PROTO_SORT_TEST) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        client.Create("//test1/output", NYT::NT_TABLE);
        TRow sample;
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetString("abcabc");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(2);
            second.SetFixed64(2);
            writer->AddRow(second, 0);
        }

        NYT::TSortOperationSpec spec;
        spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
        spec.SortBy({"Fixed64", "Int32"});
        client.Sort(spec);

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 2);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);

        spec = {};
        spec.AddInput(NYT::TRichYPath("//test1/output")).Output(NYT::TRichYPath("//test1/output2"));
        spec.SortBy({"Int32"});
        client.Sort(spec);

        reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output2"));
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 2);
    }

    Y_UNIT_TEST(COMPLEX_PROTO_SORT_TEST) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TComplexRow;
        
        auto clientHolder = NRTYT::MakeClient(".");
        auto& client = *clientHolder.Get();
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        client.Create("//test1/output", NYT::NT_TABLE);
        TComplexRow sample;
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TComplexRow first;
            first.MutableRow()->SetString("second");
            first.MutableRow()->SetInt32(0);
            first.MutableRow()->SetFixed64(5);
            first.SetKey(1);
            writer->AddRow(first, 0);

            first.MutableRow()->SetString("first");
            first.SetKey(0);
            writer->AddRow(first, 0);
        }

        NYT::TSortOperationSpec spec;
        spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
        spec.SortBy({"Key"});
        client.Sort(spec);

        auto reader = client.CreateTableReader<NRTYT::NTesting::TComplexRow>(NYT::TRichYPath("//test1/output"));
        TComplexRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetRow().GetString(), "first");
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetRow().GetString(), "second");
    }

    Y_UNIT_TEST(PROTO_SORT_TEST_EXISTING_OUTPUT) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        auto clientHolder = NRTYT::MakeClient(".");
        auto& client = *clientHolder.Get();

        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        client.Create("//test1/output", NYT::NT_TABLE);
        TRow sample;
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetString("abcabc");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);

            TRow second;
            second.SetString("");
            second.SetInt32(2);
            second.SetFixed64(2);
            writer->AddRow(second, 0);
        }

        NYT::TSortOperationSpec spec;
        spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
        spec.SortBy({"Fixed64", "Int32"});
        client.Sort(spec);

        {
            auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
            TRow row;
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
            UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 2);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
            UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        }

        spec = {};
        spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
        spec.SortBy({"Int32"});
        client.Sort(spec);

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
        TRow row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "abcabc");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 2);
    }


    Y_UNIT_TEST(PROTO_SORT_CASES_TEST) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        client.Create("//test1/output", NYT::NT_TABLE);
        TRow sample;
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetString("a");
            first.SetInt32(0);
            first.SetFixed64(5);
            writer->AddRow(first, 0);

            TRow second;
            second.SetString("b");
            second.SetInt32(2);
            second.SetFixed64(5);
            writer->AddRow(second, 0);

            TRow third;
            second.SetString("c");
            second.SetInt32(0);
            second.SetFixed64(4);
            writer->AddRow(second, 0);
        }

        NYT::TSortOperationSpec spec;
        spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
        spec.SortBy({"Fixed64", "Int32"});
        client.Sort(spec);

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "c");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 4);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "a");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "b");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);

        spec = {};
        spec.AddInput(NYT::TRichYPath("//test1/output")).Output(NYT::TRichYPath("//test1/output2"));
        spec.SortBy({"Int32", "String"});
        client.Sort(spec);

        reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output2"));
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "a");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "c");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 0);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 4);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetString(), "b");
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 5);
    }
}

Y_UNIT_TEST_SUITE(RTYT_SIMPLE_REDUCER) {
    Y_UNIT_TEST(PROTO_REDUCER) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetFixed64(0);
            first.SetInt32(1);
            writer->AddRow(first, 0);

            TRow second;
            second.SetFixed64(0);
            second.SetInt32(2);
            writer->AddRow(second, 0);

            second.SetFixed64(1);
            second.SetInt32(2);
            writer->AddRow(second, 0);
        }

        {
            NYT::TSortOperationSpec spec;
            spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
            spec.SortBy({"Fixed64", "Int32"});
            client.Sort(spec);
        }

        NYT::TReduceOperationSpec spec;
        spec.AddInput<TRow>(NYT::TRichYPath("//test1/output")).AddOutput<TRow>(NYT::TRichYPath("//test1/output2")).ReduceBy({"Fixed64"});
        client.Reduce(spec, new TSimpleProtoReducer(), NYT::TOperationOptions());

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output2"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 3);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
    }

    Y_UNIT_TEST(PROTO_MAP_REDUCE) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/input", NYT::NT_TABLE);
        {
            auto writer = client.CreateTableWriter("//test1/input", *sample.GetDescriptor());
            TRow first;
            first.SetFixed64(0);
            first.SetInt32(1);
            writer->AddRow(first, 0);

            TRow second;
            second.SetFixed64(0);
            second.SetInt32(2);
            writer->AddRow(second, 0);

            second.SetFixed64(1);
            second.SetInt32(2);
            writer->AddRow(second, 0);
        }

        {
            NYT::TMapOperationSpec spec;
            spec.AddInput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/input")).AddOutput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/test2"));
            client.Map(spec, new TProtoMultiplyingMapper(), NYT::TOperationOptions());
        }

        {
            NYT::TSortOperationSpec spec;
            spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
            spec.SortBy({"Fixed64", "Int32"});
            client.Sort(spec);
        }

        NYT::TReduceOperationSpec spec;
        spec.AddInput<TRow>(NYT::TRichYPath("//test1/output")).AddOutput<TRow>(NYT::TRichYPath("//test1/output2")).ReduceBy({"Fixed64"});
        client.Reduce(spec, new TSimpleProtoReducer(), NYT::TOperationOptions());

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output2"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 3);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 2);
    }

    Y_UNIT_TEST(PROTO_MAP_REDUCE_DUPLICATE) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions().IgnoreExisting(true));
        for (int i = 0; i < 5; i++) {
            {
                auto writer = client.CreateTableWriter("//test1/input", *sample.GetDescriptor());
                TRow first;
                first.SetFixed64(0);
                first.SetInt32(1);
                writer->AddRow(first, 0);

                TRow second;
                second.SetFixed64(0);
                second.SetInt32(2);
                writer->AddRow(second, 0);

                second.SetFixed64(1);
                second.SetInt32(2);
                writer->AddRow(second, 0);
            }

            {
                Cerr << "xxxx Map block begins" << Endl;
                NYT::TMapOperationSpec spec;
                spec.AddInput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/input")).AddOutput<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/test2"));
                client.Map(spec, new TProtoMultiplyingMapper(), NYT::TOperationOptions());
                Cerr << "xxxx Map block ends" << Endl;
            }

            {
                Cerr << "xxxx Sort block begins" << Endl;
                NYT::TSortOperationSpec spec;
                spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
                spec.SortBy({"Fixed64", "Int32"});
                client.Sort(spec);
                Cerr << "xxxx Sort block ends" << Endl;
            }

            Cerr << "xxxx Reduce block begins" << Endl;
            NYT::TReduceOperationSpec spec;
            spec.AddInput<TRow>(NYT::TRichYPath("//test1/output")).AddOutput<TRow>(NYT::TRichYPath("//test1/output2")).ReduceBy({"Fixed64"});
            client.Reduce(spec, new TSimpleProtoReducer(), NYT::TOperationOptions());

            auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output2"));
            TRow row;
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 3);
            UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
            reader->Next();
            row = reader->GetRow();
            UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
            UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 2);
            Cerr << "xxxx Reduce block ends" << Endl;
        }
    }


    Y_UNIT_TEST(PROTO_REDUCER_SAVELOAD) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        NRTYT::TClientBase client(".");
        client.Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client.Create("//test1/test2", NYT::NT_TABLE);
        {
            auto writer = client.CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetFixed64(0);
            first.SetInt32(1);
            writer->AddRow(first, 0);

            TRow second;
            second.SetFixed64(0);
            second.SetInt32(2);
            writer->AddRow(second, 0);

            second.SetFixed64(1);
            second.SetInt32(2);
            writer->AddRow(second, 0);
        }

        {
            NYT::TSortOperationSpec spec;
            spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
            spec.SortBy({"Fixed64", "Int32"});
            client.Sort(spec);
        }

        NYT::TReduceOperationSpec spec;
        spec.AddInput<TRow>(NYT::TRichYPath("//test1/test2")).AddOutput<TRow>(NYT::TRichYPath("//test1/output")).ReduceBy({"Fixed64"});
        client.Reduce(spec, new TSimpleProtoReducer(5), NYT::TOperationOptions());

        auto reader = client.CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 15);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 10);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);
    }
}

Y_UNIT_TEST_SUITE(RTYT_ROBOT_REDUCER) {
    Y_UNIT_TEST(PROTO_REDUCER) {
        Cleanup();
        path.MkDir();
        NFs::SetCurrentWorkingDirectory(path);

        using NRTYT::NTesting::TRow;
        TRow sample;
        NYT::IClientBasePtr client = new NRTYT::TClientBase(".");
        client->Create("//test1", NYT::NT_MAP, NYT::TCreateOptions());
        client->Create("//test1/test2", NYT::NT_TABLE);
        {
            auto writer = client->CreateTableWriter("//test1/test2", *sample.GetDescriptor());
            TRow first;
            first.SetFixed64(0);
            first.SetInt32(1);
            writer->AddRow(first, 0);

            TRow second;
            second.SetFixed64(0);
            second.SetInt32(2);
            writer->AddRow(second, 0);

            second.SetFixed64(1);
            second.SetInt32(2);
            writer->AddRow(second, 0);
        }

        {
            NYT::TSortOperationSpec spec;
            spec.AddInput(NYT::TRichYPath("//test1/test2")).Output(NYT::TRichYPath("//test1/output"));
            spec.SortBy({"Fixed64", "Int32"});
            client->Sort(spec);
        }

        NJupiter::TReduceCmd<TSimpleProtoReducer> cmd(client);
        NJupiter::TTable<TRow> inputTable(client, NYT::TRichYPath("//test1/output"));
        NJupiter::TTable<TRow> outputTable(client, NYT::TRichYPath("//test1/output2"));

        cmd.Input(inputTable).Output(outputTable).ReduceBy({"Fixed64"}).Do();

        auto reader = client->CreateTableReader<NRTYT::NTesting::TRow>(NYT::TRichYPath("//test1/output2"));
        TRow row;
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 3);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 0);
        reader->Next();
        row = reader->GetRow();
        UNIT_ASSERT_VALUES_EQUAL(row.GetInt32(), 2);
        UNIT_ASSERT_VALUES_EQUAL(row.GetFixed64(), 1);

    }
};
