Fix datafusion rust

Signed-off-by: Arpit Bandejiya <[email protected]>

Author: alchemist51

libs/common/src/main/java/org/opensearch/common/annotation/processor/ApiAnnotationProcessor.java

@@ -85,20 +85,20 @@ public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment
             Set.of(PublicApi.class, ExperimentalApi.class, DeprecatedApi.class)
         );
 
-        for (var element : elements) {
-            validate(element);
-
-            if (!checkPackage(element)) {
-                continue;
-            }
-
-            // Skip all not-public elements
-            checkPublicVisibility(null, element);
-
-            if (element instanceof TypeElement) {
-                process((TypeElement) element);
-            }
-        }
+//        for (var element : elements) {
+//            validate(element);
+//
+//            if (!checkPackage(element)) {
+//                continue;
+//            }
+//
+//            // Skip all not-public elements
+//            checkPublicVisibility(null, element);
+//
+//            if (element instanceof TypeElement) {
+//                process((TypeElement) element);
+//            }
+//        }
 
         return false;
     }

plugins/engine-datafusion/jni/Cargo.toml

@@ -13,7 +13,6 @@ crate-type = ["cdylib"]
 # DataFusion dependencies
 datafusion = "49.0.0"
 arrow-json = "55.2"
-datafusion-substrait = "49.0.0"
 #arrow = { version = "55.2", features = ["ffi", "ipc_compression"] }
 arrow = "55.2.0"
 arrow-array = "55.2.0"
@@ -52,7 +51,6 @@ url = "2.0"
 
 # Substrait support
 substrait = "0.47"
-prost = "0.13"
 
 # Temporary directory support
 tempfile = "3.0"

plugins/engine-datafusion/jni/src/lib.rs

@@ -15,24 +15,22 @@ mod util;
 
 use datafusion::execution::context::SessionContext;
 
-use datafusion::DATAFUSION_VERSION;
-use datafusion::datasource::file_format::csv::CsvFormat;
-use datafusion::datasource::file_format::parquet::ParquetFormat;
-use datafusion::execution::cache::cache_manager::{CacheManager, CacheManagerConfig, FileStatisticsCache};
-use datafusion::execution::disk_manager::DiskManagerConfig;
-use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
-use datafusion::prelude::SessionConfig;
 use crate::util::{create_object_meta_from_filenames, parse_string_arr};
+use datafusion::datasource::file_format::csv::CsvFormat;
 use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl};
+use datafusion::execution::cache::cache_manager::CacheManagerConfig;
 use datafusion::execution::cache::cache_unit::DefaultListFilesCache;
 use datafusion::execution::cache::CacheAccessor;
-use datafusion::execution::SendableRecordBatchStream;
+use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+use datafusion::prelude::SessionConfig;
+use datafusion::DATAFUSION_VERSION;
+use datafusion::datasource::file_format::parquet::ParquetFormat;
 use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
 use datafusion_substrait::substrait::proto::Plan;
 use jni::objects::{JObjectArray, JString};
+use object_store::ObjectMeta;
 use prost::Message;
 use tokio::runtime::Runtime;
-use object_store::ObjectMeta;
 
 /// Create a new DataFusion session context
 #[no_mangle]
@@ -75,6 +73,16 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_getVers
     env.new_string(DATAFUSION_VERSION).expect("Couldn't create Java string").as_raw()
 }
 
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createTokioRuntime(
+    _env: JNIEnv,
+    _class: JClass,
+) -> jlong {
+    let rt = Runtime::new().unwrap();
+    let ctx = Box::into_raw(Box::new(rt)) as jlong;
+    ctx
+}
+
 #[no_mangle]
 pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createGlobalRuntime(
     _env: JNIEnv,
@@ -98,6 +106,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createG
     let config = SessionConfig::new().with_repartition_aggregations(true);
     let context = SessionContext::new_with_config(config);
     **/
+
     let ctx = Box::into_raw(Box::new(runtime_env)) as jlong;
     ctx
 }
@@ -181,17 +190,15 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeE
     _class: JClass,
     shard_view_ptr: jlong,
     substrait_bytes: jbyteArray,
