apache · pablogsal · Jan 23, 2026
diff --git a/cpp/src/arrow/CMakeLists.txt b/cpp/src/arrow/CMakeLists.txt
@@ -1268,6 +1268,44 @@ add_arrow_test(misc_test
                pretty_print_test.cc
                status_test.cc)
 
+# Mimalloc interposition test for memory profiler integration (e.g., memray).
+# Linux-only: LD_PRELOAD semantics and the test interposer use Linux-specific
+# behaviour (dlsym(RTLD_NEXT), atomic constructors/destructors, ELF visibility).
+if(ARROW_MIMALLOC AND CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  # Ensure the wrapper call sites in memory_pool.cc go through the PLT so an
+  # LD_PRELOAD interposer can intercept them. Arrow's global CXX flags include
+  # -fno-semantic-interposition (see SetupCxxFlags.cmake), which lets GCC
+  # resolve same-DSO calls directly without the PLT. Override that here.
+  if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU")
+    set_source_files_properties(memory_pool.cc PROPERTIES
+                                COMPILE_OPTIONS "-fsemantic-interposition")
+  endif()
+
+  # Build the interposition shared library
+  add_library(arrow_mimalloc_interpose SHARED memory_pool_interpose.c)
+  target_link_libraries(arrow_mimalloc_interpose PRIVATE ${CMAKE_DL_LIBS})
+  set_target_properties(arrow_mimalloc_interpose PROPERTIES
+                        LIBRARY_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}")
+
+  # Build the test helper executable
+  add_executable(arrow-memory-pool-interpose-test-helper
+                 memory_pool_interpose_test_helper.cc)
+  target_link_libraries(arrow-memory-pool-interpose-test-helper PRIVATE arrow_shared)
+  set_target_properties(arrow-memory-pool-interpose-test-helper PROPERTIES
+                        RUNTIME_OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}")
+
+  # Add the interposition test
+  add_arrow_test(memory_pool_interpose_test)
+  add_dependencies(arrow-memory-pool-interpose-test
+                   arrow_mimalloc_interpose
+                   arrow-memory-pool-interpose-test-helper)
+
+  # Set environment variables for the test to find the helper and library
+  set_tests_properties(arrow-memory-pool-interpose-test PROPERTIES
+    ENVIRONMENT
+      "ARROW_INTERPOSE_TEST_HELPER=$<TARGET_FILE:arrow-memory-pool-interpose-test-helper>;ARROW_INTERPOSE_TEST_LIB=$<TARGET_FILE:arrow_mimalloc_interpose>")
+endif()
+
 add_arrow_test(public_api_test)
 set_source_files_properties(public_api_test.cc PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON)
 

diff --git a/cpp/src/arrow/memory_pool.cc b/cpp/src/arrow/memory_pool.cc
@@ -53,6 +53,35 @@
 #  include <mimalloc.h>
 #endif
 
+// PLT-exposed allocation functions for memory profiler integration (e.g., memray).
+// ARROW_EXPORT gives these default visibility so they appear in the dynamic
+// symbol table and can be overridden at load time via LD_PRELOAD. The weak
+// attribute additionally permits a stronger definition at static-link time.
+// `noinline` prevents the compiler from inlining the wrapper into call sites
+// in the same TU, which would defeat interposition. The CMake build also
+// compiles this TU with -fsemantic-interposition on GCC so the call sites
+// route through the PLT (Arrow's default is -fno-semantic-interposition,
+// which lets GCC bypass the PLT for same-DSO calls).
+#ifdef ARROW_MIMALLOC
+extern "C" {
+
+__attribute__((weak, noinline)) ARROW_EXPORT void* arrow_mimalloc_allocate(
+    size_t size, size_t alignment) {
+  return mi_malloc_aligned(size, alignment);
+}
+
+__attribute__((weak, noinline)) ARROW_EXPORT void* arrow_mimalloc_reallocate(
+    void* ptr, size_t new_size, size_t alignment) {
+  return mi_realloc_aligned(ptr, new_size, alignment);
+}
+
+__attribute__((weak, noinline)) ARROW_EXPORT void arrow_mimalloc_free(void* ptr) {
+  mi_free(ptr);
+}
+
+}  // extern "C"
+#endif  // ARROW_MIMALLOC
+
 namespace arrow {
 
 namespace memory_pool {
@@ -393,6 +422,7 @@ class SystemAllocator {
 #ifdef ARROW_MIMALLOC
 
 // Helper class directing allocations to the mimalloc allocator.
+// Uses PLT-exposed functions (arrow_mimalloc_*) to allow memory profiler interposition.
 class MimallocAllocator {
  public:
   static Status AllocateAligned(int64_t size, int64_t alignment, uint8_t** out) {
@@ -401,7 +431,7 @@ class MimallocAllocator {
       return Status::OK();
     }
     *out = reinterpret_cast<uint8_t*>(
-        mi_malloc_aligned(static_cast<size_t>(size), static_cast<size_t>(alignment)));
+        arrow_mimalloc_allocate(static_cast<size_t>(size), static_cast<size_t>(alignment)));
     if (*out == NULL) {
       return Status::OutOfMemory("malloc of size ", size, " failed");
     }
@@ -422,8 +452,8 @@ class MimallocAllocator {
       *ptr = memory_pool::internal::kZeroSizeArea;
       return Status::OK();
     }
-    *ptr = reinterpret_cast<uint8_t*>(
-        mi_realloc_aligned(previous_ptr, static_cast<size_t>(new_size), alignment));
+    *ptr = reinterpret_cast<uint8_t*>(arrow_mimalloc_reallocate(
+        previous_ptr, static_cast<size_t>(new_size), static_cast<size_t>(alignment)));
     if (*ptr == NULL) {
       *ptr = previous_ptr;
       return Status::OutOfMemory("realloc of size ", new_size, " failed");
@@ -435,7 +465,7 @@ class MimallocAllocator {
     if (ptr == memory_pool::internal::kZeroSizeArea) {
       DCHECK_EQ(size, 0);
     } else {
-      mi_free(ptr);
+      arrow_mimalloc_free(ptr);
     }
   }
 

diff --git a/cpp/src/arrow/memory_pool_interpose.c b/cpp/src/arrow/memory_pool_interpose.c
@@ -0,0 +1,89 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Interposition library for testing arrow_mimalloc_* PLT hooks.
+// This library is meant to be used with LD_PRELOAD to intercept
+// Arrow's mimalloc allocation functions.
+
+#define _GNU_SOURCE
+#include <dlfcn.h>
+#include <stdatomic.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+static atomic_size_t g_alloc_count = 0;
+static atomic_size_t g_free_count = 0;
+static atomic_size_t g_realloc_count = 0;
+static atomic_size_t g_total_allocated = 0;
+
+void* arrow_mimalloc_allocate(size_t size, size_t alignment) {
+  static void* (*real_fn)(size_t, size_t) = NULL;
+  if (!real_fn) {
+    real_fn = dlsym(RTLD_NEXT, "arrow_mimalloc_allocate");
+    if (!real_fn) {
+      fprintf(stderr, "INTERPOSE ERROR: cannot find arrow_mimalloc_allocate\n");
+      abort();
+    }
+  }
+  void* ptr = real_fn(size, alignment);
+  if (ptr) {
+    atomic_fetch_add(&g_alloc_count, 1);
+    atomic_fetch_add(&g_total_allocated, size);
+  }
+  return ptr;
+}
+
+void arrow_mimalloc_free(void* ptr) {
+  static void (*real_fn)(void*) = NULL;
+  if (!real_fn) {
+    real_fn = dlsym(RTLD_NEXT, "arrow_mimalloc_free");
+    if (!real_fn) {
+      fprintf(stderr, "INTERPOSE ERROR: cannot find arrow_mimalloc_free\n");
+      abort();
+    }
+  }
+  atomic_fetch_add(&g_free_count, 1);
+  real_fn(ptr);
+}
+
+void* arrow_mimalloc_reallocate(void* ptr, size_t new_size, size_t alignment) {
+  static void* (*real_fn)(void*, size_t, size_t) = NULL;
+  if (!real_fn) {
+    real_fn = dlsym(RTLD_NEXT, "arrow_mimalloc_reallocate");
+    if (!real_fn) {
+      fprintf(stderr, "INTERPOSE ERROR: cannot find arrow_mimalloc_reallocate\n");
+      abort();
+    }
+  }
+  void* new_ptr = real_fn(ptr, new_size, alignment);
+  atomic_fetch_add(&g_realloc_count, 1);
+  return new_ptr;
+}
+
+// Called at library unload to print statistics
+__attribute__((destructor)) static void print_interpose_stats(void) {
+  size_t allocs = atomic_load(&g_alloc_count);
+  size_t frees = atomic_load(&g_free_count);
+  size_t reallocs = atomic_load(&g_realloc_count);
+  size_t total = atomic_load(&g_total_allocated);
+
+  // Print in a parseable format
+  fprintf(stderr, "ARROW_INTERPOSE_STATS: allocs=%zu frees=%zu reallocs=%zu "
+                  "total_bytes=%zu\n",
+          allocs, frees, reallocs, total);
+}
diff --git a/cpp/src/arrow/memory_pool_interpose_test.cc b/cpp/src/arrow/memory_pool_interpose_test.cc
@@ -0,0 +1,178 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+// Test that arrow_mimalloc_* functions can be interposed via LD_PRELOAD.
+// This test spawns a subprocess with the interposition library preloaded
+// and verifies that allocations are tracked.
+
+#include <cstdlib>
+#include <regex>
+#include <sstream>
+#include <string>
+
+#ifndef _WIN32
+#  include <sys/wait.h>
+#  include <unistd.h>
+#endif
+
+#include <gtest/gtest.h>
+
+#include "arrow/testing/gtest_util.h"
+#include "arrow/util/config.h"
+#include "arrow/util/io_util.h"
+
+namespace arrow {
+
+#if defined(ARROW_MIMALLOC) && !defined(_WIN32)
+
+class TestMimallocInterpose : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    // Get paths from environment (set by CTest)
+    auto helper_path_result = internal::GetEnvVar("ARROW_INTERPOSE_TEST_HELPER");
+    auto lib_path_result = internal::GetEnvVar("ARROW_INTERPOSE_TEST_LIB");
+
+    if (!helper_path_result.ok() || !lib_path_result.ok()) {
+      GTEST_SKIP() << "ARROW_INTERPOSE_TEST_HELPER or ARROW_INTERPOSE_TEST_LIB not set";
+    }
+
+    helper_path_ = *helper_path_result;
+    lib_path_ = *lib_path_result;
+  }
+
+  // Run the helper program with LD_PRELOAD and capture output
+  struct ProcessResult {
+    int exit_code;
+    std::string stdout_str;
+    std::string stderr_str;
+  };
+
+  ProcessResult RunWithPreload() {
+    ProcessResult result;
+    result.exit_code = -1;
+
+    // Create pipes for stdout and stderr
+    int stdout_pipe[2];
+    int stderr_pipe[2];
+    if (pipe(stdout_pipe) != 0 || pipe(stderr_pipe) != 0) {
+      return result;
+    }
+
+    pid_t pid = fork();
+    if (pid == -1) {
+      return result;
+    }
+
+    if (pid == 0) {
+      // Child process
+      close(stdout_pipe[0]);
+      close(stderr_pipe[0]);
+
+      dup2(stdout_pipe[1], STDOUT_FILENO);
+      dup2(stderr_pipe[1], STDERR_FILENO);
+
+      close(stdout_pipe[1]);
+      close(stderr_pipe[1]);
+
+      // Set LD_PRELOAD
+      setenv("LD_PRELOAD", lib_path_.c_str(), 1);
+
+      // Execute helper
+      execl(helper_path_.c_str(), helper_path_.c_str(), nullptr);
+      _exit(127);  // exec failed
+    }
+
+    // Parent process
+    close(stdout_pipe[1]);
+    close(stderr_pipe[1]);
+
+    // Read stdout
+    char buffer[4096];
+    ssize_t n;
+    while ((n = read(stdout_pipe[0], buffer, sizeof(buffer) - 1)) > 0) {
+      buffer[n] = '\0';
+      result.stdout_str += buffer;
+    }
+    close(stdout_pipe[0]);
+
+    // Read stderr
+    while ((n = read(stderr_pipe[0], buffer, sizeof(buffer) - 1)) > 0) {
+      buffer[n] = '\0';
+      result.stderr_str += buffer;
+    }
+    close(stderr_pipe[0]);
+
+    // Wait for child
+    int status;
+    waitpid(pid, &status, 0);
+    if (WIFEXITED(status)) {
+      result.exit_code = WEXITSTATUS(status);
+    }
+
+    return result;
+  }
+
+  std::string helper_path_;
+  std::string lib_path_;
+};
+
+TEST_F(TestMimallocInterpose, InterpositionWorks) {
+  auto result = RunWithPreload();
+
+  // Check program succeeded
+  ASSERT_EQ(result.exit_code, 0) << "Helper exited with code " << result.exit_code
+                                 << "\nstdout: " << result.stdout_str
+                                 << "\nstderr: " << result.stderr_str;
+
+  // Check that we got the expected output from helper
+  ASSERT_NE(result.stdout_str.find("SUCCESS"), std::string::npos)
+      << "Helper did not report success\nstdout: " << result.stdout_str;
+
+  // Check that interposition library printed stats
+  ASSERT_NE(result.stderr_str.find("ARROW_INTERPOSE_STATS:"), std::string::npos)
+      << "Interposition stats not found\nstderr: " << result.stderr_str;
+
+  // Parse the stats
+  std::regex stats_regex(
+      R"(ARROW_INTERPOSE_STATS: allocs=(\d+) frees=(\d+) reallocs=(\d+) total_bytes=(\d+))");
+  std::smatch match;
+  ASSERT_TRUE(std::regex_search(result.stderr_str, match, stats_regex))
+      << "Could not parse interposition stats\nstderr: " << result.stderr_str;
+
+  int allocs = std::stoi(match[1].str());
+  int frees = std::stoi(match[2].str());
+  int reallocs = std::stoi(match[3].str());
+  int total_bytes = std::stoi(match[4].str());
+
+  // The test helper does: 3 allocs, 1 realloc, 3 frees
+  // The realloc might internally do an alloc, so we check >= expected
+  EXPECT_GE(allocs, 3) << "Expected at least 3 allocations";
+  EXPECT_GE(frees, 3) << "Expected at least 3 frees";
+  EXPECT_GE(reallocs, 1) << "Expected at least 1 realloc";
+  // Total allocated: 1024 + 2048 + 512 = 3584 (not counting realloc growth)
+  EXPECT_GE(total_bytes, 3584) << "Expected at least 3584 bytes allocated";
+}
+
+#else  // !ARROW_MIMALLOC || _WIN32
+
+TEST(TestMimallocInterpose, NotAvailable) {
+  GTEST_SKIP() << "Mimalloc interposition test requires ARROW_MIMALLOC and non-Windows";
+}
+
+#endif
+
+}  // namespace arrow