io.h

Go to the documentation of this file.

#pragma once

 
 
#include "fwd.h"
 
#include "abortable_stream.h"
#include "client_method_options.h"
#include "common.h"
#include "distributed_session.h"
#include "format.h"
#include "node.h"
#include "mpl.h"
#include "skiff_row.h"
 
#include <google/protobuf/message.h>
 
#include <util/stream/input.h>
#include <util/stream/output.h>
#include <util/generic/yexception.h>
#include <util/generic/maybe.h>
 
namespace NYT {


template<class... TProtoRowTypes>
class TProtoOneOf
{
public:
    static_assert(
        (TIsBaseOf<::google::protobuf::Message, TProtoRowTypes>::Value && ...),
        "Template parameters can only be protobuf types");
 
    TProtoOneOf() = delete;
};

template<class... TSkiffRowTypes>
class TSkiffRowOneOf
{
public:
    static_assert(
        (TIsSkiffRow<TSkiffRowTypes>::value && ...),
        "Template parameters can only be SkiffRow types");
 
    TSkiffRowOneOf() = delete;
};


namespace NDetail {

 
template <class TTuple>
struct TProtoOneOfFromTuple;
 
template <class... Ts>
struct TProtoOneOfFromTuple<std::tuple<Ts...>>
{
    using TType = TProtoOneOf<Ts...>;
};
 
template <class... Ts>
struct TProtoOneOfUnique
{
    using TTuple = typename TUniqueTypes<std::tuple<>, std::tuple<Ts...>>::TType;
    using TType = typename TProtoOneOfFromTuple<TTuple>::TType;
};

 
} // namespace NDetail

 
struct INodeReaderImpl;
struct IYaMRReaderImpl;
struct IProtoReaderImpl;
struct ISkiffRowReaderImpl;
struct INodeWriterImpl;
struct IYaMRWriterImpl;
struct IProtoWriterImpl;


class TIOException
    : public yexception
{ };


class IFileReader
    : public TThrRefBase
    , public IAbortableInputStream
{ };

class IFileWriter
    : public TThrRefBase
    , public IOutputStream
{
public:
    virtual size_t GetBufferMemoryUsage() const
    {
        return 0;
    }
};
 
class IFileFragmentWriter
    : public TThrRefBase
    , public IOutputStream
{
public:
    virtual TWriteFileFragmentResult GetWriteFragmentResult() const = 0;
};


class TRawTableReader
    : public TThrRefBase
    , public IAbortableInputStream
{
public:
    virtual bool Retry(
        const TMaybe<ui32>& rangeIndex,
        const TMaybe<ui64>& rowIndex,
        const std::exception_ptr& error) = 0;

    virtual void ResetRetries() = 0;

    virtual bool HasRangeIndices() const = 0;
};

class TRawTableWriter
    : public TThrRefBase
    , public IOutputStream
{
public:
    virtual void NotifyRowEnd() = 0;

    virtual void Abort()
    { }
 
    virtual size_t GetBufferMemoryUsage() const
    {
        return 0;
    }
};

class IProxyOutput
{
public:
    virtual ~IProxyOutput()
    { }

    virtual size_t GetStreamCount() const = 0;

    virtual IOutputStream* GetStream(size_t tableIndex) const = 0;

    virtual void OnRowFinished(size_t tableIndex) = 0;

    virtual void Abort()
    { }
 
    virtual size_t GetBufferMemoryUsage() const
    {
        return 0;
    }
};


template <class T, class>
class TTableReader
    : public TThrRefBase
{
public:
    const T& GetRow() const;

    T MoveRow();

    void MoveRow(T* result);

    bool IsValid() const;

    void Next();

    ui32 GetTableIndex() const;

    ui32 GetRangeIndex() const;

    ui64 GetRowIndex() const;

    i64 GetTabletIndex() const;

    bool IsEndOfStream() const;

    bool IsRawReaderExhausted() const;

    void Abort();

    bool IsAborted() const;
};

template <class T>
class TTableReaderIterator
{
public:
    explicit TTableReaderIterator<T>(TTableReader<T>* reader)
    {
        if (reader && reader->IsValid()) {
            Reader_ = reader;
        } else {
            Reader_ = nullptr;
        }
    }

    bool operator==(const TTableReaderIterator& it) const
    {
        return Reader_ == it.Reader_;
    }

    TTableReader<T>& operator*()
    {
        return *Reader_;
    }

    const TTableReader<T>& operator*() const
    {
        return *Reader_;
    }

