Add function inlining in a later optimizer pass

Adds a function inlining pass to the optimizer after the KeyBasedSampling rule in order to make sure that any rewritten functions are inlined.
prestodb · Mar 7, 2024 · 8c9178d · 8c9178d
1 parent 21340d8
commit 8c9178d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/PlanOptimizers.java
@@ -634,7 +634,15 @@ public PlanOptimizers(
                         ruleStats,
                         statsCalculator,
                         estimatedExchangesCostCalculator,
-                        ImmutableSet.of(new RemoveRedundantIdentityProjections())));
+                        ImmutableSet.of(new RemoveRedundantIdentityProjections())),
+                new IterativeOptimizer(
+                        metadata,
+                        ruleStats,
+                        statsCalculator,
+                        estimatedExchangesCostCalculator,
+                        ImmutableSet.<Rule<?>>builder()
+                                .addAll(new InlineSqlFunctions(metadata, sqlParser).rules())
+                                .build()));
 
         builder.add(
                 new IterativeOptimizer(

diff --git a/...tion/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java b/...tion/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeGeneralQueries.java
@@ -28,7 +28,9 @@
 import java.util.Optional;
 import java.util.UUID;
 
+import static com.facebook.presto.SystemSessionProperties.INLINE_SQL_FUNCTIONS;
 import static com.facebook.presto.SystemSessionProperties.JOIN_DISTRIBUTION_TYPE;
+import static com.facebook.presto.SystemSessionProperties.KEY_BASED_SAMPLING_ENABLED;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
 import static com.facebook.presto.common.type.VarcharType.VARCHAR;
 import static com.facebook.presto.hive.HiveStorageFormat.DWRF;
@@ -1485,6 +1487,19 @@ public void testSelectFieldsWithCapitalLetters()
         }
     }
 
+    /**
+     * See GitHub issue: <a href="https://github.com/prestodb/presto/issues/22085">link</a>
+     */
+    @Test
+    public void testKeyBasedSamplingInlined()
+    {
+        Session session = Session.builder(getSession())
+                .setSystemProperty(INLINE_SQL_FUNCTIONS, "true")
+                .setSystemProperty(KEY_BASED_SAMPLING_ENABLED, "true")
+                .build();
+        assertQuerySucceeds(session, "select count(1) from orders join lineitem using(orderkey)");
+    }
+
     private void assertQueryResultCount(String sql, int expectedResultCount)
     {
         assertEquals(getQueryRunner().execute(sql).getRowCount(), expectedResultCount);