Add dump IR capabilities.

dremio · Mar 30, 2024 · 61cb083 · 61cb083
1 parent f2aaf08
commit 61cb083
Show file tree

Hide file tree

Showing 6 changed files with 63 additions and 3 deletions.
diff --git a/java/gandiva/src/main/cpp/config_builder.cc b/java/gandiva/src/main/cpp/config_builder.cc
@@ -34,11 +34,12 @@ using gandiva::ConfigurationBuilder;
  */
 JNIEXPORT jlong JNICALL
 Java_org_apache_arrow_gandiva_evaluator_ConfigurationBuilder_buildConfigInstance(
-    JNIEnv* env, jobject configuration, jboolean optimize, jboolean target_host_cpu) {
+    JNIEnv* env, jobject configuration, jboolean optimize, jboolean target_host_cpu, jboolean dump_ir) {
   ConfigurationBuilder configuration_builder;
   std::shared_ptr<Configuration> config = configuration_builder.build();
   config->set_optimize(optimize);
   config->target_host_cpu(target_host_cpu);
+  config->set_dump_ir(dump_ir);
   return ConfigHolder::MapInsert(config);
 }
 

diff --git a/java/gandiva/src/main/cpp/jni_common.cc b/java/gandiva/src/main/cpp/jni_common.cc
@@ -1064,6 +1064,18 @@ Java_org_apache_arrow_gandiva_evaluator_JniWrapper_evaluateProjector(
   }
 }
 
+JNIEXPORT jstring JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_dumpProjectorIr(
+    JNIEnv* env, jobject cls, jlong module_id) {
+  std::shared_ptr<ProjectorHolder> holder = projector_modules_.Lookup(module_id);
+  if (holder == nullptr) {
+    env->ThrowNew(gandiva_exception_, "Unknown module id\n");
+    return env->NewStringUTF(nullptr);
+  }
+
+  std::string ir = holder->projector()->DumpIR();
+  return env->NewStringUTF(ir.c_str());
+}
+
 JNIEXPORT void JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_closeProjector(
     JNIEnv* env, jobject cls, jlong module_id) {
   projector_modules_.Erase(module_id);
@@ -1227,6 +1239,18 @@ JNIEXPORT jint JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_evalua
   }
 }
 
+JNIEXPORT jstring JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_dumpFilterIr(
+    JNIEnv* env, jobject cls, jlong module_id) {
+  std::shared_ptr<FilterHolder> holder = filter_modules_.Lookup(module_id);
+  if (holder == nullptr) {
+    env->ThrowNew(gandiva_exception_, "Unknown module id\n");
+    return env->NewStringUTF(nullptr);
+  }
+
+  std::string ir = holder->filter()->DumpIR();
+  return env->NewStringUTF(ir.c_str());
+}
+
 JNIEXPORT void JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_closeFilter(
     JNIEnv* env, jobject cls, jlong module_id) {
   filter_modules_.Erase(module_id);

diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/ConfigurationBuilder.java
@@ -25,10 +25,10 @@
 public class ConfigurationBuilder {
 
   public long buildConfigInstance(ConfigOptions configOptions) {
-    return buildConfigInstance(configOptions.optimize, configOptions.targetCPU);
+    return buildConfigInstance(configOptions.optimize, configOptions.targetCPU, configOptions.dumpIR);
   }
 
-  private native long buildConfigInstance(boolean optimize, boolean detectHostCPU);
+  private native long buildConfigInstance(boolean optimize, boolean detectHostCPU, boolean dumpIR);
 
   public native void releaseConfigInstance(long configId);
 
@@ -38,6 +38,7 @@ public long buildConfigInstance(ConfigOptions configOptions) {
   public static class ConfigOptions {
     private boolean optimize = true;
     private boolean targetCPU = true;
+    private boolean dumpIR = false;
 
     public static ConfigOptions getDefault() {
       return new ConfigOptions();
@@ -55,6 +56,11 @@ public ConfigOptions withTargetCPU(boolean targetCPU) {
       return this;
     }
 
+    public ConfigOptions withDumpIR(boolean dumpIR) {
+      this.dumpIR = dumpIR;
+      return this;
+    }
+
     @Override
     public int hashCode() {
       return Objects.hash(optimize, targetCPU);

diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Filter.java
@@ -185,6 +185,14 @@ private void evaluate(int numRows, List<ArrowBuf> buffers, List<ArrowBuffer> buf
     }
   }
 
+  /**
+   * Dumps the IR for the LLVM module representing this filter.
+   */
+  public String dumpIR() throws GandivaException {
+    return wrapper.dumpFilterIr(this.moduleId);
+  }
+
+
   /**
    * Closes the LLVM module representing this filter.
    */

diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/JniWrapper.java
@@ -68,6 +68,13 @@ native void evaluateProjector(Object expander, Object listExpander, long moduleI
                                 int selectionVectorType, int selectionVectorSize,
                                 long selectionVectorBufferAddr, long selectionVectorBufferSize,
                                 long[] outAddrs, long[] outSizes) throws GandivaException;
+
+  /**
+   * Dumps the IR of the projector referenced by moduleId.
+   *
+   * @param moduleId moduleId to dump the IR
+   */
+  native String dumpProjectorIr(long moduleId);
 
   /**
    * Closes the projector referenced by moduleId.
@@ -112,6 +119,13 @@ native int evaluateFilter(long moduleId, int numRows, long[] bufAddrs, long[] bu
                             int selectionVectorType,
                             long outAddr, long outSize) throws GandivaException;
 
+  /**
+   * Dumps the IR of the filter referenced by moduleId.
+   *
+   * @param moduleId moduleId to dump the IR
+   */
+  native String dumpFilterIr(long moduleId);
+
   /**
    * Closes the filter referenced by moduleId.
    *

diff --git a/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java b/java/gandiva/src/main/java/org/apache/arrow/gandiva/evaluator/Projector.java
@@ -378,6 +378,13 @@ private void evaluate(int numRows, List<ArrowBuf> buffers, List<ArrowBuffer> buf
     }
   }
 
+  /**
+   * Dumps the IR for the LLVM module representing this projector.
+   */
+  public String dumpIR() throws GandivaException {
+    return wrapper.dumpProjectorIr(this.moduleId);
+  }
+
   /**
    * Closes the LLVM module representing this evaluator.
    */