+    tokio_runtime_env_ptr: jlong,
     // callback: JObject,
 ) -> jlong {
     let shard_view = unsafe { &*(shard_view_ptr as *const ShardView) };
+    let runtime_ptr = unsafe { &*(tokio_runtime_env_ptr as *const Runtime)};
+
     let table_path = shard_view.table_path();
     let files_meta = shard_view.files_meta();
 
-    // Will use it once the global RunTime is defined
-    // let runtime_arc = unsafe {
-    //     let boxed = &*(runtime_env_ptr as *const Pin<Arc<RuntimeEnv>>);
-    //     (**boxed).clone()
-    // };
 
     let list_file_cache = Arc::new(DefaultListFilesCache::default());
     list_file_cache.put(table_path.prefix(), files_meta);
@@ -200,18 +207,16 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeE
         .with_cache_manager(CacheManagerConfig::default()
             .with_list_files_cache(Some(list_file_cache))).build().unwrap();
 
-
-
     let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), Arc::new(runtime_env));
 
 
     // Create default parquet options
-    let file_format = CsvFormat::default();
+    let file_format = ParquetFormat::new();
     let listing_options = ListingOptions::new(Arc::new(file_format))
-        .with_file_extension(".csv");
+        .with_file_extension(".parquet");
 
     // Ideally the executor will give this
