[Feature/datafusion] Add SingleNodeTestCase with clickbench queries (#19953)

* Handle empty flush and ignore mapped types not defined in vsr.

Signed-off-by: Marc Handalian <[email protected]>

* add setup to run clickbench tests as SingleNodeTestCases

Signed-off-by: Marc Handalian <[email protected]>

* fix test name to match convention

Signed-off-by: Marc Handalian <[email protected]>

* fix resource loading when executing DataFusionSingleNodeTests outside of intellij

Signed-off-by: Marc Handalian <[email protected]>

* bump parquet version to match data format plugin

Signed-off-by: Marc Handalian <[email protected]>

* Add 41 clickbench queries as DSL - fix 25 to not blow up the cluster

Signed-off-by: Marc Handalian <[email protected]>

* fix build

Signed-off-by: Marc Handalian <[email protected]>

---------

Signed-off-by: Marc Handalian <[email protected]>

Author: mch2

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/fields/ParquetField.java

@@ -75,7 +75,11 @@ public final void createField(final MappedFieldType mappedFieldType,
         Objects.requireNonNull(managedVSR, "ManagedVSR cannot be null");
 
         if (mappedFieldType.isColumnar()) {
-            addToGroup(mappedFieldType, managedVSR, parseValue);
+            // TODO: support dynamic mapping update
+            // for now ignore the field
+            if (managedVSR.getVector(mappedFieldType.name()) != null) {
+                addToGroup(mappedFieldType, managedVSR, parseValue);
+            }
         }
     }
 

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/writer/ParquetWriter.java

@@ -54,6 +54,10 @@ public WriteResult addDoc(ParquetDocumentInput d) throws IOException {
     public FileInfos flush(FlushIn flushIn) throws IOException {
         String fileName = vsrManager.flush(flushIn);
         FileInfos fileInfos = new FileInfos();
+        // no data flushed
+        if (fileName == null) {
+            return fileInfos;
+        }
         WriterFileSet writerFileSet = new WriterFileSet(Path.of(fileName).getParent(), writerGeneration);
         writerFileSet.add(fileName);
         fileInfos.putWriterFileSet(PARQUET_DATA_FORMAT, writerFileSet);

plugins/engine-datafusion/Cargo.toml

@@ -40,7 +40,7 @@ thiserror = "1.0"
 # Logging
 log = "0.4"
 # Parquet support
-parquet = "53.0.0"
+parquet = "54.0.0"
 
 # Object store for file access
 object_store = "=0.12.3"

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

@@ -385,6 +385,8 @@ pub extern "system" fn Java_org_opensearch_datafusion_DataFusionQueryJNI_execute
         };
 
         let dataframe = ctx.execute_logical_plan(logical_plan).await.expect("Failed to execute logical plan");
+        let physical_plan = dataframe.clone().create_physical_plan().await.unwrap();
+        println!("Physical Plan:\n{}", datafusion::physical_plan::displayable(physical_plan.as_ref()).indent(true));
 
         let stream = match dataframe.execute_stream().await {
             Ok(stream) => { stream }

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

@@ -8,8 +8,6 @@
 
 package org.opensearch.datafusion;
 
-import org.opensearch.datafusion.search.QueryCallback;
-
 /**
  * JNI wrapper for DataFusion operations
  */

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

@@ -0,0 +1,113 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.datafusion;
+
+import com.parquet.parquetdataformat.ParquetDataFormatPlugin;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.xcontent.json.JsonXContent;
+import org.opensearch.core.xcontent.MediaTypeRegistry;
+import org.opensearch.core.xcontent.XContentParser;
+import org.opensearch.plugins.Plugin;
+import org.opensearch.search.builder.SearchSourceBuilder;
+import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.test.OpenSearchSingleNodeTestCase;
+
+import java.io.BufferedReader;
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.io.Reader;
+import java.nio.charset.StandardCharsets;
+import java.util.Collection;
+import java.util.List;
+import java.util.Locale;
+
+@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST)
+public class DataFusionSingleNodeTests extends OpenSearchSingleNodeTestCase {
+
+    private static final String INDEX_MAPPING_JSON = "clickbench_index_mapping.json";
+    private static final String DATA = "clickbench.json";
+    private final String indexName = "hits";
+
+    @Override
+    protected Collection<Class<? extends Plugin>> getPlugins() {
+        return List.of(DataFusionPlugin.class, ParquetDataFormatPlugin.class);
+    }
+
+    public void testClickBenchQueries() throws IOException {
+        String mappings = fileToString(
+            INDEX_MAPPING_JSON,
+            false
+        );
+        createIndexWithMappingSource(
+            indexName,
+            Settings.builder()
+                .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0)
+                .put("index.refresh_interval", -1)
+                .build(),
+            mappings
+        );
+        String req = fileToString(
+            DATA,
+            false
+        );
+        System.out.println(req.trim());
+        client().prepareIndex("hits").setSource(req, MediaTypeRegistry.JSON).get();
+        client().admin().indices().prepareRefresh().get();
+        client().admin().indices().prepareFlush().get();
+        client().admin().indices().prepareFlush().get();
+
+        // TODO: run in a loop
+        String sourceFile = fileToString(
+            "q25.json",
+            false
+        );
+        SearchSourceBuilder source = new SearchSourceBuilder();
+        XContentParser parser = createParser(JsonXContent.jsonXContent,
+            sourceFile);
+        source.parseXContent(parser);
+        SearchResponse response = client().prepareSearch(indexName).setSource(source).get();
+        // TODO: Match expected results...
+        System.out.println(response);
+    }
+
+    static String getResourceFilePath(String relPath) {
+        return DataFusionSingleNodeTests.class.getClassLoader().getResource(relPath).getPath();
+    }
+
+    static String fileToString(
+        final String filePathFromProjectRoot, final boolean removeNewLines) throws IOException {
+
+        final String absolutePath = getResourceFilePath(filePathFromProjectRoot);
+
+        try (final InputStream stream = new FileInputStream(absolutePath);
+             final Reader streamReader = new InputStreamReader(stream, StandardCharsets.UTF_8);
+             final BufferedReader br = new BufferedReader(streamReader)) {
+
+            final StringBuilder stringBuilder = new StringBuilder();
+            String line = br.readLine();
+
+            while (line != null) {
+
+                stringBuilder.append(line);
+                if (!removeNewLines) {
+                    stringBuilder.append(String.format(Locale.ROOT, "%n"));
+                }
+                line = br.readLine();
+            }
+
+            return stringBuilder.toString();
+        }
+    }
+
+}

plugins/engine-datafusion/src/test/resources/clickbench.json

@@ -0,0 +1 @@
+{"WatchID":"9110818468285196899","JavaEnable":0,"Title":"","GoodEvent":1,"EventTime":"2013-07-14 20:38:47","EventDate":"2013-07-15","CounterID":17,"ClientIP":-1216690514,"RegionID":839,"UserID":"-2461439046089301801","CounterClass":0,"OS":0,"UserAgent":0,"URL":"","Referer":"https://example.org/about","IsRefresh":0,"RefererCategoryID":0,"RefererRegionID":0,"URLCategoryID":0,"URLRegionID":0,"ResolutionWidth":0,"ResolutionHeight":0,"ResolutionDepth":0,"FlashMajor":0,"FlashMinor":0,"FlashMinor2":"","NetMajor":0,"NetMinor":0,"UserAgentMajor":0,"UserAgentMinor":"�O","CookieEnable":0,"JavascriptEnable":0,"IsMobile":0,"MobilePhone":0,"MobilePhoneModel":"","Params":"","IPNetworkID":3793327,"TraficSourceID":4,"SearchEngineID":0,"SearchPhrase":"ha","AdvEngineID":0,"IsArtifical":0,"WindowClientWidth":0,"WindowClientHeight":0,"ClientTimeZone":-1,"ClientEventTime":"1971-01-01 14:16:06","SilverlightVersion1":0,"SilverlightVersion2":0,"SilverlightVersion3":0,"SilverlightVersion4":0,"PageCharset":"","CodeVersion":0,"IsLink":0,"IsDownload":0,"IsNotBounce":0,"FUniqID":"0","OriginalURL":"","HID":0,"IsOldCounter":0,"IsEvent":0,"IsParameter":0,"DontCountHits":0,"WithHash":0,"HitColor":"5","LocalEventTime":"2013-07-15 10:47:34","Age":0,"Sex":0,"Income":0,"Interests":0,"Robotness":0,"RemoteIP":-1001831330,"WindowName":-1,"OpenerName":-1,"HistoryLength":-1,"BrowserLanguage":"�","BrowserCountry":"�\f","SocialNetwork":"","SocialAction":"","HTTPError":0,"SendTiming":0,"DNSTiming":0,"ConnectTiming":0,"ResponseStartTiming":0,"ResponseEndTiming":0,"FetchTiming":0,"SocialSourceNetworkID":0,"SocialSourcePage":"","ParamPrice":"0","ParamOrderID":"","ParamCurrency":"NH\u001C","ParamCurrencyID":0,"OpenstatServiceName":"","OpenstatCampaignID":"","OpenstatAdID":"","OpenstatSourceID":"","UTMSource":"","UTMMedium":"","UTMCampaign":"","UTMContent":"","UTMTerm":"","FromTag":"","HasGCLID":0,"RefererHash":"-296158784638538920","URLHash":"-8417682003818480435","CLID":0}