Skip to content

Commit

Permalink
Fix aggregation metadata optimization when existing Limit/TopN
Browse files Browse the repository at this point in the history
  • Loading branch information
hantangwangd authored and tdcmeehan committed Mar 15, 2024
1 parent 4f67087 commit 938cd1d
Show file tree
Hide file tree
Showing 2 changed files with 51 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,50 @@ protected QueryRunner createQueryRunner()
Optional.empty());
}

@Test
public void testMetadataQueryOptimizationWithLimit()
{
QueryRunner queryRunner = getQueryRunner();
Session sessionWithOptimizeMetadataQueries = getSessionWithOptimizeMetadataQueries();
Session defaultSession = queryRunner.getDefaultSession();
try {
queryRunner.execute("CREATE TABLE test_metadata_query_optimization_with_limit(a varchar, b int, c int) WITH (partitioned_by = ARRAY['b', 'c'])");
queryRunner.execute("INSERT INTO test_metadata_query_optimization_with_limit VALUES" +
" ('1001', 1, 1), ('1002', 1, 1), ('1003', 1, 1)," +
" ('1004', 1, 2), ('1005', 1, 2), ('1006', 1, 2)," +
" ('1007', 2, 1), ('1008', 2, 1), ('1009', 2, 1)");

// Could do metadata optimization when `limit` existing above `aggregation`
assertQuery(sessionWithOptimizeMetadataQueries, "select distinct b, c from test_metadata_query_optimization_with_limit order by c desc limit 3",
"values(1, 2), (1, 1), (2, 1)");
assertPlan(sessionWithOptimizeMetadataQueries, "select distinct b, c from test_metadata_query_optimization_with_limit order by c desc limit 3",
anyTree(values(ImmutableList.of("b", "c"),
ImmutableList.of(
ImmutableList.of(new LongLiteral("1"), new LongLiteral("2")),
ImmutableList.of(new LongLiteral("1"), new LongLiteral("1")),
ImmutableList.of(new LongLiteral("2"), new LongLiteral("1"))))));
// Compare with default session which do not enable metadata optimization
assertQuery(defaultSession, "select distinct b, c from test_metadata_query_optimization_with_limit order by c desc limit 3",
"values(1, 2), (1, 1), (2, 1)");
assertPlan(defaultSession, "select distinct b, c from test_metadata_query_optimization_with_limit order by c desc limit 3",
anyTree(strictTableScan("test_metadata_query_optimization_with_limit", identityMap("b", "c"))));

// Should not do metadata optimization when `limit` existing below `aggregation`
assertQuery(sessionWithOptimizeMetadataQueries, "with tt as (select b, c from test_metadata_query_optimization_with_limit order by c desc limit 3) select b, min(c), max(c) from tt group by b",
"values(1, 2, 2)");
assertPlan(sessionWithOptimizeMetadataQueries, "with tt as (select b, c from test_metadata_query_optimization_with_limit order by c desc limit 3) select b, min(c), max(c) from tt group by b",
anyTree(strictTableScan("test_metadata_query_optimization_with_limit", identityMap("b", "c"))));
// Compare with default session which do not enable metadata optimization
assertQuery(defaultSession, "with tt as (select b, c from test_metadata_query_optimization_with_limit order by c desc limit 3) select b, min(c), max(c) from tt group by b",
"values(1, 2, 2)");
assertPlan(defaultSession, "with tt as (select b, c from test_metadata_query_optimization_with_limit order by c desc limit 3) select b, min(c), max(c) from tt group by b",
anyTree(strictTableScan("test_metadata_query_optimization_with_limit", identityMap("b", "c"))));
}
finally {
queryRunner.execute("DROP TABLE test_metadata_query_optimization_with_limit");
}
}

@Test
public void testRepeatedFilterPushdown()
{
Expand Down Expand Up @@ -1987,6 +2031,13 @@ private void assertParquetDereferencePushDown(Session session, String query, Str
assertPlan(session, query, anyTree(tableScanParquetDeferencePushDowns(tableName, expectedDeferencePushDowns)));
}

protected Session getSessionWithOptimizeMetadataQueries()
{
return Session.builder(getQueryRunner().getDefaultSession())
.setSystemProperty(OPTIMIZE_METADATA_QUERIES, "true")
.build();
}

private Session pushdownFilterEnabled()
{
return Session.builder(getQueryRunner().getDefaultSession())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,13 +34,11 @@
import com.facebook.presto.spi.plan.AggregationNode.Aggregation;
import com.facebook.presto.spi.plan.Assignments;
import com.facebook.presto.spi.plan.FilterNode;
import com.facebook.presto.spi.plan.LimitNode;
import com.facebook.presto.spi.plan.MarkDistinctNode;
import com.facebook.presto.spi.plan.PlanNode;
import com.facebook.presto.spi.plan.PlanNodeIdAllocator;
import com.facebook.presto.spi.plan.ProjectNode;
import com.facebook.presto.spi.plan.TableScanNode;
import com.facebook.presto.spi.plan.TopNNode;
import com.facebook.presto.spi.plan.ValuesNode;
import com.facebook.presto.spi.relation.ConstantExpression;
import com.facebook.presto.spi.relation.RowExpression;
Expand Down Expand Up @@ -384,8 +382,6 @@ private static Optional<TableScanNode> findTableScan(PlanNode source, RowExpress
// allow any chain of linear transformations
if (source instanceof MarkDistinctNode ||
source instanceof FilterNode ||
source instanceof LimitNode ||
source instanceof TopNNode ||
source instanceof SortNode) {
source = source.getSources().get(0);
}
Expand Down

0 comments on commit 938cd1d

Please sign in to comment.