    TTableReaderIterator& operator++()
    {
        Reader_->Next();
        if (!Reader_->IsValid()) {
            Reader_ = nullptr;
        }
        return *this;
    }
 
private:
    TTableReader<T>* Reader_;
};

template <class T>
TTableReaderIterator<T> begin(TTableReader<T>& reader) // NOLINT
{
    return TTableReaderIterator<T>(&reader);
}

template <class T>
TTableReaderIterator<T> end(TTableReader<T>&) // NOLINT
{
    return TTableReaderIterator<T>(nullptr);
}


template <class T, class>
class TTableRangesReader
    : public TThrRefBase
{
public:
    TTableReader<T>& GetRange();

    bool IsValid() const;

    void Next();
};


template <class T, class>
class TTableWriter
    : public TThrRefBase
{
public:
    void AddRow(const T& row);

    void Finish();
 
    size_t GetBufferMemoryUsage() const;
};


template <class T>
class ITableFragmentWriter
    : public TThrRefBase
{
public:
    virtual TWriteTableFragmentResult GetWriteFragmentResult() const = 0;

    virtual void AddRow(const T& row) = 0;

    virtual void Finish() = 0;
};


struct TYaMRRow
{
    TStringBuf Key;

    TStringBuf SubKey;

    TStringBuf Value;
};


class IIOClient
{
public:
    virtual ~IIOClient() = default;

    virtual IFileReaderPtr CreateFileReader(
        const TRichYPath& path,
        const TFileReaderOptions& options = TFileReaderOptions()) = 0;

    virtual IFileReaderPtr CreateFilePartitionReader(
        const TString& cookie,
        const TFilePartitionReaderOptions& options = {}) = 0;

    virtual IFileWriterPtr CreateFileWriter(
        const TRichYPath& path,
        const TFileWriterOptions& options = TFileWriterOptions()) = 0;

    virtual IFileFragmentWriterPtr CreateFileFragmentWriter(
        const TDistributedWriteFileCookie& cookie,
        const TFileFragmentWriterOptions& options = TFileFragmentWriterOptions()) = 0;

    template <class T>
    TTableReaderPtr<T> CreateTableReader(
        const TRichYPath& path,
        const TTableReaderOptions& options = TTableReaderOptions());

    template <class T>
    TTableWriterPtr<T> CreateTableWriter(
        const TRichYPath& path,
        const TTableWriterOptions& options = TTableWriterOptions());

    virtual TTableWriterPtr<::google::protobuf::Message> CreateTableWriter(
        const TRichYPath& path,
        const ::google::protobuf::Descriptor& descriptor,
        const TTableWriterOptions& options = TTableWriterOptions()) = 0;

    template <class T>
    ITableFragmentWriterPtr<T> CreateTableFragmentWriter(
        const TDistributedWriteTableCookie& cookie,
        const TTableFragmentWriterOptions& options = {});

    virtual TRawTableReaderPtr CreateRawReader(
        const TRichYPath& path,
        const TFormat& format,
        const TTableReaderOptions& options = TTableReaderOptions()) = 0;

    virtual TRawTableWriterPtr CreateRawWriter(
        const TRichYPath& path,
        const TFormat& format,
        const TTableWriterOptions& options = TTableWriterOptions()) = 0;

    virtual TRawTableReaderPtr CreateRawTablePartitionReader(
        const TString& cookie,
        const TFormat& format,
        const TTablePartitionReaderOptions& options = {}) = 0;

    template <class T>
    TTableReaderPtr<T> CreateTablePartitionReader(
        const TString& cookie,
        const TTablePartitionReaderOptions& options = {});

    virtual IFileReaderPtr CreateBlobTableReader(
        const TYPath& path,
        const TKey& blobId,
        const TBlobTableReaderOptions& options = TBlobTableReaderOptions()) = 0;
 
private:
    virtual ::TIntrusivePtr<INodeReaderImpl> CreateNodeReader(
        const TRichYPath& path, const TTableReaderOptions& options) = 0;
 
    virtual ::TIntrusivePtr<IYaMRReaderImpl> CreateYaMRReader(
        const TRichYPath& path, const TTableReaderOptions& options) = 0;
 
    virtual ::TIntrusivePtr<IProtoReaderImpl> CreateProtoReader(
        const TRichYPath& path,
        const TTableReaderOptions& options,
        const ::google::protobuf::Message* prototype) = 0;
 
