Introduce SparkParameterComposerCollection (opensearch-project#2774)

* Introduce SparkParameterComposerCollection Signed-off-by: Tomoyuki Morita <moritato@amazon.com> * Fix comments Signed-off-by: Tomoyuki Morita <moritato@amazon.com> * Fix integ test Signed-off-by: Tomoyuki Morita <moritato@amazon.com> --------- Signed-off-by: Tomoyuki Morita <moritato@amazon.com>
manasvinibs · Aug 14, 2024 · 3617afe · 3617afe
1 parent ebaaf13
commit 3617afe
Show file tree

Hide file tree

Showing 37 changed files with 1,387 additions and 706 deletions.
diff --git a/async-query-core/README.md b/async-query-core/README.md
@@ -27,6 +27,8 @@ Following is the list of extension points where the consumer of the library need
   - [QueryIdProvider](src/main/java/org/opensearch/sql/spark/dispatcher/QueryIdProvider.java)
   - [SessionIdProvider](src/main/java/org/opensearch/sql/spark/execution/session/SessionIdProvider.java)
   - [SessionConfigSupplier](src/main/java/org/opensearch/sql/spark/execution/session/SessionConfigSupplier.java)
-  - [SparkExecutionEngineConfigSupplier](src/main/java/org/opensearch/sql/spark/config/SparkExecutionEngineConfigSupplier.java)
-  - [SparkSubmitParameterModifier](src/main/java/org/opensearch/sql/spark/config/SparkSubmitParameterModifier.java)
   - [EMRServerlessClientFactory](src/main/java/org/opensearch/sql/spark/client/EMRServerlessClientFactory.java)
+  - [SparkExecutionEngineConfigSupplier](src/main/java/org/opensearch/sql/spark/config/SparkExecutionEngineConfigSupplier.java)
+  - [DataSourceSparkParameterComposer](src/main/java/org/opensearch/sql/spark/parameter/DataSourceSparkParameterComposer.java)
+  - [GeneralSparkParameterComposer](src/main/java/org/opensearch/sql/spark/parameter/GeneralSparkParameterComposer.java)
+  - [SparkSubmitParameterModifier](src/main/java/org/opensearch/sql/spark/config/SparkSubmitParameterModifier.java) To be deprecated in favor of GeneralSparkParameterComposer
diff --git a/...y-core/src/main/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParameters.java b/...y-core/src/main/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParameters.java
diff --git a/...uery-core/src/main/java/org/opensearch/sql/spark/config/SparkSubmitParameterModifier.java b/...uery-core/src/main/java/org/opensearch/sql/spark/config/SparkSubmitParameterModifier.java
@@ -1,11 +1,12 @@
 package org.opensearch.sql.spark.config;
 
-import org.opensearch.sql.spark.asyncquery.model.SparkSubmitParameters;
+import org.opensearch.sql.spark.parameter.SparkSubmitParametersBuilder;
 
 /**
  * Interface for extension point to allow modification of spark submit parameter. modifyParameter
- * method is called after the default spark submit parameter is build.
+ * method is called after the default spark submit parameter is build. To be deprecated in favor of
+ * {@link org.opensearch.sql.spark.parameter.GeneralSparkParameterComposer}
  */
 public interface SparkSubmitParameterModifier {
-  void modifyParameters(SparkSubmitParameters parameters);
+  void modifyParameters(SparkSubmitParametersBuilder parametersBuilder);
 }
diff --git a/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/BatchQueryHandler.java b/async-query-core/src/main/java/org/opensearch/sql/spark/dispatcher/BatchQueryHandler.java
@@ -16,7 +16,6 @@
 import org.json.JSONObject;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
-import org.opensearch.sql.spark.asyncquery.model.SparkSubmitParameters;
 import org.opensearch.sql.spark.client.EMRServerlessClient;
 import org.opensearch.sql.spark.client.StartJobRequest;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryContext;
@@ -25,6 +24,7 @@
 import org.opensearch.sql.spark.dispatcher.model.JobType;
 import org.opensearch.sql.spark.leasemanager.LeaseManager;
 import org.opensearch.sql.spark.metrics.MetricsService;
+import org.opensearch.sql.spark.parameter.SparkSubmitParametersBuilderProvider;
 import org.opensearch.sql.spark.response.JobExecutionResponseReader;
 
 /**
@@ -37,6 +37,7 @@ public class BatchQueryHandler extends AsyncQueryHandler {
   protected final JobExecutionResponseReader jobExecutionResponseReader;
   protected final LeaseManager leaseManager;
   protected final MetricsService metricsService;
+  protected final SparkSubmitParametersBuilderProvider sparkSubmitParametersBuilderProvider;
 
   @Override
   protected JSONObject getResponseFromResultIndex(AsyncQueryJobMetadata asyncQueryJobMetadata) {
@@ -80,12 +81,16 @@ public DispatchQueryResponse submit(
             dispatchQueryRequest.getAccountId(),
             dispatchQueryRequest.getApplicationId(),
             dispatchQueryRequest.getExecutionRoleARN(),
-            SparkSubmitParameters.builder()
+            sparkSubmitParametersBuilderProvider
+                .getSparkSubmitParametersBuilder()
                 .clusterName(clusterName)
-                .dataSource(context.getDataSourceMetadata())
                 .query(dispatchQueryRequest.getQuery())
-                .build()
+                .dataSource(
+                    context.getDataSourceMetadata(),
+                    dispatchQueryRequest,
+                    context.getAsyncQueryRequestContext())
                 .acceptModifier(dispatchQueryRequest.getSparkSubmitParameterModifier())
+                .acceptComposers(dispatchQueryRequest, context.getAsyncQueryRequestContext())
                 .toString(),
             tags,
             false,

diff --git a/...query-core/src/main/java/org/opensearch/sql/spark/dispatcher/InteractiveQueryHandler.java b/...query-core/src/main/java/org/opensearch/sql/spark/dispatcher/InteractiveQueryHandler.java
@@ -16,7 +16,6 @@
 import org.json.JSONObject;
 import org.opensearch.sql.datasource.model.DataSourceMetadata;
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
-import org.opensearch.sql.spark.asyncquery.model.SparkSubmitParameters;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryContext;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse;
@@ -32,6 +31,7 @@
 import org.opensearch.sql.spark.leasemanager.model.LeaseRequest;
 import org.opensearch.sql.spark.metrics.EmrMetrics;
 import org.opensearch.sql.spark.metrics.MetricsService;
+import org.opensearch.sql.spark.parameter.SparkSubmitParametersBuilderProvider;
 import org.opensearch.sql.spark.response.JobExecutionResponseReader;
 
 /**
@@ -46,6 +46,7 @@ public class InteractiveQueryHandler extends AsyncQueryHandler {
   private final JobExecutionResponseReader jobExecutionResponseReader;
   private final LeaseManager leaseManager;
   private final MetricsService metricsService;
+  protected final SparkSubmitParametersBuilderProvider sparkSubmitParametersBuilderProvider;
 
   @Override
   protected JSONObject getResponseFromResultIndex(AsyncQueryJobMetadata asyncQueryJobMetadata) {
@@ -112,12 +113,16 @@ public DispatchQueryResponse submit(
                   dispatchQueryRequest.getAccountId(),
                   dispatchQueryRequest.getApplicationId(),
                   dispatchQueryRequest.getExecutionRoleARN(),
-                  SparkSubmitParameters.builder()
+                  sparkSubmitParametersBuilderProvider
+                      .getSparkSubmitParametersBuilder()
                       .className(FLINT_SESSION_CLASS_NAME)
                       .clusterName(clusterName)
-                      .dataSource(dataSourceMetadata)
-                      .build()
-                      .acceptModifier(dispatchQueryRequest.getSparkSubmitParameterModifier()),
+                      .dataSource(
+                          dataSourceMetadata,
+                          dispatchQueryRequest,
+                          context.getAsyncQueryRequestContext())
+                      .acceptModifier(dispatchQueryRequest.getSparkSubmitParameterModifier())
+                      .acceptComposers(dispatchQueryRequest, context.getAsyncQueryRequestContext()),
                   tags,
                   dataSourceMetadata.getResultIndex(),
                   dataSourceMetadata.getName()),