Register dataframe as table to allow further sql queries over df (#63)

* register dataframe as table to allow further sql queries over df

* indent

Author: yjshen

datafusion-examples/src/main/java/org/apache/arrow/datafusion/examples/ExampleMain.java

@@ -33,7 +33,7 @@ public static void main(String[] args) throws Exception {
           .registerParquet(
               "test_parquet", Paths.get("src/main/resources/aggregate_test_100.parquet"))
           .join();
-      context.sql("select * from test_parquet limit 3").thenComposeAsync(DataFrame::show).join();
+      context.sql("select * from test_parquet limit 5").thenComposeAsync(DataFrame::show).join();
 
       context
           .sql("select * from test_csv")
@@ -52,6 +52,13 @@ public static void main(String[] args) throws Exception {
           .sql("select * from test_parquet limit 3")
           .thenComposeAsync(df -> df.writeParquet(tempPath.resolve("parquet-out")))
           .join();
+
+      context
+          .sql("select * from test_parquet limit 3")
+          .thenComposeAsync(df -> df.registerTable(context, "test_parquet_limited"))
+          .join();
+
+      context.sql("select * from test_parquet_limited").thenComposeAsync(DataFrame::show).join();
     }
   }
 

datafusion-java/src/main/java/org/apache/arrow/datafusion/DataFrame.java

@@ -43,4 +43,12 @@ public interface DataFrame extends NativeProxy {
    * @return null
    */
   CompletableFuture<Void> writeCsv(Path path);
+
+  /**
+   * Register this dataframe as a temporary table.
+   * @param context SessionContext to register table to
+   * @param name name of the tmp table
+   * @return null
+   */
+  CompletableFuture<Void> registerTable(SessionContext context, String name);
 }

datafusion-java/src/main/java/org/apache/arrow/datafusion/DataFrames.java

@@ -19,4 +19,7 @@ static native void writeParquet(
       long runtime, long dataframe, String path, Consumer<String> callback);
 
   static native void writeCsv(long runtime, long dataframe, String path, Consumer<String> callback);
+
+  static native void registerTable(
+      long runtime, long dataframe, long context, String name, Consumer<String> callback);
 }

datafusion-java/src/main/java/org/apache/arrow/datafusion/DefaultDataFrame.java

@@ -104,6 +104,27 @@ public CompletableFuture<Void> writeCsv(Path path) {
     return future;
   }
 
+  public CompletableFuture<Void> registerTable(SessionContext ctx, String name) {
+      Runtime runtime = context.getRuntime();
+      long runtimePointer = runtime.getPointer();
+      long dataframe = getPointer();
+      long contextPointer = ctx.getPointer();
+      CompletableFuture<Void> future = new CompletableFuture<>();
+      DataFrames.registerTable(
+          runtimePointer,
+          dataframe,
+          contextPointer,
+          name,
+          (String errString) -> {
+            if (containsError(errString)) {
+              future.completeExceptionally(new RuntimeException(errString));
+            } else {
+              future.complete(null);
+            }
+          });
+      return future;
+  }
+
   @Override
   void doClose(long pointer) {
     DataFrames.destroyDataFrame(pointer);

datafusion-jni/src/dataframe.rs

@@ -7,6 +7,7 @@ use std::convert::Into;
 use std::io::BufWriter;
 use std::io::Cursor;
 use std::sync::Arc;
+use datafusion::prelude::SessionContext;
 use tokio::runtime::Runtime;
 
 #[no_mangle]
@@ -153,6 +154,44 @@ pub extern "system" fn Java_org_apache_arrow_datafusion_DataFrames_writeCsv(
     });
 }
 
+#[no_mangle]
+pub extern "system" fn Java_org_apache_arrow_datafusion_DataFrames_registerTable(
+    env: JNIEnv,
+    _class: JClass,
+    runtime: jlong,
+    dataframe: jlong,
+    session: jlong,
+    name: JString,
+    callback: JObject,
+) {
+    let runtime = unsafe { &mut *(runtime as *mut Runtime) };
+    let dataframe = unsafe { &mut *(dataframe as *mut Arc<DataFrame>) };
+    let context = unsafe { &mut *(session as *mut SessionContext) };
+    let name: String = env
+        .get_string(name)
+        .expect("Couldn't get name as string!")
+        .into();
+    runtime.block_on(async {
+        let r = context.register_table(name.as_str(), dataframe.clone());
+        let err_message: JValue = match r {
+            Ok(_) => JValue::Void,
+            Err(err) => {
+                let err_message = env
+                    .new_string(err.to_string())
+                    .expect("Couldn't create java string!");
+                err_message.into()
+            }
+        };
+        env.call_method(
+            callback,
+            "accept",
+            "(Ljava/lang/Object;)V",
+            &[err_message.into()],
+        )
+        .expect("failed to call method");
+    });
+}
+
 #[no_mangle]
 pub extern "system" fn Java_org_apache_arrow_datafusion_DataFrames_destroyDataFrame(
     _env: JNIEnv,