From 0dd8d827e727c8a294ce431944a86b75019ae25e Mon Sep 17 00:00:00 2001
From: Sergii Druzkin <65374769+sdruzkin@users.noreply.github.com>
Date: Fri, 15 Mar 2024 10:08:42 -0700
Subject: [PATCH] Minor refactoring of HivePageSourceProvider

---
 .../presto/hive/HivePageSourceProvider.java   | 35 ++++++++++---------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HivePageSourceProvider.java b/presto-hive/src/main/java/com/facebook/presto/hive/HivePageSourceProvider.java
index c950dd7de5320..dff54f345ecb7 100644
--- a/presto-hive/src/main/java/com/facebook/presto/hive/HivePageSourceProvider.java
+++ b/presto-hive/src/main/java/com/facebook/presto/hive/HivePageSourceProvider.java
@@ -98,7 +98,7 @@ public class HivePageSourceProvider
     private final Set<HiveRecordCursorProvider> cursorProviders;
     private final Set<HiveBatchPageSourceFactory> pageSourceFactories;
     private final Set<HiveSelectivePageSourceFactory> selectivePageSourceFactories;
-    private Set<HiveAggregatedPageSourceFactory> aggregatedPageSourceFactories;
+    private final Set<HiveAggregatedPageSourceFactory> aggregatedPageSourceFactories;
     private final TypeManager typeManager;
     private final RowExpressionService rowExpressionService;
     private final LoadingCache<RowExpressionCacheKey, RowExpression> optimizedRowExpressionCache;
@@ -255,16 +255,17 @@ private ConnectorPageSource createAggregatedPageSource(
             HiveFileContext fileContext,
             Optional<EncryptionInformation> encryptionInformation)
     {
-        for (HiveAggregatedPageSourceFactory pageSourceFactory : aggregatedPageSourceFactories) {
-            List<ColumnMapping> columnMappings = ColumnMapping.buildColumnMappings(
-                    hiveSplit.getPartitionKeys(),
-                    selectedColumns,
-                    hiveSplit.getBucketConversion().map(BucketConversion::getBucketColumnHandles).orElse(ImmutableList.of()),
-                    hiveSplit.getTableToPartitionMapping(),
-                    hiveSplit.getFileSplit(),
-                    hiveSplit.getTableBucketNumber());
+        List<ColumnMapping> columnMappings = ColumnMapping.buildColumnMappings(
+                hiveSplit.getPartitionKeys(),
+                selectedColumns,
+                hiveSplit.getBucketConversion().map(BucketConversion::getBucketColumnHandles).orElse(ImmutableList.of()),
+                hiveSplit.getTableToPartitionMapping(),
+                hiveSplit.getFileSplit(),
+                hiveSplit.getTableBucketNumber());
+
+        List<ColumnMapping> regularAndInterimColumnMappings = ColumnMapping.extractRegularAndInterimColumnMappings(columnMappings);
 
-            List<ColumnMapping> regularAndInterimColumnMappings = ColumnMapping.extractRegularAndInterimColumnMappings(columnMappings);
+        for (HiveAggregatedPageSourceFactory pageSourceFactory : aggregatedPageSourceFactories) {
             Optional<? extends ConnectorPageSource> pageSource = pageSourceFactory.createPageSource(
                     configuration,
                     session,
@@ -364,6 +365,10 @@ private static Optional<ConnectorPageSource> createSelectivePageSource(
             return Optional.of(new HiveEmptySplitPageSource());
         }
 
+        TupleDomain<Subfield> domainPredicate = splitContext.getDynamicFilterPredicate()
+                .map(filter -> filter.transform(handle -> new Subfield(((HiveColumnHandle) handle).getName())).intersect(layout.getDomainPredicate()))
+                .orElse(layout.getDomainPredicate());
+
         for (HiveSelectivePageSourceFactory pageSourceFactory : selectivePageSourceFactories) {
             Optional<? extends ConnectorPageSource> pageSource = pageSourceFactory.createPageSource(
                     configuration,
@@ -375,8 +380,7 @@ private static Optional<ConnectorPageSource> createSelectivePageSource(
                     coercers,
                     bucketAdaptation,
                     outputColumns,
-                    splitContext.getDynamicFilterPredicate().map(filter -> filter.transform(
-                            handle -> new Subfield(((HiveColumnHandle) handle).getName())).intersect(layout.getDomainPredicate())).orElse(layout.getDomainPredicate()),
+                    domainPredicate,
                     optimizedRemainingPredicate,
                     hiveStorageTimeZone,
                     fileContext,
@@ -459,7 +463,6 @@ public static Optional<ConnectorPageSource> createHivePageSource(
                     fileSplit,
                     storage,
                     effectivePredicate,
-                    hiveColumns,
                     hiveStorageTimeZone,
                     typeManager,
                     tableName,
@@ -522,7 +525,6 @@ public static Optional<ConnectorPageSource> createHivePageSource(
                 fileSplit,
                 storage,
                 effectivePredicate,
-                hiveColumns,
                 hiveStorageTimeZone,
                 typeManager,
                 tableName,
@@ -549,7 +551,6 @@ private static Optional<ConnectorPageSource> getPageSourceFromCursorProvider(
             HiveFileSplit fileSplit,
             Storage storage,
             TupleDomain<HiveColumnHandle> effectivePredicate,
-            List<HiveColumnHandle> hiveColumns,
             DateTimeZone hiveStorageTimeZone,
             TypeManager typeManager,
             SchemaTableName tableName,
@@ -585,8 +586,8 @@ private static Optional<ConnectorPageSource> getPageSourceFromCursorProvider(
                     tableParameters,
                     tableName.getSchemaName(),
                     tableName.getTableName(),
-                    partitionKeyColumnHandles.stream().map(column -> column.getName()).collect(toImmutableList()),
-                    partitionKeyColumnHandles.stream().map(column -> column.getHiveType()).collect(toImmutableList()));
+                    partitionKeyColumnHandles.stream().map(BaseHiveColumnHandle::getName).collect(toImmutableList()),
+                    partitionKeyColumnHandles.stream().map(HiveColumnHandle::getHiveType).collect(toImmutableList()));
 
             Optional<RecordCursor> cursor = provider.createRecordCursor(
                     configuration,