Merge branch 'main' of https://github.com/apache/arrow into fix-mingw-json-test-segfault

Author: vanshaj2023

cpp/src/arrow/filesystem/s3fs.cc

@@ -119,7 +119,6 @@
 #include "arrow/util/string.h"
 #include "arrow/util/task_group.h"
 #include "arrow/util/thread_pool.h"
-#include "arrow/util/value_parsing.h"
 
 namespace arrow::fs {
 
@@ -3579,9 +3578,10 @@ S3GlobalOptions S3GlobalOptions::Defaults() {
     log_level = S3LogLevel::Off;
   }
 
-  value = arrow::internal::GetEnvVar("ARROW_S3_THREADS").ValueOr("1");
-  if (uint32_t u; ::arrow::internal::ParseUnsigned(value.data(), value.size(), &u)) {
-    num_event_loop_threads = u;
+  auto maybe_num_threads =
+      arrow::internal::GetEnvVarInteger("ARROW_S3_THREADS", /*min_value=*/1);
+  if (maybe_num_threads.ok()) {
+    num_event_loop_threads = static_cast<int>(*maybe_num_threads);
   }
   return S3GlobalOptions{log_level, num_event_loop_threads};
 }

cpp/src/arrow/io/interfaces.cc

@@ -390,23 +390,15 @@ namespace {
 constexpr int kDefaultNumIoThreads = 8;
 
 std::shared_ptr<ThreadPool> MakeIOThreadPool() {
-  int threads = 0;
-  auto maybe_env_var = ::arrow::internal::GetEnvVar("ARROW_IO_THREADS");
-  if (maybe_env_var.ok()) {
-    auto str = *std::move(maybe_env_var);
-    if (!str.empty()) {
-      try {
-        threads = std::stoi(str);
-      } catch (...) {
-      }
-      if (threads <= 0) {
-        ARROW_LOG(WARNING)
-            << "ARROW_IO_THREADS does not contain a valid number of threads "
-               "(should be an integer > 0)";
-      }
-    }
+  int threads = kDefaultNumIoThreads;
+  auto maybe_num_threads = ::arrow::internal::GetEnvVarInteger(
+      "ARROW_IO_THREADS", /*min_value=*/1, /*max_value=*/std::numeric_limits<int>::max());
+  if (maybe_num_threads.ok()) {
+    threads = static_cast<int>(*maybe_num_threads);
+  } else if (!maybe_num_threads.status().IsKeyError()) {
+    maybe_num_threads.status().Warn();
   }
-  auto maybe_pool = ThreadPool::MakeEternal(threads > 0 ? threads : kDefaultNumIoThreads);
+  auto maybe_pool = ThreadPool::MakeEternal(threads);
   if (!maybe_pool.ok()) {
     maybe_pool.status().Abort("Failed to create global IO thread pool");
   }

cpp/src/arrow/ipc/message.cc

@@ -21,6 +21,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -363,9 +364,13 @@ Result<std::unique_ptr<Message>> ReadMessage(std::shared_ptr<Buffer> metadata,
   }
 }
 
-Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_length,
-                                             io::RandomAccessFile* file,
-                                             const FieldsLoaderFunction& fields_loader) {
+// Common helper for the two ReadMessage overloads that take a file + offset.
+// When body_length is provided, metadata and body are read in a single IO.
+// When body_length is absent, metadata is read first, then the body is read
+// separately.
+static Result<std::unique_ptr<Message>> ReadMessageInternal(
+    int64_t offset, int32_t metadata_length, std::optional<int64_t> body_length,
+    io::RandomAccessFile* file, const FieldsLoaderFunction& fields_loader) {
   std::unique_ptr<Message> result;
   auto listener = std::make_shared<AssignMessageDecoderListener>(&result);
   MessageDecoder decoder(listener);
@@ -375,15 +380,18 @@ Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_le
                            decoder.next_required_size());
   }
 
-  // TODO(GH-48846): we should take a body_length just like ReadMessageAsync
-  // and read metadata + body in one go.
-  ARROW_ASSIGN_OR_RAISE(auto metadata, file->ReadAt(offset, metadata_length));
+  // When body_length is known, read metadata + body in one IO call.
+  // Otherwise, read only metadata first.
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Buffer> metadata,
+                        file->ReadAt(offset, metadata_length + body_length.value_or(0)));
+
   if (metadata->size() < metadata_length) {
     return Status::Invalid("Expected to read ", metadata_length,
                            " metadata bytes at offset ", offset, " but got ",
                            metadata->size());
   }
-  ARROW_RETURN_NOT_OK(decoder.Consume(metadata));
+
+  ARROW_RETURN_NOT_OK(decoder.Consume(SliceBuffer(metadata, 0, metadata_length)));
 
   switch (decoder.state()) {
     case MessageDecoder::State::INITIAL:
@@ -398,14 +406,23 @@ Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_le
     case MessageDecoder::State::BODY: {
       std::shared_ptr<Buffer> body;
       if (fields_loader) {
+        // Selective field loading: allocate a body buffer and read only the
+        // requested field ranges into it.
         ARROW_ASSIGN_OR_RAISE(
             body, AllocateBuffer(decoder.next_required_size(), default_memory_pool()));
         RETURN_NOT_OK(ReadFieldsSubset(offset, metadata_length, file, fields_loader,
-                                       metadata, decoder.next_required_size(), body));
+                                       SliceBuffer(metadata, 0, metadata_length),
+                                       decoder.next_required_size(), body));
+      } else if (body_length.has_value()) {
+        // Body was already read as part of the combined IO; just slice it out.
+        body = SliceBuffer(metadata, metadata_length,
+                           std::min(*body_length, metadata->size() - metadata_length));
       } else {
+        // Body length was unknown; do a separate IO to read the body.
         ARROW_ASSIGN_OR_RAISE(
             body, file->ReadAt(offset + metadata_length, decoder.next_required_size()));
       }
+
       if (body->size() < decoder.next_required_size()) {
         return Status::IOError("Expected to be able to read ",
                                decoder.next_required_size(),
@@ -421,6 +438,21 @@ Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_le
   }
 }
 
+Result<std::unique_ptr<Message>> ReadMessage(int64_t offset, int32_t metadata_length,
+                                             io::RandomAccessFile* file,
+                                             const FieldsLoaderFunction& fields_loader) {
+  return ReadMessageInternal(offset, metadata_length, /*body_length=*/std::nullopt, file,
+                             fields_loader);
+}
+
+Result<std::unique_ptr<Message>> ReadMessage(const int64_t offset,
+                                             const int32_t metadata_length,
+                                             const int64_t body_length,
+                                             io::RandomAccessFile* file) {
+  return ReadMessageInternal(offset, metadata_length, body_length, file,
+                             /*fields_loader=*/{});
+}
+
 Future<std::shared_ptr<Message>> ReadMessageAsync(int64_t offset, int32_t metadata_length,
                                                   int64_t body_length,
                                                   io::RandomAccessFile* file,

cpp/src/arrow/ipc/message.h

@@ -449,7 +449,7 @@ class ARROW_EXPORT MessageReader {
 // org::apache::arrow::flatbuf::RecordBatch*)
 using FieldsLoaderFunction = std::function<Status(const void*, io::RandomAccessFile*)>;
 
-/// \brief Read encapsulated RPC message from position in file
+/// \brief Read encapsulated IPC message from position in file
 ///
 /// Read a length-prefixed message flatbuffer starting at the indicated file
 /// offset. If the message has a body with non-zero length, it will also be
@@ -469,7 +469,27 @@ Result<std::unique_ptr<Message>> ReadMessage(
     const int64_t offset, const int32_t metadata_length, io::RandomAccessFile* file,
     const FieldsLoaderFunction& fields_loader = {});
 
-/// \brief Read encapsulated RPC message from cached buffers
+/// \brief Read encapsulated IPC message from position in file
+///
+/// Read a length-prefixed message flatbuffer starting at the indicated file
+/// offset.
+///
+/// The metadata_length includes at least the length prefix and the flatbuffer
+///
+/// \param[in] offset the position in the file where the message starts. The
+/// first 4 bytes after the offset are the message length
+/// \param[in] metadata_length the total number of bytes to read from file
+/// \param[in] body_length the number of bytes for the message body
+/// \param[in] file the seekable file interface to read from
+/// \return the message read
+
+ARROW_EXPORT
+Result<std::unique_ptr<Message>> ReadMessage(const int64_t offset,
+                                             const int32_t metadata_length,
+                                             const int64_t body_length,
+                                             io::RandomAccessFile* file);
+
+/// \brief Read encapsulated IPC message from cached buffers
 ///
 /// The buffers should contain an entire message.  Partial reads are not handled.
 ///

cpp/src/arrow/ipc/read_write_test.cc

@@ -552,9 +552,15 @@ class TestIpcRoundTrip : public ::testing::TestWithParam<MakeRecordBatch*>,
     ASSERT_OK(WriteRecordBatch(*batch, buffer_offset, mmap_.get(), &metadata_length,
                                &body_length, options_));
 
-    ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message1,
                          ReadMessage(0, metadata_length, mmap_.get()));
-    ASSERT_EQ(expected_version, message->metadata_version());
+    ASSERT_EQ(expected_version, message1->metadata_version());
+
+    ASSERT_OK_AND_ASSIGN(auto message2,
+                         ReadMessage(0, metadata_length, body_length, mmap_.get()));
+    ASSERT_EQ(expected_version, message2->metadata_version());
+
+    ASSERT_TRUE(message1->Equals(*message2));
   }
 };
 
