Tying searcher and reader with rust

Co-authored-by: Arpit Bandejiya <[email protected]>
Signed-off-by: bharath-techie <[email protected]>

Author: bharath-techie

libs/vectorized-exec-spi/src/main/java/org/opensearch/vectorized/execution/search/DataFormat.java

@@ -8,6 +8,9 @@
 
 package org.opensearch.vectorized.execution.search;
 
+import org.opensearch.common.annotation.ExperimentalApi;
+
+@ExperimentalApi
 public enum DataFormat {
     CSV,
     Text

plugins/dataformat-csv/src/main/java/org/opensearch/datafusion/csv/CsvDataFormatPlugin.java

@@ -48,7 +48,16 @@ public Optional<Map<org.opensearch.vectorized.execution.search.DataFormat, DataS
 
     @Override
     public <T extends DataFormat> IndexingExecutionEngine<T> indexingEngine() {
-        return (IndexingExecutionEngine<T>) new CsvEngine();
+        if (CsvDataFormat.class.equals(getDataFormatType())) {
+            @SuppressWarnings("unchecked")
+            IndexingExecutionEngine<T> engine = (IndexingExecutionEngine<T>) new CsvEngine();
+            return engine;
+        }
+        throw new IllegalArgumentException("Unsupported data format type: " + getDataFormatType());
+    }
+
+    private Class<? extends DataFormat> getDataFormatType() {
+        return CsvDataFormat.class;
     }
 
     @Override

plugins/engine-datafusion/jni/Cargo.toml

@@ -47,7 +47,7 @@ log = "0.4"
 parquet = "53.0.0"
 
 # Object store for file access
-object_store = "0.11"
+object_store = "=0.12.3"
 url = "2.0"
 
 # Substrait support
@@ -56,6 +56,7 @@ prost = "0.13"
 
 # Temporary directory support
 tempfile = "3.0"
+chrono = "0.4.41"
 
 [build-dependencies]
 cbindgen = "0.27"

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

@@ -6,18 +6,33 @@
  * compatible open source license.
  */
 
-use jni::objects::JClass;
-use jni::sys::{jlong, jstring};
+use jni::objects::{JByteArray, JClass};
+use jni::sys::{jbyteArray, jlong, jstring};
 use jni::JNIEnv;
 use std::sync::Arc;
 
+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::listing::{ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl};
+use datafusion::execution::cache::cache_unit::DefaultListFilesCache;
+use datafusion::execution::cache::CacheAccessor;
+use datafusion::execution::SendableRecordBatchStream;
+use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
+use datafusion_substrait::substrait::proto::Plan;
+use jni::objects::{JObjectArray, JString};
+use prost::Message;
+use tokio::runtime::Runtime;
+use object_store::ObjectMeta;
 
 /// Create a new DataFusion session context
 #[no_mangle]
@@ -110,6 +125,232 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_closeSe
 }
 
 
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_createReader(
+    mut env: JNIEnv,
+    _class: JClass,
+    table_path: JString,
+    files: JObjectArray
+) -> jlong {
+
+    let table_path: String = env.get_string(&table_path).expect("Couldn't get java string!").into();
+    let files: Vec<String> = parse_string_arr(&mut env, files).expect("Expected list of files");
+    let files_meta = create_object_meta_from_filenames(&table_path, files);
+
+    let table_path = ListingTableUrl::parse(table_path).unwrap();
+    let shard_view = ShardView::new(table_path, files_meta);
+    Box::into_raw(Box::new(shard_view)) as jlong
+}
+
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_destroyReader(
+    mut env: JNIEnv,
+    _class: JClass,
+    ptr: jlong
+)  {
+    let _ = unsafe { Box::from_raw(ptr as *mut ShardView) };
+}
+
+pub struct ShardView {
+    table_path: ListingTableUrl,
+    files_meta: Arc<Vec<ObjectMeta>>
+}
+
+impl ShardView {
+    pub fn new(table_path: ListingTableUrl, files_meta: Vec<ObjectMeta>) -> Self {
+        let files_meta = Arc::new(files_meta);
+        ShardView {
+            table_path,
+            files_meta
+        }
+    }
+
+    pub fn table_path(&self) -> ListingTableUrl {
+        self.table_path.clone()
+    }
+
+    pub fn files_meta(&self) -> Arc<Vec<ObjectMeta>> {
+        self.files_meta.clone()
+    }
+}
+
+
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeExecuteSubstraitQuery(
+    mut env: JNIEnv,
+    _class: JClass,
+    shard_view_ptr: jlong,
+    substrait_bytes: jbyteArray,
+    // callback: JObject,
+) -> jlong {
+    let shard_view = unsafe { &*(shard_view_ptr as *const ShardView) };
+    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);
+
+    let runtime_env = RuntimeEnvBuilder::new()
+        .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 listing_options = ListingOptions::new(Arc::new(file_format))
+        .with_file_extension(".csv");
+
+    // Ideally the executor will give this
+    Runtime::new().expect("Failed to create Tokio Runtime").block_on(async {
+        let resolved_schema = listing_options
+            .infer_schema(&ctx.state(), &table_path.clone())
+            .await.unwrap();
+
+
+        let config = ListingTableConfig::new(table_path.clone())
+            .with_listing_options(listing_options)
+            .with_schema(resolved_schema);
+
+        // Create a new TableProvider
+        let provider = Arc::new(ListingTable::try_new(config).unwrap());
+        let shard_id = table_path.prefix().filename().expect("error in fetching Path");
+        ctx.register_table(shard_id, provider)
+            .expect("Failed to attach the Table");
+
+    });
+
+    // TODO : how to close ctx ?
+    // Convert Java byte array to Rust Vec<u8>
+    let plan_bytes_obj = unsafe { JByteArray::from_raw(substrait_bytes) };
+    let plan_bytes_vec = match env.convert_byte_array(plan_bytes_obj) {
+        Ok(bytes) => bytes,
+        Err(e) => {
+            let error_msg = format!("Failed to convert plan bytes: {}", e);
+            env.throw_new("java/lang/Exception", error_msg);
+            return 0;
+        }
+    };
+
+    let substrait_plan = match Plan::decode(plan_bytes_vec.as_slice()) {
+        Ok(plan) => {
+            println!("SUBSTRAIT rust: Decoding is successful, Plan has {} relations", plan.relations.len());
+            plan
+        },
+        Err(e) => {
+            return 0;
+        }
+    };
+
+    //let runtime = unsafe { &mut *(runtime_ptr as *mut Runtime) };
+    Runtime::new().expect("Failed to create Tokio Runtime").block_on(async {
+
+        let logical_plan = match from_substrait_plan(&ctx.state(), &substrait_plan).await {
+            Ok(plan) => {
+                println!("SUBSTRAIT Rust: LogicalPlan: {:?}", plan);
+                plan
+            },
+            Err(e) => {
+                println!("SUBSTRAIT Rust: Failed to convert Substrait plan: {}", e);
+                return;
+            }
+        };
+
+        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
+
+
+}
+
+// If we need to create session context separately
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_nativeCreateSessionContext(
+    mut env: JNIEnv,
+    _class: JClass,
+    runtime_ptr: jlong,
+    shard_view_ptr: jlong,
+    global_runtime_env_ptr: jlong,
+) -> jlong {
+    let shard_view = unsafe { &*(shard_view_ptr as *const ShardView) };
+    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);
+
+    let runtime_env = RuntimeEnvBuilder::new()
+        .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 listing_options = ListingOptions::new(Arc::new(file_format))
+        .with_file_extension(".csv");
+
+
+    // let runtime = unsafe { &mut *(runtime_ptr as *mut Runtime) };
+    let mut session_context_ptr = 0;
+
+    // Ideally the executor will give this
+    Runtime::new().expect("Failed to create Tokio Runtime").block_on(async {
+        let resolved_schema = listing_options
+            .infer_schema(&ctx.state(), &table_path.clone())
+            .await.unwrap();
+
+
+        let config = ListingTableConfig::new(table_path.clone())
+            .with_listing_options(listing_options)
+            .with_schema(resolved_schema);
+
+        // Create a new TableProvider
+        let provider = Arc::new(ListingTable::try_new(config).unwrap());
+        let shard_id = table_path.prefix().filename().expect("error in fetching Path");
+        ctx.register_table(shard_id, provider)
+            .expect("Failed to attach the Table");
+
+        // Return back after wrapping in Box
+        session_context_ptr = Box::into_raw(Box::new(ctx)) as jlong
+    });
+
+    session_context_ptr
+}
+