1 files changed, 13 insertions, 3 deletions
diff --git a/yt/cpp/mapreduce/client/partition_reader.cpp b/yt/cpp/mapreduce/client/partition_reader.cpp
index a2ca7eb9758..d849040d854 100644
--- a/yt/cpp/mapreduce/client/partition_reader.cpp
+++ b/yt/cpp/mapreduce/client/partition_reader.cpp
@@ -12,7 +12,7 @@ class TPartitionTableReader
     : public TRawTableReader
 {
 public:
-    TPartitionTableReader(std::unique_ptr<IInputStream> input)
+    TPartitionTableReader(std::unique_ptr<IAbortableInputStream> input)
         : Input_(std::move(input))
     { }
 
@@ -32,6 +32,16 @@ public:
         return false;
     }
 
+    void Abort() override
+    {
+        Input_->Abort();
+    }
+
+    bool IsAborted() const override
+    {
+        return Input_->IsAborted();
+    }
+
 protected:
     size_t DoRead(void* buf, size_t len) override
     {
@@ -39,7 +49,7 @@ protected:
     }
 
 private:
-    std::unique_ptr<IInputStream> Input_;
+    std::unique_ptr<IAbortableInputStream> Input_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -52,7 +62,7 @@ TRawTableReaderPtr CreateTablePartitionReader(
     const TTablePartitionReaderOptions& options)
 {
 
-    auto stream = NDetail::RequestWithRetry<std::unique_ptr<IInputStream>>(
+    auto stream = NDetail::RequestWithRetry<std::unique_ptr<IAbortableInputStream>>(
         retryPolicy,
         [&] (TMutationId /*mutationId*/) {
             return rawClient->ReadTablePartition(cookie, format, options);