[FEA] Add JIT cache management functions

cpp/include/cudf/context.hpp

@@ -78,4 +78,28 @@ void initialize(init_flags flags = init_flags::INIT_JIT_CACHE);
 /// teardown and that only one thread calls teardown at a time.
 void teardown();
 
+/**
+ * @brief Enable or disable the JIT program cache
+ *
+ * When disabled, the cache will not be used for
+ * storing or retrieving compiled programs, effectively bypassing the cache. When enabled, the
+ * cache will be used as normal. This can be used to temporarily disable caching without clearing
+ * the existing cache contents, allowing for easy re-enabling of the cache later.
+ *
+ * @param enable If `true`, the JIT program cache is enabled; if `false`, it is disabled.
+ */
+void enable_jit_cache(bool enable);
+
+/**
+ * @brief Clear the JIT program cache, removing all cached programs from memory and disk.
+ *
+ * This is a more expensive operation than simply disabling the cache, as it involves deleting
+ * cached files from disk, but it also frees up any memory used by the cached programs. Use
+ * `enable_jit_cache(false)` if you want to temporarily disable caching without clearing existing
+ * cache contents.
+ *
+ * @warning For benchmarking or testing purposes, prefer `enable_jit_cache`.
+ */
+void clear_jit_cache();
+
 }  // namespace CUDF_EXPORT cudf

cpp/src/jit/cache.cpp

@@ -1,8 +1,9 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include "io/utilities/getenv_or.hpp"
 #include "runtime/context.hpp"
 
 #include <cudf/context.hpp>
@@ -95,40 +96,63 @@ std::string get_program_cache_dir()
 #endif
 }
 