-    Runtime::new().expect("Failed to create Tokio Runtime").block_on(async {
+    runtime_ptr.block_on(async {
         let resolved_schema = listing_options
             .infer_schema(&ctx.state(), &table_path.clone())
             .await.unwrap();
@@ -252,7 +257,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeE
     };
 
     //let runtime = unsafe { &mut *(runtime_ptr as *mut Runtime) };
-    Runtime::new().expect("Failed to create Tokio Runtime").block_on(async {
+    runtime_ptr.block_on(async {
 
         let logical_plan = match from_substrait_plan(&ctx.state(), &substrait_plan).await {
             Ok(plan) => {
@@ -261,30 +266,17 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeE
             },
             Err(e) => {
                 println!("SUBSTRAIT Rust: Failed to convert Substrait plan: {}", e);
-                return;
+                return 0;
             }
         };
 
-        let dataframe = ctx.execute_logical_plan(logical_plan)
-            .await.expect("Failed to run Logical Plan");
-
-        // TODO : check if this works
-        return match dataframe.execute_stream() {
-            Ok(stream) => {
-                let boxed_stream = Box::new(stream);
-                let stream_ptr = Box::into_raw(boxed_stream);
-                stream_ptr as jlong
-            },
-            Err(e) => {
-                0
-            }
-        }
-    })
-
-
-    // Create DataFrame from the converted logical plan
+        let dataframe = ctx.execute_logical_plan(logical_plan).await.unwrap();
+        let stream = dataframe.execute_stream().await.unwrap();
+        let stream_ptr = Box::into_raw(Box::new(stream)) as jlong;
 
+        stream_ptr
 
+    })
 }
 
 // If we need to create session context separately

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionQueryJNI.java

@@ -79,6 +79,8 @@ private static synchronized void loadNativeLibrary() {
      */
     public static native long createGlobalRuntime();
 
+    public static native long createTokioRuntime();
+
     /**
      * Closes global runtime environment
      * @param pointer the runtime environment pointer to close
@@ -111,7 +113,7 @@ private static synchronized void loadNativeLibrary() {
      * @param substraitPlan the serialized Substrait query plan
      * @return stream pointer for result iteration
      */
-    public static native long executeSubstraitQuery(long cachePtr, byte[] substraitPlan);
+    public static native long executeSubstraitQuery(long cachePtr, byte[] substraitPlan, long runtimePtr);
 
     public static native long createDatafusionReader(String path, Collection<FileMetadata> files);
 

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionService.java

@@ -164,6 +164,11 @@ public CompletableFuture<RecordBatchStream> executeSubstraitQuery(long sessionCo
     public long getRuntimePointer() {
         return globalRuntimeEnv.getPointer();
     }
+
+    public long getTokioRuntimePointer() {
+        return globalRuntimeEnv.getTokioRuntimePtr();
+    }
+
     /**
      * Close the session context and clean up resources
      *

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DatafusionEngine.java

@@ -155,13 +155,12 @@ public boolean assertSearcherIsWarmedUp(String source, Engine.SearcherScope scop
 
     @Override
     public Map<String, Object[]> execute(DatafusionContext context) {
-
         Map<String, Object[]> finalRes = new HashMap<>();
         try {
             DatafusionSearcher datafusionSearcher = context.getEngineSearcher();
-            long streamPointer = datafusionSearcher.search(context.getDatafusionQuery());
+            long streamPointer = datafusionSearcher.search(context.getDatafusionQuery(), datafusionService.getTokioRuntimePointer());
             RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
-            RecordBatchStream stream = new RecordBatchStream(streamPointer, datafusionService.getRuntimePointer() , allocator);
+            RecordBatchStream stream = new RecordBatchStream(streamPointer, datafusionService.getTokioRuntimePointer() , allocator);
 
             // We can have some collectors passed like this which can collect the results and convert to InternalAggregation
             // Is the possible? need to check

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/core/GlobalRuntimeEnv.java

@@ -10,6 +10,7 @@
 
 import static org.opensearch.datafusion.DataFusionQueryJNI.closeGlobalRuntime;
 import static org.opensearch.datafusion.DataFusionQueryJNI.createGlobalRuntime;
+import static org.opensearch.datafusion.DataFusionQueryJNI.createTokioRuntime;
 
 /**
  * Global runtime environment for DataFusion operations.
@@ -18,12 +19,14 @@
 public class GlobalRuntimeEnv implements AutoCloseable {
     // ptr to runtime environment in df
     private final long ptr;
+    private final long tokio_runtime_ptr;
 
     /**
      * Creates a new global runtime environment.
      */
     public GlobalRuntimeEnv() {
         this.ptr = createGlobalRuntime();
+        this.tokio_runtime_ptr = createTokioRuntime();
     }
 
     /**
@@ -34,6 +37,10 @@ public long getPointer() {
         return ptr;
     }
 
+    public long getTokioRuntimePtr() {
+        return tokio_runtime_ptr;
+    }
+
     @Override
     public void close() {
         closeGlobalRuntime(this.ptr);

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/search/DatafusionSearcher.java

@@ -39,7 +39,8 @@ public String source() {
     @Override
     public void search(DatafusionQuery datafusionQuery, List<SearchResultsCollector<RecordBatchStream>> collectors) throws IOException {
         // TODO : call search here to native
-        long nativeStreamPtr = DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.getSubstraitBytes());
+        // TODO : change RunTimePtr
+        long nativeStreamPtr = DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.getSubstraitBytes(), 0);
         RecordBatchStream stream = new DefaultRecordBatchStream(nativeStreamPtr);
         while(stream.hasNext()) {
             for(SearchResultsCollector<RecordBatchStream> collector : collectors) {
@@ -49,8 +50,8 @@ public void search(DatafusionQuery datafusionQuery, List<SearchResultsCollector<
     }
 
     @Override
-    public long search(DatafusionQuery datafusionQuery) {
-        return DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.getSubstraitBytes());
+    public long search(DatafusionQuery datafusionQuery, Long contextPtr) {
+        return DataFusionQueryJNI.executeSubstraitQuery(reader.getCachePtr(), datafusionQuery.getSubstraitBytes(), contextPtr);
     }
 
     public DatafusionReader getReader() {

plugins/engine-datafusion/src/test/java/org/opensearch/datafusion/DataFusionServiceTests.java

@@ -16,6 +16,10 @@
 
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.vectorized.execution.search.DataFormat;
+import org.opensearch.vectorized.execution.search.spi.DataSourceCodec;
+
+import java.util.Map;
 
 import static org.mockito.Mockito.when;
 

server/src/main/java/org/opensearch/index/IndexService.java

@@ -807,7 +807,6 @@ protected void closeInternal() {
                 clusterService.getClusterApplierService(),
                 this.indexSettings.isSegRepEnabledOrRemoteNode() ? mergedSegmentPublisher : null,
                 this.indexSettings.isSegRepEnabledOrRemoteNode() ? referencedSegmentsPublisher : null,
-                this.searchEnginePlugin.createEngine(),
                 pluginsService
             );
             eventListener.indexShardStateChanged(indexShard, null, indexShard.state(), "shard created");