Fix writing of compressed ORC files with large stripe footers (rapidsai#17700)

In ORC, stripe footers can be compressed, the same way as the data. This means that compressed footers need to be written in multiple blocks if they are larger than the maximum block size. This applies even if the footer is actually uncompressed (in this case a flag in the block header is raised).
Currently, the ORC writer does not take into account that footer can be larger than max block size, and writes the entire thing in a single block, which is not valid.
The issue only applies to compressed files. Uncompressed files do not apply this limitation to the footers.

This PR changes the way the stripe footers are written to account for this case. The output hasn't changed for files with small stripe footers.

Authors:
  - Vukasin Milovanovic (https://github.com/vuule)

Approvers:
  - Nghia Truong (https://github.com/ttnghia)
  - https://github.com/nvdbaranec

URL: rapidsai#17700

Author: vuule

cpp/src/io/orc/writer_impl.cu

@@ -2595,16 +2595,30 @@ void writer::impl::write_orc_data_to_sink(encoded_data const& enc_data,
           : 0;
       if (orc_table.column(i - 1).orc_kind() == TIMESTAMP) { sf.writerTimezone = "UTC"; }
     }
-    ProtobufWriter pbw((_compression != compression_type::NONE) ? 3 : 0);
+
+    ProtobufWriter pbw;
     pbw.write(sf);
-    stripe.footerLength = pbw.size();
-    if (_compression != compression_type::NONE) {
-      uint32_t uncomp_sf_len = (stripe.footerLength - 3) * 2 + 1;
-      pbw.buffer()[0]        = static_cast<uint8_t>(uncomp_sf_len >> 0);
-      pbw.buffer()[1]        = static_cast<uint8_t>(uncomp_sf_len >> 8);
-      pbw.buffer()[2]        = static_cast<uint8_t>(uncomp_sf_len >> 16);
+    if (_compression == compression_type::NONE) {
+      _out_sink->host_write(pbw.data(), pbw.size());
+      stripe.footerLength = pbw.size();
+    } else {
+      std::size_t bytes_written  = 0;
+      std::size_t written_sf_len = 0;
+      while (written_sf_len < pbw.size()) {
+        auto const block_size = std::min(_compression_blocksize, pbw.size() - written_sf_len);
+        auto const header_val = block_size * 2 + 1;  // 1 means uncompressed
+        CUDF_EXPECTS(header_val >> 24 == 0, "Block length exceeds maximum size");
+        std::array const header{static_cast<uint8_t>(header_val >> 0),
+                                static_cast<uint8_t>(header_val >> 8),
+                                static_cast<uint8_t>(header_val >> 16)};
+
+        _out_sink->host_write(header.data(), header.size());
+        _out_sink->host_write(pbw.data() + written_sf_len, block_size);
+        written_sf_len += block_size;
+        bytes_written += header.size() + block_size;
+      }
+      stripe.footerLength = bytes_written;
     }
-    _out_sink->host_write(pbw.data(), pbw.size());
   }
   for (auto const& task : write_tasks) {
     task.wait();

cpp/tests/io/orc_test.cpp

@@ -2196,4 +2196,23 @@ TEST_F(OrcChunkedWriterTest, NoDataInSinkWhenNoWrite)
   EXPECT_EQ(out_buffer.size(), 0);
 }
 
+TEST_F(OrcWriterTest, MultipleBlocksInStripeFooter)
+{
+  std::vector<std::string> vals_col(8, "a");
+  str_col col{vals_col.begin(), vals_col.end()};
+  cudf::column_view col_view = col;
+  table_view expected(std::vector<cudf::column_view>{6400, col_view});
+
+  std::vector<char> out_buffer;
+  cudf::io::orc_writer_options out_opts =
+    cudf::io::orc_writer_options::builder(cudf::io::sink_info{&out_buffer}, expected);
+  // Write with compression on (default)
+  cudf::io::write_orc(out_opts);
+
+  cudf::io::orc_reader_options in_opts = cudf::io::orc_reader_options::builder(
+    cudf::io::source_info{out_buffer.data(), out_buffer.size()});
+  auto result = cudf::io::read_orc(in_opts);
+  CUDF_TEST_EXPECT_TABLES_EQUAL(expected, result.tbl->view());
+}
+
 CUDF_TEST_PROGRAM_MAIN()