-std::size_t try_parse_numeric_env_var(char const* const env_name, std::size_t default_val)
-{
-  auto const value = std::getenv(env_name);
-  return value != nullptr ? std::stoull(value) : default_val;
-}
 }  // namespace
 
 jitify2::ProgramCache<>& jit::program_cache::get(jitify2::PreprocessedProgramData const& preprog)
 {
   CUDF_FUNC_RANGE();
-  std::lock_guard<std::mutex> const caches_lock(_caches_mutex);
+  std::lock_guard caches_lock(_caches_mutex);
 
   auto existing_cache = _caches.find(preprog.name());
 
-  if (existing_cache == _caches.end()) {
-    auto const kernel_limit_proc =
-      try_parse_numeric_env_var("LIBCUDF_KERNEL_CACHE_LIMIT_PER_PROCESS", 10'000);
-    auto const kernel_limit_disk =
-      try_parse_numeric_env_var("LIBCUDF_KERNEL_CACHE_LIMIT_DISK", 100'000);
-
-    // if kernel_limit_disk is zero, jitify will assign it the value of kernel_limit_proc.
-    // to avoid this, we treat zero as "disable disk caching" by not providing the cache dir.
-    auto const cache_dir = kernel_limit_disk == 0 ? std::string{} : get_program_cache_dir();
-
-    auto const res =
-      _caches.insert({preprog.name(),
+  if (existing_cache == _caches.end() || _disabled.load(std::memory_order_seq_cst)) {
+    auto res =
+      _caches.emplace(preprog.name(),
                       std::make_unique<jitify2::ProgramCache<>>(
-                        kernel_limit_proc, preprog, nullptr, cache_dir, kernel_limit_disk)});
+                        _kernel_limit_proc, preprog, nullptr, _cache_dir, _kernel_limit_disk));
     existing_cache = res.first;
   }
 
   return *(existing_cache->second);
 }
 
+void jit::program_cache::clear()
+{
+  CUDF_FUNC_RANGE();
+  std::lock_guard caches_lock(_caches_mutex);
+
+  _caches.clear();
+
+  // non-atomic
+  std::filesystem::remove_all(_cache_dir);
+}
+
+void jit::program_cache::enable(bool enable)
+{
+  _disabled.store(!enable, std::memory_order_seq_cst);
+}
+
+bool jit::program_cache::is_enabled() const { return !_disabled.load(std::memory_order_seq_cst); }
+
+std::unique_ptr<jit::program_cache> jit::program_cache::create()
+{
+  auto const kernel_limit_proc = getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_PER_PROCESS", 10'000);
+  auto const kernel_limit_disk = getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_DISK", 100'000);
+  auto const disabled          = get_bool_env_or("LIBCUDF_KERNEL_CACHE_DISABLED", false);
+  auto const clear_cache       = get_bool_env_or("LIBCUDF_KERNEL_CACHE_CLEAR", false);
+
+  // if kernel_limit_disk is zero, jitify will assign it the value of kernel_limit_proc.
+  // to avoid this, we treat zero as "disable disk caching" by not providing the cache dir.
+  auto cache_dir = kernel_limit_disk == 0 ? std::string{} : get_program_cache_dir();
+
+  auto cache =
+    std::make_unique<jit::program_cache>(kernel_limit_proc, kernel_limit_disk, cache_dir, disabled);
+
+  if (clear_cache) { cache->clear(); }
+
+  return cache;
+}
+
 jitify2::ProgramCache<>& jit::get_program_cache(jitify2::PreprocessedProgramData const& preprog)
 {
   return cudf::get_context().program_cache().get(preprog);

cpp/src/jit/cache.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -11,6 +11,8 @@
 
 #include <jitify2.hpp>
 
+#include <atomic>
+#include <filesystem>
 #include <memory>
 #include <mutex>
 #include <string>
@@ -21,16 +23,38 @@ namespace jit {
 class program_cache {
   std::mutex _caches_mutex;
   std::unordered_map<std::string, std::unique_ptr<jitify2::ProgramCache<>>> _caches;
+  int32_t _kernel_limit_proc;
+  int32_t _kernel_limit_disk;
+  std::filesystem::path _cache_dir;
+  std::atomic<bool> _disabled;
 
  public:
-  program_cache()                                = default;
+  program_cache(int32_t kernel_limit_proc,
+                int32_t kernel_limit_disk,
+                std::filesystem::path cache_dir,
+                bool disabled)
+    : _kernel_limit_proc{kernel_limit_proc},
+      _kernel_limit_disk{kernel_limit_disk},
+      _cache_dir{std::move(cache_dir)},
+      _disabled{disabled}
+  {
+  }
+
   program_cache(program_cache const&)            = delete;
   program_cache(program_cache&&)                 = delete;
   program_cache& operator=(program_cache const&) = delete;
   program_cache& operator=(program_cache&&)      = delete;
   ~program_cache()                               = default;
 
   jitify2::ProgramCache<>& get(jitify2::PreprocessedProgramData const& preprog);
+
+  void clear();
+
+  void enable(bool enable);
+
+  bool is_enabled() const;
+
+  static std::unique_ptr<jit::program_cache> create();
 };
 
 jitify2::ProgramCache<>& get_program_cache(jitify2::PreprocessedProgramData const& preprog);

cpp/src/runtime/context.cpp

@@ -27,7 +27,7 @@ void context::ensure_nvcomp_loaded() { io::detail::nvcomp::load_nvcomp_library()
 void context::ensure_jit_cache_initialized()
 {
   std::call_once(_program_cache_init_flag,
-                 [&]() { _program_cache = std::make_unique<jit::program_cache>(); });
+                 [&]() { _program_cache = jit::program_cache::create(); });
 }
 
 jit::program_cache& context::program_cache()
@@ -85,6 +85,10 @@ void teardown()
   });
 }
 
+void enable_jit_cache(bool enable) { get_context().program_cache().enable(enable); }
+
+void clear_jit_cache() { get_context().program_cache().clear(); }
+
 context& get_context()
 {
   cudf::initialize();

cpp/src/transform/compute_column.cu

@@ -1,9 +1,10 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include "compute_column_kernel.hpp"
+#include "runtime/context.hpp"
 
 #include <cudf/ast/detail/expression_evaluator.cuh>
 #include <cudf/ast/detail/expression_parser.hpp>
@@ -31,6 +32,8 @@ std::unique_ptr<column> compute_column(table_view const& table,
                                        rmm::cuda_stream_view stream,
                                        rmm::device_async_resource_ref mr)
 {
+  if (get_context().use_jit()) { return compute_column_jit(table, expr, stream, mr); }
+
   // If evaluating the expression may produce null outputs we create a nullable
   // output column and follow the null-supporting expression evaluation code
   // path.

cpp/src/transform/transform.cu

@@ -148,6 +148,7 @@ auto to_device_input_arg(InputsView inputs,
                          rmm::device_async_resource_ref mr)
 {
   std::vector<column_view> columns;
+
   for (auto const& input : inputs) {
     columns.emplace_back(std::visit([](auto const& col) { return to_column_view(col); }, input));
   }