YT-28650: Fix crash on empty read_query_result and emit a spec-valid arrow stream for zero-row data

* Changelog entry
  Type: feature
  Component: proxy

Earlier, the arrow writer emitted an empty (0-byte) stream when there were no rows. This is not a valid Arrow
stream (<https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format>). Moreover, `read_query_result --format arrow` crashed.

This PR fixes both issues: in these situations the writer now emits a valid, empty Arrow stream (a scheme message and zero rows), so upstream parsers work correctly.

This is a behavior change.
commit_hash:2c7487061d4bb5cadcefac3c9037edb76014e9f7

4 files changed, 20 insertions, 7 deletions

diff --git a/yt/yql/providers/yt/codec/yt_codec_io.cpp b/yt/yql/providers/yt/codec/yt_codec_io.cpp
index a4016f7358c..197c0d9bb5b 100644
--- a/yt/yql/providers/yt/codec/yt_codec_io.cpp
+++ b/yt/yql/providers/yt/codec/yt_codec_io.cpp
@@ -1498,6 +1498,7 @@ public:
             auto streamReaderResult = arrow::ipc::RecordBatchStreamReader::Open(InputStream_.get());
             if (!streamReaderResult.ok() && InputStream_->EOSReached() && InputStream_->Tell().ValueOrDie() == 0) {
                 // Workaround for YT-23495
+                // TODO(dagorokhov): remove the 0-byte workaround (YT-28650)
                 return false;
             }
             StreamReader_ = ARROW_RESULT(streamReaderResult);
diff --git a/yt/yt/library/column_converters/column_converter.cpp b/yt/yt/library/column_converters/column_converter.cpp
index b361c96c86e..9371a3ebff7 100644
--- a/yt/yt/library/column_converters/column_converter.cpp
+++ b/yt/yt/library/column_converters/column_converter.cpp
@@ -67,9 +67,6 @@ TConvertedColumnRange TColumnConverters::ConvertRowsToColumns(
     const THashMap<int, TColumnSchema>& columnSchemas)
 {
     TConvertedColumnRange convertedColumnsRange;
-    if (rows.size() == 0) {
-        return convertedColumnsRange;
-    }
 
     if (IsFirstBatch_) {
         // Initialize mapping column ids to indexes.
diff --git a/yt/yt/library/formats/arrow_parser.cpp b/yt/yt/library/formats/arrow_parser.cpp
index 947d6378019..fb6526400c8 100644
--- a/yt/yt/library/formats/arrow_parser.cpp
+++ b/yt/yt/library/formats/arrow_parser.cpp
@@ -2158,6 +2158,7 @@ public:
 
     void Finish() override
     {
+        // TODO(dagorokhov): reject empty (0-byte) input (YT-28650)
         if (LastState_ == EListenerState::InProgress) {
             THROW_ERROR_EXCEPTION("Unexpected end of stream");
         }
diff --git a/yt/yt/library/formats/arrow_writer.cpp b/yt/yt/library/formats/arrow_writer.cpp
index 5615de42823..54756b7cff6 100644
--- a/yt/yt/library/formats/arrow_writer.cpp
+++ b/yt/yt/library/formats/arrow_writer.cpp
@@ -2424,7 +2424,7 @@ public:
         TableCount_ = tableSchemas.size();
         ColumnSchemas_.resize(tableSchemas.size());
         TableIdToIndex_.resize(tableSchemas.size());
-        IsFirstBatchForSpecificTable_.assign(tableSchemas.size(), false);
+        IsTableInitialized_.assign(tableSchemas.size(), false);
 
         for (int tableIndex = 0; tableIndex < std::ssize(tableSchemas); ++tableIndex) {
             THashSet<std::string> columnNames;
@@ -2456,6 +2456,20 @@ public:
         }
     }
 
+    TFuture<void> Close() override
+    {
+        try {
+            for (int tableIndex = 0; tableIndex < TableCount_; ++tableIndex) {
+                if (!IsTableInitialized_[tableIndex]) {
+                    WriteRowsForSingleTable(TRange<TUnversionedRow>(), tableIndex);
+                }
+            }
+        } catch (const std::exception& ex) {
+            SetError(TError(ex));
+        }
+        return TSchemalessFormatWriterBase::Close();
+    }
+
 private:
     void Reset()
     {
@@ -2593,7 +2607,7 @@ private:
     std::vector<IUnversionedColumnarRowBatch::TDictionaryId> ArrowDictionaryIds_;
     std::vector<TColumnConverters> ColumnConverters_;
     std::vector<THashMap<int, int>> TableIdToIndex_;
-    std::vector<bool> IsFirstBatchForSpecificTable_;
+    std::vector<bool> IsTableInitialized_;
     TConvertedColumnRange MissingColumns_;
 
     std::vector<TArrowWriterBuffer> Buffers_;
@@ -2634,7 +2648,7 @@ private:
 
     void PrepareColumns(const TRange<const TBatchColumn*>& batchColumns, int tableIndex)
     {
-        if (!IsFirstBatchForSpecificTable_[tableIndex]) {
+        if (!IsTableInitialized_[tableIndex]) {
             int currentIndex = 0;
             for (const auto& columnSchema : ColumnSchemas_[tableIndex]) {
                 auto columnId = columnSchema.first;
@@ -2644,7 +2658,7 @@ private:
                 }
             }
 
-            IsFirstBatchForSpecificTable_[tableIndex] = true;
+            IsTableInitialized_[tableIndex] = true;
         }
 
         TypedColumns_.resize(TableIdToIndex_[tableIndex].size());