Add IR dumping to JNI.

dremio · Mar 30, 2024 · e4d2bd3 · e4d2bd3
1 parent 9f0101e
commit e4d2bd3
Show file tree

Hide file tree

Showing 6 changed files with 62 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
@@ -697,6 +697,18 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_build
   return module_id;
 }
 
+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());
+}
+
 ///
 /// \brief Resizable buffer which resizes by doing a callback into java.
 ///
@@ -977,6 +989,18 @@ JNIEXPORT jlong JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_build
   return module_id;
 }
 
+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 jint JNICALL Java_org_apache_arrow_gandiva_evaluator_JniWrapper_evaluateFilter(
     JNIEnv* env, jobject cls, jlong module_id, jint num_rows, jlongArray buf_addrs,
     jlongArray buf_sizes, jint jselection_vector_type, jlong out_buf_addr,

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();
@@ -56,6 +57,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,13 @@ 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, long moduleId, int numRows,
                                 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.
    *
@@ -111,6 +118,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
@@ -350,6 +350,13 @@ private void evaluate(int numRows, List<ArrowBuf> buffers, List<ArrowBuffer> buf
         outAddrs, outSizes);
   }
 
+  /**
+   * 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.
    */