aboutsummaryrefslogtreecommitdiffstats
path: root/yt/cpp/mapreduce/io/job_writer.cpp
blob: 21b5ef6b918144f9a1661e7973fb666e3cc2cc26 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
#include "job_writer.h"

#include <yt/cpp/mapreduce/interface/io.h>

#include <util/system/file.h>

namespace NYT {
namespace NDetail {

////////////////////////////////////////////////////////////////////////////////

TJobWriterStream::TJobWriterStream(int fd)
    : TJobWriterStream(Duplicate(fd))
{ }

TJobWriterStream::TJobWriterStream(const TFile& file)
    : FDFile(file)
    , FDOutput(FDFile)
    , BufferedOutput(&FDOutput, BufferSize)
{ }

////////////////////////////////////////////////////////////////////////////////

} // namespace NDetail

////////////////////////////////////////////////////////////////////////////////

TJobWriter::TJobWriter(size_t outputTableCount)
{
    for (size_t i = 0; i < outputTableCount; ++i) {
        Streams_.emplace_back(std::make_unique<NDetail::TJobWriterStream>(static_cast<int>(i * 3 + 1)));
    }
}

TJobWriter::TJobWriter(const TVector<TFile>& fileList)
{
    for (const auto& f : fileList) {
        Streams_.emplace_back(std::make_unique<NDetail::TJobWriterStream>(f));
    }
}

size_t TJobWriter::GetStreamCount() const
{
    return Streams_.size();
}

IOutputStream* TJobWriter::GetStream(size_t tableIndex) const
{
    if (tableIndex >= Streams_.size()) {
        ythrow TIOException() <<
            "Table index " << tableIndex <<
            " is out of range [0, " << Streams_.size() << ")";
    }
    return &Streams_[tableIndex]->BufferedOutput;
}

void TJobWriter::OnRowFinished(size_t)
{ }

size_t TJobWriter::GetBufferMemoryUsage() const
{
    return NDetail::TJobWriterStream::BufferSize * GetStreamCount();
}

////////////////////////////////////////////////////////////////////////////////

THolder<IProxyOutput> CreateRawJobWriter(size_t outputTableCount)
{
    return ::MakeHolder<TJobWriter>(outputTableCount);
}

////////////////////////////////////////////////////////////////////////////////

TSingleStreamJobWriter::TSingleStreamJobWriter(size_t tableIndex)
    : TableIndex_(tableIndex)
    , Stream_(std::make_unique<NDetail::TJobWriterStream>(static_cast<int>(tableIndex * 3 + 1)))
{ }

size_t TSingleStreamJobWriter::GetStreamCount() const
{
    return 1;
}

IOutputStream* TSingleStreamJobWriter::GetStream(size_t tableIndex) const
{
    if (tableIndex != TableIndex_) {
        ythrow TIOException() <<
            "Table index " << tableIndex <<
            " does not match this SignleTableJobWriter with index " << TableIndex_;
    }
    return &Stream_->BufferedOutput;
}

void TSingleStreamJobWriter::OnRowFinished(size_t)
{ }

size_t TSingleStreamJobWriter::GetBufferMemoryUsage() const
{
    return NDetail::TJobWriterStream::BufferSize * GetStreamCount();
}

////////////////////////////////////////////////////////////////////////////////

} // namespace NYT