feat: Cost-based MV candidate selection for query rewriting (#27222)

presto-docs/src/main/sphinx/admin/properties-session.rst

@@ -754,6 +754,20 @@ in the server configuration.
 
 The corresponding configuration property is :ref:`admin/properties:\`\`experimental.legacy-materialized-views\`\``.
 
+``materialized_view_query_rewrite_cost_based_selection_enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``boolean``
+* **Default value:** ``false``
+
+Enable cost-based selection when multiple materialized views are available for query
+rewriting. When enabled, the optimizer evaluates all compatible materialized view rewrites
+and selects the plan with the lowest estimated cost, instead of using the first compatible
+view.
+
+The corresponding configuration property is
+:ref:`admin/properties:\`\`materialized-view-query-rewrite-cost-based-selection-enabled\`\``.
+
 ``materialized_view_stale_read_behavior``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

presto-docs/src/main/sphinx/admin/properties.rst

@@ -1540,6 +1540,20 @@ migration purposes only.
 
     This should only be enabled in non-production environments.
 
+``materialized-view-query-rewrite-cost-based-selection-enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``boolean``
+* **Default value:** ``false``
+
+Enable cost-based selection when multiple materialized views are available for query
+rewriting. When enabled, the optimizer evaluates all compatible materialized view rewrites
+and selects the plan with the lowest estimated cost, instead of using the first compatible
+view.
+
+The corresponding session property is
+:ref:`admin/properties-session:\`\`materialized_view_query_rewrite_cost_based_selection_enabled\`\``.
+
 ``materialized-view-stale-read-behavior``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

presto-hive/src/test/java/com/facebook/presto/hive/TestHiveMaterializedViewLogicalPlanner.java

@@ -29,6 +29,7 @@
 import com.facebook.presto.testing.QueryRunner;
 import com.facebook.presto.tests.AbstractTestQueryFramework;
 import com.facebook.presto.tests.DistributedQueryRunner;
+import com.facebook.presto.tests.QueryAssertions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
@@ -42,6 +43,7 @@
 import static com.facebook.presto.SystemSessionProperties.JOIN_DISTRIBUTION_TYPE;
 import static com.facebook.presto.SystemSessionProperties.JOIN_REORDERING_STRATEGY;
 import static com.facebook.presto.SystemSessionProperties.MATERIALIZED_VIEW_DATA_CONSISTENCY_ENABLED;
+import static com.facebook.presto.SystemSessionProperties.MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED;
 import static com.facebook.presto.SystemSessionProperties.PREFER_PARTIAL_AGGREGATION;
 import static com.facebook.presto.SystemSessionProperties.QUERY_OPTIMIZATION_WITH_MATERIALIZED_VIEW_ENABLED;
 import static com.facebook.presto.SystemSessionProperties.SIMPLIFY_PLAN_WITH_EMPTY_INPUT;
@@ -3447,6 +3449,142 @@ public void testMaterializedViewRefreshedInNonLegacyMode()
         }
     }
 