@@ -613,6 +619,27 @@ TEST(TestReadMessage, CorruptedSmallInput) {
   ASSERT_EQ(nullptr, message);
 }
 
+TEST(TestReadMessage, ReadBodyWithLength) {
+  // Test the optimized ReadMessage(offset, meta_len, body_len, file) overload
+  std::shared_ptr<RecordBatch> batch;
+  ASSERT_OK(MakeIntRecordBatch(&batch));
+
+  ASSERT_OK_AND_ASSIGN(auto stream, io::BufferOutputStream::Create(0));
+  int32_t metadata_length;
+  int64_t body_length;
+  ASSERT_OK(WriteRecordBatch(*batch, 0, stream.get(), &metadata_length, &body_length,
+                             IpcWriteOptions::Defaults()));
+
+  ASSERT_OK_AND_ASSIGN(auto buffer, stream->Finish());
+  io::BufferReader reader(buffer);
+
+  ASSERT_OK_AND_ASSIGN(auto message,
+                       ReadMessage(0, metadata_length, body_length, &reader));
+
+  ASSERT_EQ(body_length, message->body_length());
+  ASSERT_TRUE(message->Verify());
+}
+
 TEST(TestMetadata, GetMetadataVersion) {
   ASSERT_EQ(MetadataVersion::V1, ipc::internal::GetMetadataVersion(
                                      flatbuf::MetadataVersion::MetadataVersion_V1));
@@ -1094,7 +1121,7 @@ TEST_F(RecursionLimits, ReadLimit) {
                         &schema));
 
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,
-                       ReadMessage(0, metadata_length, mmap_.get()));
+                       ReadMessage(0, metadata_length, body_length, mmap_.get()));
 
   io::BufferReader reader(message->body());
 
