diff --git a/core/src/main/java/org/apache/calcite/materialize/MaterializationService.java b/core/src/main/java/org/apache/calcite/materialize/MaterializationService.java
index fc9ac2b28ffd..bae0d51164db 100644
--- a/core/src/main/java/org/apache/calcite/materialize/MaterializationService.java
+++ b/core/src/main/java/org/apache/calcite/materialize/MaterializationService.java
@@ -334,6 +334,9 @@ public static void setThreadLocal() {
     THREAD_INSTANCE.set(new MaterializationService());
   }
 
+  public static void setThreadLocal(MaterializationService service) {
+    THREAD_INSTANCE.set(service);
+  }
   /** Returns the instance of the materialization service. Usually the global
    * one, but returns a thread-local one during testing (when
    * {@link #setThreadLocal()} has been called by the current thread). */