+    @Test
+    public void testCostBasedMVSelectionWithMultipleCandidates()
+    {
+        Session costBasedMVSession = Session.builder(getSession())
+                .setSystemProperty(QUERY_OPTIMIZATION_WITH_MATERIALIZED_VIEW_ENABLED, "true")
+                .setSystemProperty(MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED, "true")
+                .setSystemProperty(MATERIALIZED_VIEW_DATA_CONSISTENCY_ENABLED, "false")
+                .setSystemProperty(SIMPLIFY_PLAN_WITH_EMPTY_INPUT, "false")
+                .build();
+        QueryRunner queryRunner = getQueryRunner();
+        String table = "orders_cost_based";
+        String wideView = "orders_wide_mv";
+        String narrowView = "orders_narrow_mv";
+        try {
+            queryRunner.execute(format("CREATE TABLE %s WITH (partitioned_by = ARRAY['ds']) AS " +
+                    "SELECT orderkey, orderpriority, orderdate, totalprice, '2020-01-01' as ds FROM orders WHERE orderkey < 1000 " +
+                    "UNION ALL " +
+                    "SELECT orderkey, orderpriority, orderdate, totalprice, '2020-01-02' as ds FROM orders WHERE orderkey > 1000", table));
+
+            // Wide MV: all columns
+            assertUpdate(format("CREATE MATERIALIZED VIEW %s WITH (partitioned_by = ARRAY['ds']) " +
+                    "AS SELECT orderkey, orderpriority, orderdate, totalprice, ds FROM %s", wideView, table));
+            // Narrow MV: only orderkey + totalprice (fewer columns, should be cheaper for queries needing only these)
+            assertUpdate(format("CREATE MATERIALIZED VIEW %s WITH (partitioned_by = ARRAY['ds']) " +
+                    "AS SELECT orderkey, totalprice, ds FROM %s", narrowView, table));
+
+            setReferencedMaterializedViews((DistributedQueryRunner) queryRunner, table, ImmutableList.of(wideView, narrowView));
+
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-01'", wideView), 255);
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-02'", wideView), 14745);
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-01'", narrowView), 255);
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-02'", narrowView), 14745);
+
+            // Query needs orderkey + totalprice: both MVs are compatible, optimizer should pick one
+            String baseQuery = format("SELECT orderkey, totalprice FROM %s WHERE orderkey < 1000", table);
+
+            // Verify correctness: optimized results match non-optimized
+            QueryAssertions.assertQuery(queryRunner, costBasedMVSession, baseQuery, queryRunner, getSession(), baseQuery, false, false);
+
+            // Should rewrite to use the narrow MV (fewer columns, cheaper)
+            assertPlan(costBasedMVSession, baseQuery, anyTree(
+                    filter("orderkey_0 < BIGINT'1000'", constrainedTableScan(narrowView,
+                            ImmutableMap.of(),
+                            ImmutableMap.of("orderkey_0", "orderkey")))));
+        }
+        finally {
+            queryRunner.execute("DROP MATERIALIZED VIEW IF EXISTS " + wideView);
+            queryRunner.execute("DROP MATERIALIZED VIEW IF EXISTS " + narrowView);
+            queryRunner.execute("DROP TABLE IF EXISTS " + table);
+        }
+    }
+
+    @Test
+    public void testCostBasedMVSelectionSingleCandidate()
+    {
+        Session costBasedMVSession = Session.builder(getSession())
+                .setSystemProperty(QUERY_OPTIMIZATION_WITH_MATERIALIZED_VIEW_ENABLED, "true")
+                .setSystemProperty(MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED, "true")
+                .setSystemProperty(MATERIALIZED_VIEW_DATA_CONSISTENCY_ENABLED, "false")
+                .setSystemProperty(SIMPLIFY_PLAN_WITH_EMPTY_INPUT, "false")
+                .build();
+        QueryRunner queryRunner = getQueryRunner();
+        String table = "orders_cost_single";
+        String view = "orders_single_mv";
+        try {
+            queryRunner.execute(format("CREATE TABLE %s WITH (partitioned_by = ARRAY['ds']) AS " +
+                    "SELECT orderkey, orderpriority, orderdate, totalprice, '2020-01-01' as ds FROM orders WHERE orderkey < 1000 " +
+                    "UNION ALL " +
+                    "SELECT orderkey, orderpriority, orderdate, totalprice, '2020-01-02' as ds FROM orders WHERE orderkey > 1000", table));
+
+            assertUpdate(format("CREATE MATERIALIZED VIEW %s WITH (partitioned_by = ARRAY['ds']) " +
+                    "AS SELECT orderkey, orderdate, ds FROM %s", view, table));
+
+            setReferencedMaterializedViews((DistributedQueryRunner) queryRunner, table, ImmutableList.of(view));
+
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-01'", view), 255);
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-02'", view), 14745);
+
+            String baseQuery = format("SELECT orderkey, orderdate FROM %s WHERE orderkey < 1000", table);
+
+            // Verify correctness: optimized results match non-optimized
+            QueryAssertions.assertQuery(queryRunner, costBasedMVSession, baseQuery, queryRunner, getSession(), baseQuery, false, false);
+
+            // Should rewrite to use the single compatible MV
+            assertPlan(costBasedMVSession, baseQuery, anyTree(
+                    filter("orderkey_0 < BIGINT'1000'", constrainedTableScan(view,
+                            ImmutableMap.of(),
+                            ImmutableMap.of("orderkey_0", "orderkey")))));
+        }
+        finally {
+            queryRunner.execute("DROP MATERIALIZED VIEW IF EXISTS " + view);
+            queryRunner.execute("DROP TABLE IF EXISTS " + table);
+        }
+    }
+
+    @Test
+    public void testCostBasedMVSelectionNoCompatibleCandidates()
+    {
+        Session costBasedMVSession = Session.builder(getSession())
+                .setSystemProperty(QUERY_OPTIMIZATION_WITH_MATERIALIZED_VIEW_ENABLED, "true")
+                .setSystemProperty(MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED, "true")
+                .setSystemProperty(MATERIALIZED_VIEW_DATA_CONSISTENCY_ENABLED, "false")
+                .setSystemProperty(SIMPLIFY_PLAN_WITH_EMPTY_INPUT, "false")
+                .build();
+        QueryRunner queryRunner = getQueryRunner();
+        String table = "orders_cost_nocompat";
+        String view = "orders_nocompat_mv";
+        try {
+            queryRunner.execute(format("CREATE TABLE %s WITH (partitioned_by = ARRAY['ds']) AS " +
+                    "SELECT orderkey, orderpriority, orderdate, totalprice, '2020-01-01' as ds FROM orders WHERE orderkey < 1000 " +
+                    "UNION ALL " +
+                    "SELECT orderkey, orderpriority, orderdate, totalprice, '2020-01-02' as ds FROM orders WHERE orderkey > 1000", table));
+
+            // MV has orderpriority but query needs orderdate — incompatible
+            assertUpdate(format("CREATE MATERIALIZED VIEW %s WITH (partitioned_by = ARRAY['ds']) " +
+                    "AS SELECT orderkey, orderpriority, ds FROM %s", view, table));
+
+            setReferencedMaterializedViews((DistributedQueryRunner) queryRunner, table, ImmutableList.of(view));
+
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-01'", view), 255);
+            assertUpdate(format("REFRESH MATERIALIZED VIEW %s WHERE ds='2020-01-02'", view), 14745);
+
+            // Query needs orderdate which the MV doesn't have — should fall back to base table
+            String baseQuery = format("SELECT orderkey, orderdate FROM %s WHERE orderkey < 1000", table);
+
+            assertPlan(costBasedMVSession, baseQuery, anyTree(
+                    filter("orderkey < BIGINT'1000'", constrainedTableScan(table,
+                            ImmutableMap.of(),
+                            ImmutableMap.of("orderkey", "orderkey")))));
+        }
+        finally {
+            queryRunner.execute("DROP MATERIALIZED VIEW IF EXISTS " + view);
+            queryRunner.execute("DROP TABLE IF EXISTS " + table);
+        }
+    }
+
     private void setReferencedMaterializedViews(DistributedQueryRunner queryRunner, String tableName, List<String> referencedMaterializedViews)
     {
         appendTableParameter(replicateHiveMetastore(queryRunner),

presto-main-base/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -258,6 +258,7 @@ public final class SystemSessionProperties
     public static final String MATERIALIZED_VIEW_DATA_CONSISTENCY_ENABLED = "materialized_view_data_consistency_enabled";
     public static final String CONSIDER_QUERY_FILTERS_FOR_MATERIALIZED_VIEW_PARTITIONS = "consider-query-filters-for-materialized-view-partitions";
     public static final String QUERY_OPTIMIZATION_WITH_MATERIALIZED_VIEW_ENABLED = "query_optimization_with_materialized_view_enabled";
+    public static final String MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED = "materialized_view_query_rewrite_cost_based_selection_enabled";
     public static final String LEGACY_MATERIALIZED_VIEWS = "legacy_materialized_views";
     public static final String MATERIALIZED_VIEW_ALLOW_FULL_REFRESH_ENABLED = "materialized_view_allow_full_refresh_enabled";
     public static final String MATERIALIZED_VIEW_STALE_READ_BEHAVIOR = "materialized_view_stale_read_behavior";
@@ -1442,6 +1443,11 @@ public SystemSessionProperties(
                         "Enable query optimization with materialized view",
                         featuresConfig.isQueryOptimizationWithMaterializedViewEnabled(),
                         true),
+                booleanProperty(
+                        MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED,
+                        "When enabled, collect all compatible MV candidates and defer selection to cost-based optimizer instead of using the first compatible MV",
+                        featuresConfig.isMaterializedViewQueryRewriteCostBasedSelectionEnabled(),
+                        false),
                 new PropertyMetadata<>(
                         LEGACY_MATERIALIZED_VIEWS,
                         "Experimental: Use legacy materialized views.  This feature is under active development and may change " +
@@ -3148,6 +3154,11 @@ public static boolean isQueryOptimizationWithMaterializedViewEnabled(Session ses
         return session.getSystemProperty(QUERY_OPTIMIZATION_WITH_MATERIALIZED_VIEW_ENABLED, Boolean.class);
     }
 
+    public static boolean isMaterializedViewQueryRewriteCostBasedSelectionEnabled(Session session)
+    {
+        return session.getSystemProperty(MATERIALIZED_VIEW_QUERY_REWRITE_COST_BASED_SELECTION_ENABLED, Boolean.class);
+    }
+
     public static boolean isLegacyMaterializedViews(Session session)
     {
         return session.getSystemProperty(LEGACY_MATERIALIZED_VIEWS, Boolean.class);

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java

@@ -238,6 +238,7 @@ public class FeaturesConfig
     private boolean materializedViewDataConsistencyEnabled = true;
     private boolean materializedViewPartitionFilteringEnabled = true;
     private boolean queryOptimizationWithMaterializedViewEnabled;
+    private boolean materializedViewQueryRewriteCostBasedSelectionEnabled;
     private boolean legacyMaterializedViewRefresh = true;
     private boolean allowLegacyMaterializedViewsToggle;
     private boolean materializedViewAllowFullRefreshEnabled;
@@ -2333,6 +2334,19 @@ public FeaturesConfig setQueryOptimizationWithMaterializedViewEnabled(boolean va
         return this;
     }
 
+    public boolean isMaterializedViewQueryRewriteCostBasedSelectionEnabled()
+    {
+        return materializedViewQueryRewriteCostBasedSelectionEnabled;
+    }
+
+    @Config("materialized-view-query-rewrite-cost-based-selection-enabled")
+    @ConfigDescription("When enabled, collect all compatible MV candidates and defer selection to cost-based optimizer instead of using the first compatible MV")
+    public FeaturesConfig setMaterializedViewQueryRewriteCostBasedSelectionEnabled(boolean value)
+    {
+        this.materializedViewQueryRewriteCostBasedSelectionEnabled = value;
+        return this;
+    }
+
     public boolean isLegacyMaterializedViews()
     {
         return legacyMaterializedViewRefresh;

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/MaterializedViewQueryOptimizer.java

@@ -62,6 +62,7 @@
 import com.facebook.presto.sql.tree.Query;
 import com.facebook.presto.sql.tree.QueryBody;
 import com.facebook.presto.sql.tree.QuerySpecification;
+import com.facebook.presto.sql.tree.QueryWithMVRewriteCandidates;
 import com.facebook.presto.sql.tree.Relation;
 import com.facebook.presto.sql.tree.SampledRelation;
 import com.facebook.presto.sql.tree.Select;
@@ -86,6 +87,7 @@
 
 import static com.facebook.presto.SystemSessionProperties.isMaterializedViewDataConsistencyEnabled;
 import static com.facebook.presto.SystemSessionProperties.isMaterializedViewPartitionFilteringEnabled;
+import static com.facebook.presto.SystemSessionProperties.isMaterializedViewQueryRewriteCostBasedSelectionEnabled;
 import static com.facebook.presto.common.RuntimeMetricName.MANY_PARTITIONS_MISSING_IN_MATERIALIZED_VIEW_COUNT;
 import static com.facebook.presto.common.RuntimeMetricName.OPTIMIZED_WITH_MATERIALIZED_VIEW_SUBQUERY_COUNT;
 import static com.facebook.presto.common.RuntimeUnit.NONE;
@@ -335,7 +337,7 @@ private <T extends Node> T processSameType(T node)
         return (T) process(node);
     }
 
-    private QuerySpecification rewriteQuerySpecificationIfCompatible(QuerySpecification querySpecification, Table baseTable)
+    private Node rewriteQuerySpecificationIfCompatible(QuerySpecification querySpecification, Table baseTable)
     {
         Optional<String> errorMessage = MaterializedViewRewriteQueryShapeValidator.validate(querySpecification);
 
@@ -351,6 +353,11 @@ private QuerySpecification rewriteQuerySpecificationIfCompatible(QuerySpecificat
         // TODO: Select the most compatible and efficient materialized view for query rewrite optimization https://github.com/prestodb/presto/issues/16431
         // TODO: Refactor query optimization code https://github.com/prestodb/presto/issues/16759
 
+        if (isMaterializedViewQueryRewriteCostBasedSelectionEnabled(session)) {
+            return rewriteWithAllCandidates(querySpecification, referencedMaterializedViews);
+        }
+
+        // Default path: return the first compatible MV
         for (QualifiedObjectName materializedViewName : referencedMaterializedViews) {
             QuerySpecification rewrittenQuerySpecification = getRewrittenQuerySpecification(materializedViewName, querySpecification);
             if (rewrittenQuerySpecification != querySpecification) {
@@ -360,6 +367,31 @@ private QuerySpecification rewriteQuerySpecificationIfCompatible(QuerySpecificat
         return querySpecification;
     }
 
+    private QueryBody rewriteWithAllCandidates(
+            QuerySpecification originalQuery,
+            List<QualifiedObjectName> referencedMaterializedViews)
+    {
+        ImmutableList.Builder<QueryWithMVRewriteCandidates.MVRewriteCandidate> candidates = ImmutableList.builder();
+
+        for (QualifiedObjectName materializedViewName : referencedMaterializedViews) {
+            QuerySpecification rewrittenQuerySpecification = getRewrittenQuerySpecification(materializedViewName, originalQuery);
+            if (rewrittenQuerySpecification != originalQuery) {
+                candidates.add(new QueryWithMVRewriteCandidates.MVRewriteCandidate(
+                        rewrittenQuerySpecification,
+                        materializedViewName.getCatalogName(),
+                        materializedViewName.getSchemaName(),
+                        materializedViewName.getObjectName()));
+            }
+        }
+
+        List<QueryWithMVRewriteCandidates.MVRewriteCandidate> candidateList = candidates.build();
+        if (candidateList.isEmpty()) {
+            return originalQuery;
+        }
+
+        return new QueryWithMVRewriteCandidates(originalQuery, candidateList);
+    }
+
     private QuerySpecification getRewrittenQuerySpecification(QualifiedObjectName materializedViewName, QuerySpecification originalQuerySpecification)
     {
         MaterializedViewDefinition materializedViewDefinition = metadataResolver.getMaterializedView(materializedViewName).orElseThrow(() ->

presto-main-base/src/main/java/com/facebook/presto/sql/analyzer/StatementAnalyzer.java

@@ -172,6 +172,7 @@
 import com.facebook.presto.sql.tree.Query;
 import com.facebook.presto.sql.tree.QueryBody;
 import com.facebook.presto.sql.tree.QuerySpecification;
+import com.facebook.presto.sql.tree.QueryWithMVRewriteCandidates;
 import com.facebook.presto.sql.tree.RefreshMaterializedView;
 import com.facebook.presto.sql.tree.Relation;
 import com.facebook.presto.sql.tree.RenameColumn;
@@ -3076,6 +3077,18 @@ protected Scope visitQuerySpecification(QuerySpecification node, Optional<Scope>
             return outputScope;
         }
 
+        @Override
+        protected Scope visitQueryWithMVRewriteCandidates(QueryWithMVRewriteCandidates node, Optional<Scope> scope)
+        {
+            Scope originalScope = process(node.getOriginalQuery(), scope);
+
+            for (QueryWithMVRewriteCandidates.MVRewriteCandidate candidate : node.getCandidates()) {
+                process(candidate.getRewrittenQuery(), scope);
+            }
+
+            return createAndAssignScope(node, scope, originalScope.getRelationType());
+        }
+
         @Override
         protected Scope visitSetOperation(SetOperation node, Optional<Scope> scope)
         {