@@ -1119,7 +1146,7 @@ TEST_F(RecursionLimits, StressLimit) {
                           &schema));
 
     ASSERT_OK_AND_ASSIGN(std::unique_ptr<Message> message,
-                         ReadMessage(0, metadata_length, mmap_.get()));
+                         ReadMessage(0, metadata_length, body_length, mmap_.get()));
 
     DictionaryMemo empty_memo;
 
@@ -3018,25 +3045,56 @@ void GetReadRecordBatchReadRanges(
 
   auto read_ranges = tracked->get_read_ranges();
 
-  // there are 3 read IOs before reading body:
-  // 1) read magic and footer length IO
-  // 2) read footer IO
-  // 3) read record batch metadata IO
-  EXPECT_EQ(read_ranges.size(), 3 + expected_body_read_lengths.size());
   const int32_t magic_size = static_cast<int>(strlen(ipc::internal::kArrowMagicBytes));
   // read magic and footer length IO
   auto file_end_size = magic_size + sizeof(int32_t);
   auto footer_length_offset = buffer->size() - file_end_size;
   auto footer_length = bit_util::FromLittleEndian(
       util::SafeLoadAs<int32_t>(buffer->data() + footer_length_offset));
+
+  // there are at least 2 read IOs before reading body:
+  // 1) read magic and footer length IO
+  // 2) footer IO
+  EXPECT_GE(read_ranges.size(), 2);
+
+  // read magic and footer length IO
   EXPECT_EQ(read_ranges[0].length, file_end_size);
   // read footer IO
   EXPECT_EQ(read_ranges[1].length, footer_length);
-  // read record batch metadata.  The exact size is tricky to determine but it doesn't
-  // matter for this test and it should be smaller than the footer.
-  EXPECT_LE(read_ranges[2].length, footer_length);
-  for (uint32_t i = 0; i < expected_body_read_lengths.size(); i++) {
-    EXPECT_EQ(read_ranges[3 + i].length, expected_body_read_lengths[i]);
+
+  if (included_fields.empty()) {
+    // When no fields are explicitly included, the reader optimizes by
+    // reading metadata and the entire body in a single IO.
+    // Thus, there are exactly 3 read IOs in total:
+    // 1) magic and footer length
+    // 2) footer
+    // 3) record batch metadata + body
+    EXPECT_EQ(read_ranges.size(), 3);
+
+    int64_t total_body = 0;
+    for (auto len : expected_body_read_lengths) total_body += len;
+
+    // In the optimized path (included_fields is empty), the 3rd read operation
+    // fetches both the message metadata (flatbuffer) and the entire message body
+    // in one contiguous block. Therefore, its length must at least exceed the
+    // total body length by the size of the metadata.
+    EXPECT_GT(read_ranges[2].length, total_body);
+    EXPECT_LE(read_ranges[2].length, total_body + footer_length);
+  } else {
+    // When fields are filtered, we see 3 initial reads followed by N body reads
+    // (one for each field/buffer range):
+    // 1) magic and footer length
+    // 2) footer
+    // 3) record batch metadata
+    // 4) individual body buffer reads
+    EXPECT_EQ(read_ranges.size(), 3 + expected_body_read_lengths.size());
+
+    // read record batch metadata.  The exact size is tricky to determine but it doesn't
+    // matter for this test and it should be smaller than the footer.
+    EXPECT_LE(read_ranges[2].length, footer_length);
+    for (uint32_t i = 0; i < expected_body_read_lengths.size(); i++) {
+      EXPECT_EQ(read_ranges[3 + i].length, expected_body_read_lengths[i]);
+    }
   }
 }
 
@@ -3186,7 +3244,9 @@ class PreBufferingTest : public ::testing::TestWithParam<bool> {
         metadata_reads++;
       }
     }
