Dynamic index name (#19853)

* Dynamic index name

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

* Remap paths

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

---------

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

Author: alchemist51

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

@@ -5,42 +5,47 @@
  * this file be licensed under the Apache-2.0 license or a
  * compatible open source license.
  */
-use std::ptr::addr_of_mut;
+use arrow_array::ffi::FFI_ArrowArray;
+use arrow_array::{Array, StructArray};
+use arrow_schema::ffi::FFI_ArrowSchema;
 use jni::objects::{JByteArray, JClass, JObject};
 use jni::sys::{jbyteArray, jlong, jstring};
 use jni::JNIEnv;
+use std::ptr::addr_of_mut;
 use std::sync::Arc;
-use arrow_array::{Array, StructArray};
-use arrow_array::ffi::FFI_ArrowArray;
-use arrow_schema::DataType;
-use arrow_schema::ffi::FFI_ArrowSchema;
 
 mod util;
 mod row_id_optimizer;
 mod listing_table;
+mod memory;
 
 use datafusion::execution::context::SessionContext;
 
+use crate::listing_table::{ListingOptions, ListingTable, ListingTableConfig};
 use crate::util::{create_object_meta_from_filenames, parse_string_arr, set_object_result_error, set_object_result_ok};
 use datafusion::datasource::file_format::csv::CsvFormat;
-use datafusion::datasource::listing::{ListingTableUrl};
+use datafusion::datasource::file_format::parquet::ParquetFormat;
+use datafusion::datasource::listing::ListingTableUrl;
 use datafusion::execution::cache::cache_manager::CacheManagerConfig;
 use datafusion::execution::cache::cache_unit::DefaultListFilesCache;
 use datafusion::execution::cache::CacheAccessor;
 use datafusion::execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder};
+use datafusion::physical_plan::SendableRecordBatchStream;
 use datafusion::prelude::SessionConfig;
 use datafusion::DATAFUSION_VERSION;
-use datafusion::datasource::file_format::parquet::ParquetFormat;
-use datafusion::physical_plan::SendableRecordBatchStream;
 use datafusion_substrait::logical_plan::consumer::from_substrait_plan;
 use datafusion_substrait::substrait::proto::Plan;
 use futures::TryStreamExt;
 use jni::objects::{JObjectArray, JString};
 use object_store::ObjectMeta;
 use prost::Message;
 use tokio::runtime::Runtime;
-use crate::listing_table::{ListingOptions, ListingTable, ListingTableConfig};
-use crate::row_id_optimizer::FilterRowIdOptimizer;
+use std::thread;
+use std::time::{Duration, SystemTime, UNIX_EPOCH};
+use std::fs::OpenOptions;
+use std::io::Write;
+use std::sync::atomic::{AtomicBool, Ordering};
+
 
 /// Create a new DataFusion session context
 #[no_mangle]
@@ -199,12 +204,14 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute
     mut env: JNIEnv,
     _class: JClass,
     shard_view_ptr: jlong,
+    table_name: JString,
     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_name: String = env.get_string(&table_name).expect("Couldn't get java string!").into();
 
     let table_path = shard_view.table_path();
     let files_meta = shard_view.files_meta();
@@ -254,7 +261,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute
         // 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("index-7", provider)
+        ctx.register_table(table_name, provider)
             .expect("Failed to attach the Table");
 
     });
@@ -273,7 +280,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute
 
     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());
+            // println!("SUBSTRAIT rust: Decoding is successful, Plan has {} relations", plan.relations.len());
             plan
         },
         Err(e) => {
@@ -286,7 +293,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute
 
         let logical_plan = match from_substrait_plan(&ctx.state(), &substrait_plan).await {
             Ok(plan) => {
-                println!("SUBSTRAIT Rust: LogicalPlan: {:?}", plan);
+                // println!("SUBSTRAIT Rust: LogicalPlan: {:?}", plan);
                 plan
             },
             Err(e) => {
@@ -298,7 +305,7 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute
         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;
-
+        // println!("The memory used currently right now: {:?}", jemalloc_stats::refresh_allocated());
         stream_ptr
 
     })

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

@@ -113,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, long runtimePtr);
+    public static native long executeSubstraitQuery(long cachePtr, String tableName, byte[] substraitPlan, long runtimePtr);
 
     public static native long createDatafusionReader(String path, String[] files);
 

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

