Merge pull request #6062 from gchq/6061-remove-datafusion-fallback-to-java

Issue 6061 - Remove DataFusion query fallback to Java

Author: patchwork01

java/clients/src/main/java/sleeper/clients/api/aws/SleeperClientQueryProvider.java

@@ -23,6 +23,7 @@
 import sleeper.foreign.datafusion.DataFusionAwsConfig;
 import sleeper.parquet.utils.TableHadoopConfigurationProvider;
 import sleeper.query.core.rowretrieval.LeafPartitionRowRetrieverProvider;
+import sleeper.query.core.rowretrieval.QueryEngineSelector;
 import sleeper.query.datafusion.DataFusionQueryContext;
 import sleeper.query.runner.rowretrieval.LeafPartitionRowRetrieverImpl;
 
@@ -66,8 +67,9 @@ static SleeperClientQueryProvider withThreadPoolForEachClient(int threadPoolSize
             ExecutorService executorService = Executors.newFixedThreadPool(threadPoolSize);
             DataFusionQueryContext dataFusionContext = DataFusionQueryContext.createIfLoaded(RootAllocator::new);
             LeafPartitionRowRetrieverProvider javaProvider = new LeafPartitionRowRetrieverImpl.Provider(executorService, hadoopProvider);
+            LeafPartitionRowRetrieverProvider dataFusionProvider = dataFusionContext.createDataFusionProvider(DataFusionAwsConfig::getDefault);
             return ShutdownWrapper.shutdown(
-                    dataFusionContext.createQueryEngineSelectorWithFallback(DataFusionAwsConfig::getDefault, javaProvider),
+                    QueryEngineSelector.javaAndDataFusion(javaProvider, dataFusionProvider),
                     new UncheckedAutoCloseables(List.of(
                             dataFusionContext::close,
                             executorService::shutdown)));
@@ -95,7 +97,8 @@ private PersistentThreadPool(ExecutorService executorService) {
         @Override
         public ShutdownWrapper<LeafPartitionRowRetrieverProvider> getRowRetrieverProvider(TableHadoopConfigurationProvider hadoopProvider) {
             LeafPartitionRowRetrieverProvider javaProvider = new LeafPartitionRowRetrieverImpl.Provider(executorService, hadoopProvider);
-            return ShutdownWrapper.noShutdown(dataFusionContext.createQueryEngineSelectorWithFallback(DataFusionAwsConfig::getDefault, javaProvider));
+            LeafPartitionRowRetrieverProvider dataFusionProvider = dataFusionContext.createDataFusionProvider(DataFusionAwsConfig::getDefault);
+            return ShutdownWrapper.noShutdown(QueryEngineSelector.javaAndDataFusion(javaProvider, dataFusionProvider));
         }
 
         @Override

java/query/query-datafusion/src/main/java/sleeper/query/datafusion/DataFusionQueryContext.java

@@ -22,12 +22,12 @@
 import sleeper.foreign.bridge.FFIContext;
 import sleeper.foreign.datafusion.DataFusionAwsConfig;
 import sleeper.query.core.rowretrieval.LeafPartitionRowRetrieverProvider;
-import sleeper.query.core.rowretrieval.QueryEngineSelector;
 
 import java.util.function.Supplier;
 
 /**
- * A wrapper for the DataFusion FFI context. Allows for fallback if DataFusion could not be loaded.
+ * A wrapper for the DataFusion FFI context. Allows for delayed failure if DataFusion could not be loaded, in case you
+ * only want to use the Java data engine.
  */
 public class DataFusionQueryContext implements AutoCloseable {
 
@@ -44,7 +44,7 @@ private DataFusionQueryContext(FFIContext<DataFusionQueryFunctions> context, Buf
     /**
      * Creates the DataFusion FFI context and Arrow allocator if the DataFusion functions were loaded. The FFI context
      * and Arrow allocator will be closed when the query context is closed. If the DataFusion functions were not loaded,
-     * this allows falling back to the Java implementation.
+     * this avoids immediate failure in case you only want to use the Java data engine.
      *
      * @param  allocator a constructor for the Arrow allocator
      * @return           the context
@@ -56,35 +56,19 @@ public static DataFusionQueryContext createIfLoaded(Supplier<BufferAllocator> al
     }
 
     /**
-     * Creates a context where no DataFusion functions are used. Will always fall back to the Java implementation.
+     * Creates a context where no DataFusion functions are used. Supports cases where an instance of this class is
+     * required but we will always use the Java data engine.
      *
      * @return the context
      */
     public static DataFusionQueryContext none() {
         return new DataFusionQueryContext(null, null);
     }
 
-    /**
-     * Creates a row retriever provider for use in queries, using the Java or DataFusion implementation depending on
-     * configuration. If the DataFusion functions could not be loaded, the Java implementation will always be used.
-     *
-     * @param  awsConfig    a constructor for the AWS configuration
-     * @param  javaProvider the Java implementation
-     * @return              the row retriever provider
-     */
-    public LeafPartitionRowRetrieverProvider createQueryEngineSelectorWithFallback(Supplier<DataFusionAwsConfig> awsConfig, LeafPartitionRowRetrieverProvider javaProvider) {
-        if (context == null) {
-            LOGGER.warn("Falling back to Java row retriever as DataFusion was not loaded");
-            return javaProvider;
-        } else {
-            return QueryEngineSelector.javaAndDataFusion(javaProvider,
-                    new DataFusionLeafPartitionRowRetriever.Provider(awsConfig.get(), allocator, context));
-        }
-    }
-
     /**
      * Creates a DataFusion row retriever provider for use in queries. If the DataFusion functions could not be loaded,
-     * the returned row retriever will fail if it is used.
+     * the returned row retriever will fail if it is used. This is usually wrapped with an engine selector to allow
+     * switching between different data engines.
      *
      * @param  awsConfig a constructor for the AWS configuration
      * @return           the row retriever provider

java/query/query-lambda/src/main/java/sleeper/query/lambda/SqsLeafPartitionQueryLambda.java

@@ -33,6 +33,8 @@
 import sleeper.foreign.datafusion.DataFusionAwsConfig;
 import sleeper.parquet.utils.TableHadoopConfigurationProvider;
 import sleeper.query.core.rowretrieval.LeafPartitionQueryExecutor;
+import sleeper.query.core.rowretrieval.LeafPartitionRowRetrieverProvider;
+import sleeper.query.core.rowretrieval.QueryEngineSelector;
 import sleeper.query.datafusion.DataFusionQueryContext;
 import sleeper.query.runner.rowretrieval.LeafPartitionRowRetrieverImpl;
 import sleeper.query.runner.tracker.DynamoDBQueryTracker;
@@ -74,13 +76,14 @@ public SqsLeafPartitionQueryLambda(
         TablePropertiesProvider tablePropertiesProvider = S3TableProperties.createProvider(instanceProperties, s3Client, dynamoClient);
         TableHadoopConfigurationProvider hadoopProvider = TableHadoopConfigurationProvider.withCache(
                 TableHadoopConfigurationProvider.forQueryLambdas(instanceProperties));
+        LeafPartitionRowRetrieverProvider javaProvider = new LeafPartitionRowRetrieverImpl.Provider(
+                Executors.newFixedThreadPool(instanceProperties.getInt(QUERY_PROCESSOR_LAMBDA_ROW_RETRIEVAL_THREADS)), hadoopProvider);
+        LeafPartitionRowRetrieverProvider dataFusionProvider = dataFusionContext.createDataFusionProvider(DataFusionAwsConfig::getDefault);
         messageHandler = new QueryMessageHandler(tablePropertiesProvider, new DynamoDBQueryTracker(instanceProperties, dynamoClient));
         processor = SqsLeafPartitionQueryProcessor.builder()
                 .sqsClient(sqsClient).s3Client(s3Client).dynamoClient(dynamoClient)
                 .instanceProperties(instanceProperties).tablePropertiesProvider(tablePropertiesProvider).hadoopProvider(hadoopProvider)
-                .rowRetrieverProvider(dataFusionContext.createQueryEngineSelectorWithFallback(awsConfig,
-                        new LeafPartitionRowRetrieverImpl.Provider(
-                                Executors.newFixedThreadPool(instanceProperties.getInt(QUERY_PROCESSOR_LAMBDA_ROW_RETRIEVAL_THREADS)), hadoopProvider)))
+                .rowRetrieverProvider(QueryEngineSelector.javaAndDataFusion(javaProvider, dataFusionProvider))
                 .build();
     }