-    ASSERT_EQ(metadata_reads, reader_->num_record_batches() - num_indices_pre_buffered);
+    // With ReadMessage optimization, non-prebuffered reads verify metadata and body
+    // in a single large read, so we no longer see small metadata-only reads here.
+    ASSERT_EQ(metadata_reads, 0);
     ASSERT_EQ(data_reads, reader_->num_record_batches());
   }
 

cpp/src/arrow/ipc/reader.cc

@@ -1236,9 +1236,15 @@ Result<std::unique_ptr<Message>> ReadMessageFromBlock(
     const FileBlock& block, io::RandomAccessFile* file,
     const FieldsLoaderFunction& fields_loader) {
   RETURN_NOT_OK(CheckAligned(block));
-  ARROW_ASSIGN_OR_RAISE(auto message, ReadMessage(block.offset, block.metadata_length,
-                                                  file, fields_loader));
-  return CheckBodyLength(std::move(message), block);
+  if (fields_loader) {
+    ARROW_ASSIGN_OR_RAISE(auto message, ReadMessage(block.offset, block.metadata_length,
+                                                    file, fields_loader));
+    return CheckBodyLength(std::move(message), block);
+  } else {
+    ARROW_ASSIGN_OR_RAISE(auto message, ReadMessage(block.offset, block.metadata_length,
+                                                    block.body_length, file));
+    return CheckBodyLength(std::move(message), block);
+  }
 }
 
 Future<std::shared_ptr<Message>> ReadMessageFromBlockAsync(

cpp/src/arrow/json/reader_test.cc

@@ -37,12 +37,11 @@
 // On MinGW, install a crash handler to log the exception address for debugging
 // intermittent segfaults. See https://github.com/apache/arrow/issues/49272
 #if defined(__MINGW32__) || defined(__MINGW64__)
-#  include <cstdio>
 #  include <windows.h>
+#  include <cstdio>
 namespace {
 LONG WINAPI ArrowCrashHandler(EXCEPTION_POINTERS* ep) {
-  fprintf(stderr, "CRASH: code=0x%lx addr=%p\n",
-          ep->ExceptionRecord->ExceptionCode,
+  fprintf(stderr, "CRASH: code=0x%lx addr=%p\n", ep->ExceptionRecord->ExceptionCode,
           ep->ExceptionRecord->ExceptionAddress);
   fflush(stderr);
   return EXCEPTION_CONTINUE_SEARCH;

cpp/src/arrow/testing/gtest_util.cc

@@ -660,21 +660,24 @@ LocaleGuard::LocaleGuard(const char* new_locale) : impl_(new Impl(new_locale)) {
 
 LocaleGuard::~LocaleGuard() {}
 
-EnvVarGuard::EnvVarGuard(const std::string& name, const std::string& value)
-    : name_(name) {
-  auto maybe_value = arrow::internal::GetEnvVar(name);
+EnvVarGuard::EnvVarGuard(std::string name, std::optional<std::string> value)
+    : name_(std::move(name)) {
+  auto maybe_value = arrow::internal::GetEnvVar(name_);
   if (maybe_value.ok()) {
-    was_set_ = true;
     old_value_ = *std::move(maybe_value);
   } else {
-    was_set_ = false;
+    old_value_ = std::nullopt;
+  }
+  if (value.has_value()) {
+    ARROW_CHECK_OK(arrow::internal::SetEnvVar(name_, *value));
+  } else {
+    ARROW_CHECK_OK(arrow::internal::DelEnvVar(name_));
   }
-  ARROW_CHECK_OK(arrow::internal::SetEnvVar(name, value));
 }
 
 EnvVarGuard::~EnvVarGuard() {
-  if (was_set_) {
-    ARROW_CHECK_OK(arrow::internal::SetEnvVar(name_, old_value_));
+  if (old_value_.has_value()) {
+    ARROW_CHECK_OK(arrow::internal::SetEnvVar(name_, *old_value_));
   } else {
     ARROW_CHECK_OK(arrow::internal::DelEnvVar(name_));
   }