@@ -80,7 +80,7 @@ public QueryPhaseExecutor<DatafusionContext> getQueryPhaseExecutor() {
     public DatafusionContext createContext(ReaderContext readerContext, ShardSearchRequest request, SearchShardTarget searchShardTarget, SearchShardTask task, BigArrays bigArrays, SearchContext originalContext) throws IOException {
         DatafusionContext datafusionContext = new DatafusionContext(readerContext, request, searchShardTarget, task, this, bigArrays, originalContext);
         // Parse source
-        datafusionContext.datafusionQuery(new DatafusionQuery(request.source().queryPlanIR(), new ArrayList<>()));
+        datafusionContext.datafusionQuery(new DatafusionQuery(request.shardId().getIndexName(), request.source().queryPlanIR(), new ArrayList<>()));
         return datafusionContext;
     }
 
@@ -191,10 +191,10 @@ public void collect(RecordBatchStream value) {
                 collector.collect(stream);
             }
 
-            logger.info("Final Results:");
-            for (Map.Entry<String, Object[]> entry : finalRes.entrySet()) {
-                logger.info("{}: {}", entry.getKey(), java.util.Arrays.toString(entry.getValue()));
-            }
+//            logger.info("Final Results:");
+//            for (Map.Entry<String, Object[]> entry : finalRes.entrySet()) {
+//                logger.info("{}: {}", entry.getKey(), java.util.Arrays.toString(entry.getValue()));
+//            }
 
         } catch (Exception exception) {
             logger.error("Failed to execute Substrait query plan", exception);

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

@@ -12,12 +12,14 @@
 import java.util.List;
 
 public class DatafusionQuery {
+    private String indexName;
     private final byte[] substraitBytes;
 
     // List of Search executors which returns a result iterator which contains row id which can be joined in datafusion
     private final List<SearchExecutor> searchExecutors;
 
-    public DatafusionQuery(byte[] substraitBytes, List<SearchExecutor> searchExecutors) {
+    public DatafusionQuery(String indexName, byte[] substraitBytes, List<SearchExecutor> searchExecutors) {
+        this.indexName = indexName;
         this.substraitBytes = substraitBytes;
         this.searchExecutors = searchExecutors;
     }
@@ -29,4 +31,8 @@ public byte[] getSubstraitBytes() {
     public List<SearchExecutor> getSearchExecutors() {
         return searchExecutors;
     }
+
+    public String getIndexName() {
+        return indexName;
+    }
 }

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

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

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

@@ -113,7 +113,7 @@ public void testQueryPhaseExecutor() throws IOException {
                 throw new RuntimeException(e);
             }
 
-            long streamPointer = datafusionSearcher.search(new DatafusionQuery(protoContent, new ArrayList<>()), service.getTokioRuntimePointer());
+            long streamPointer = datafusionSearcher.search(new DatafusionQuery("test-index",protoContent, new ArrayList<>()), service.getTokioRuntimePointer());
             RootAllocator allocator = new RootAllocator(Long.MAX_VALUE);
             RecordBatchStream stream = new RecordBatchStream(streamPointer, service.getTokioRuntimePointer() , allocator);
 

server/src/main/java/org/opensearch/action/search/TransportSearchAction.java

@@ -324,7 +324,7 @@ protected void doExecute(Task task, SearchRequest searchRequest, ActionListener<
             );
         }
         executeRequest(task, searchRequest, this::searchAsyncAction, listener);
-        logger.info("Search request received is {}", searchRequest.source());
+//        logger.info("Search request received is {}", searchRequest.source());
     }
 
     /**

server/src/main/java/org/opensearch/index/engine/exec/WriterFileSet.java

@@ -41,6 +41,20 @@ public WriterFileSet(StreamInput in) throws IOException {
         }
     }
 
+    public WriterFileSet withDirectory(String newDirectory) {
+        WriterFileSet newFileSet = new WriterFileSet(Path.of(newDirectory), this.writerGeneration);
+
+        // Extract just the filename and reconstruct with new directory
+        for (String oldFilePath : this.files) {
+            Path oldPath = Path.of(oldFilePath);
+            String fileName = oldPath.getFileName().toString();
+            String newFilePath = Path.of(newDirectory, fileName).toString();
+            newFileSet.files.add(newFilePath);
+        }
+
+        return newFileSet;
+    }
+
     /**
      * Serialize this WriterFileSet to StreamOutput
      */

server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshot.java

@@ -18,6 +18,7 @@
 import org.opensearch.index.engine.exec.FileMetadata;
 import org.opensearch.index.engine.exec.RefreshResult;
 import org.opensearch.index.engine.exec.WriterFileSet;
+import java.nio.file.Path;
 
 import java.io.IOException;
 import java.io.Serializable;
@@ -43,6 +44,31 @@ public CatalogSnapshot(RefreshResult refreshResult, long id) {
         refreshResult.getRefreshedFiles().forEach((dataFormat, writerFiles) -> dfGroupedSearchableFiles.put(dataFormat.name(), writerFiles));
     }
 
+    private CatalogSnapshot(long id, Map<String, Collection<WriterFileSet>> dfGroupedSearchableFiles) {
+        super("catalog_snapshot");
+        this.id = id;
+        this.dfGroupedSearchableFiles = dfGroupedSearchableFiles;
+    }
+
+    public CatalogSnapshot remapPaths(Path newShardDataPath) {
+        Map<String, Collection<WriterFileSet>> remappedFiles = new HashMap<>();
+
+        for (Map.Entry<String, Collection<WriterFileSet>> entry : dfGroupedSearchableFiles.entrySet()) {
+            String dataFormat = entry.getKey();
+            List<WriterFileSet> remappedFileSets = new ArrayList<>();
+
+            for (WriterFileSet fileSet : entry.getValue()) {
+                // Create new WriterFileSet with updated directory and file paths
+                WriterFileSet remappedFileSet = fileSet.withDirectory(newShardDataPath.toString());
+                remappedFileSets.add(remappedFileSet);
+            }
+
+            remappedFiles.put(dataFormat, remappedFileSets);
+        }
+
+        return new CatalogSnapshot(this.id, remappedFiles);
+    }
+
     public CatalogSnapshot(StreamInput in) throws IOException {
         super("catalog_snapshot");
         this.id = in.readLong();

server/src/main/java/org/opensearch/index/engine/exec/coord/CompositeEngine.java

@@ -71,6 +71,10 @@ public CompositeEngine(MapperService mapperService, PluginsService pluginsServic
         this.engine = new CompositeIndexingExecutionEngine(mapperService, pluginsService, shardPath, 0);
         this.compositeEngineCommitter = new LuceneCommitEngine(shardPath.getDataPath());
         this.catalogSnapshot = ((LuceneCommitEngine)this.compositeEngineCommitter).readCatalogSnapshot();
+
+        // remap the paths
+        this.catalogSnapshot = this.catalogSnapshot.remapPaths(shardPath.getDataPath());
+
         this.mergeHandler = new ParquetMergeHandler(this, this.engine, this.engine.getDataFormat());
         mergeScheduler = new MergeScheduler(this.mergeHandler, this);