    virtual ::TIntrusivePtr<ISkiffRowReaderImpl> CreateSkiffRowReader(
        const TRichYPath& path,
        const TTableReaderOptions& options,
        const ISkiffRowSkipperPtr& skipper,
        const NSkiff::TSkiffSchemaPtr& requestedSchema,
        const NSkiff::TSkiffSchemaPtr& parserSchema) = 0;
 
    virtual ::TIntrusivePtr<INodeReaderImpl> CreateNodeTablePartitionReader(
        const TString& cookie, const TTablePartitionReaderOptions& options) = 0;
 
    virtual ::TIntrusivePtr<IProtoReaderImpl> CreateProtoTablePartitionReader(
        const TString& cookie,
        const TTablePartitionReaderOptions& options,
        const ::google::protobuf::Message* prototype) = 0;
 
    virtual ::TIntrusivePtr<ISkiffRowReaderImpl> CreateSkiffRowTablePartitionReader(
        const TString& cookie,
        const TTablePartitionReaderOptions& options,
        const ISkiffRowSkipperPtr& skipper,
        const NSkiff::TSkiffSchemaPtr& requestedSchema,
        const NSkiff::TSkiffSchemaPtr& parserSchema) = 0;
 
    virtual ::TIntrusivePtr<INodeWriterImpl> CreateNodeWriter(
        const TRichYPath& path, const TTableWriterOptions& options) = 0;
 
    virtual ::TIntrusivePtr<IYaMRWriterImpl> CreateYaMRWriter(
        const TRichYPath& path, const TTableWriterOptions& options) = 0;
 
    virtual ::TIntrusivePtr<IProtoWriterImpl> CreateProtoWriter(
        const TRichYPath& path,
        const TTableWriterOptions& options,
        const ::google::protobuf::Message* prototype) = 0;
 
    virtual ::TIntrusivePtr<ITableFragmentWriter<TNode>> CreateNodeFragmentWriter(
        const TDistributedWriteTableCookie& cookie,
        const TTableFragmentWriterOptions& options) = 0;
 
    virtual ::TIntrusivePtr<ITableFragmentWriter<TYaMRRow>> CreateYaMRFragmentWriter(
        const TDistributedWriteTableCookie& cookie,
        const TTableFragmentWriterOptions& options) = 0;
 
    virtual ::TIntrusivePtr<ITableFragmentWriter<Message>> CreateProtoFragmentWriter(
        const TDistributedWriteTableCookie& cookie,
        const TTableFragmentWriterOptions& options,
        const ::google::protobuf::Message* prototype) = 0;
};


template <typename T>
TTableReaderPtr<T> CreateTableReader(
    IInputStream* stream,
    const TTableReaderOptions& options = {});

template <class... Ts>
TTableReaderPtr<typename NDetail::TProtoOneOfUnique<Ts...>::TType> CreateProtoMultiTableReader(
    IInputStream* stream,
    const TTableReaderOptions& options = {});

template <class T>
TTableReaderPtr<T> CreateProtoMultiTableReader(
    IInputStream* stream,
    int tableCount,
    const TTableReaderOptions& options = {});

template <>
TTableReaderPtr<TNode> CreateTableReader<TNode>(
    IInputStream* stream, const TTableReaderOptions& options);

template <>
TTableReaderPtr<TYaMRRow> CreateTableReader<TYaMRRow>(
    IInputStream* stream, const TTableReaderOptions& options);
 
namespace NDetail {

::TIntrusivePtr<IProtoReaderImpl> CreateProtoReader(
    IInputStream* stream,
    const TTableReaderOptions& options,
    const ::google::protobuf::Descriptor* descriptor);
 

::TIntrusivePtr<IProtoReaderImpl> CreateProtoReader(
    IInputStream* stream,
    const TTableReaderOptions& options,
    TVector<const ::google::protobuf::Descriptor*> descriptors);
 
} // namespace NDetail


template <typename T>
TTableReaderPtr<T> CreateConcreteProtobufReader(TTableReader<Message>* reader);

template <typename T>
TTableReaderPtr<T> CreateConcreteProtobufReader(const TTableReaderPtr<Message>& reader);

template <typename T>
TTableReaderPtr<Message> CreateGenericProtobufReader(TTableReader<T>* reader);

template <typename T>
TTableReaderPtr<Message> CreateGenericProtobufReader(const TTableReaderPtr<T>& reader);

 
} // namespace NYT
 
#define IO_INL_H_
#include "io-inl.h"
#undef IO_INL_H_