Query Plan Optimizer (Part 2) (#1919)

* Removed source from ColumnProjection so it can be built with a builder

* Minor fixes

* Fleshed out logic to split filters for subquery optimization

* Added unit tests

* Finished unit tests

* Cleanup

* Checkstyle fix

Co-authored-by: Aaron Klish <klish@verizonmedia.com>

elide-datastore/elide-datastore-aggregation/src/main/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/query/SubqueryFilterSplitter.java

@@ -0,0 +1,134 @@
+/*
+ * Copyright 2021, Yahoo Inc.
+ * Licensed under the Apache License, Version 2.0
+ * See LICENSE file in project root for terms.
+ */
+
+package com.yahoo.elide.datastores.aggregation.queryengines.sql.query;
+
+import com.yahoo.elide.core.filter.expression.AndFilterExpression;
+import com.yahoo.elide.core.filter.expression.FilterExpression;
+import com.yahoo.elide.core.filter.expression.FilterExpressionVisitor;
+import com.yahoo.elide.core.filter.expression.NotFilterExpression;
+import com.yahoo.elide.core.filter.expression.OrFilterExpression;
+import com.yahoo.elide.core.filter.predicates.FilterPredicate;
+import com.yahoo.elide.core.filter.visitors.FilterExpressionNormalizationVisitor;
+import com.yahoo.elide.core.type.Type;
+import com.yahoo.elide.datastores.aggregation.metadata.MetaDataStore;
+import com.yahoo.elide.datastores.aggregation.queryengines.sql.metadata.SQLReferenceTable;
+import com.yahoo.elide.datastores.aggregation.queryengines.sql.metadata.SQLTable;
+
+import lombok.Builder;
+import lombok.Data;
+
+import java.util.Set;
+
+/**
+ * Splits a filter expression into two parts:
+ * 1. A part that can be safely run on a nested subquery where no joins take place.
+ * 2. A part that must run on the outer query because of joins to other tables.
+ */
+public class SubqueryFilterSplitter
+        implements FilterExpressionVisitor<SubqueryFilterSplitter.SplitFilter> {
+
+    @Data
+    @Builder
+    public static class SplitFilter {
+        FilterExpression outer;
+        FilterExpression inner;
+    }
+
+    private SQLReferenceTable lookupTable;
+    private MetaDataStore metaDataStore;
+
+    public SubqueryFilterSplitter(MetaDataStore metaDataStore, SQLReferenceTable lookupTable) {
+        this.metaDataStore = metaDataStore;
+        this.lookupTable = lookupTable;
+    }
+
+    public static SplitFilter splitFilter(
+            SQLReferenceTable lookupTable,
+            MetaDataStore metaDataStore,
+            FilterExpression expression) {
+        FilterExpressionNormalizationVisitor normalizer = new FilterExpressionNormalizationVisitor();
+        FilterExpression normalizedExpression = expression.accept(normalizer);
+
+        return normalizedExpression.accept(new SubqueryFilterSplitter(metaDataStore, lookupTable));
+    }
+
+    @Override
+    public SplitFilter visitPredicate(FilterPredicate filterPredicate) {
+        Type<?> tableType = filterPredicate.getEntityType();
+        String fieldName = filterPredicate.getField();
+
+        SQLTable table = (SQLTable) metaDataStore.getTable(tableType);
+
+        Set<String> joins = lookupTable.getResolvedJoinExpressions(table, fieldName);
+
+        if (joins.size() > 0) {
+            return SplitFilter.builder().outer(filterPredicate).build();
+        } else {
+            return SplitFilter.builder().inner(filterPredicate).build();
+        }
+    }
+
+    @Override
+    public SplitFilter visitAndExpression(AndFilterExpression expression) {
+        SplitFilter lhs = expression.getLeft().accept(this);
+        SplitFilter rhs = expression.getRight().accept(this);
+
+        return SplitFilter.builder()
+                .outer(AndFilterExpression.fromPair(lhs.getOuter(), rhs.getOuter()))
+                .inner(AndFilterExpression.fromPair(lhs.getInner(), rhs.getInner()))
+                .build();
+    }
+
+    @Override
+    public SplitFilter visitOrExpression(OrFilterExpression expression) {
+        SplitFilter lhs = expression.getLeft().accept(this);
+        SplitFilter rhs = expression.getRight().accept(this);
+
+        //If either the left or right side of the expression require an outer query, the entire
+        //expression must be run as an outer query.
+        if (lhs.getOuter() != null || rhs.getOuter() != null) {
+
+            //The only operation that splits a filter into inner & outer is AND.  We can recombine
+            //each side using AND and then OR the results.
+            FilterExpression combined = OrFilterExpression.fromPair(
+                    AndFilterExpression.fromPair(lhs.getOuter(), lhs.getInner()),
+                    AndFilterExpression.fromPair(rhs.getOuter(), rhs.getInner()));
+
+            return SplitFilter.builder()
+                    .outer(combined)
+                    .build();
+
+        //Both left and right are inner queries.
+        } else {
+            FilterExpression combined = OrFilterExpression.fromPair(
+                    lhs.getInner(),
+                    rhs.getInner());
+
+            return SplitFilter.builder()
+                    .inner(combined)
+                    .build();
+        }
+    }
+
+    @Override
+    public SplitFilter visitNotExpression(NotFilterExpression expression) {
+        SplitFilter negated = expression.getNegated().accept(this);
+
+        FilterExpression outerFilter = negated.getOuter() == null
+                ? null
+                : new NotFilterExpression(negated.getOuter());
+
+        FilterExpression innerFilter = negated.getInner() == null
+                ? null
+                : new NotFilterExpression(negated.getInner());
+
+        return SplitFilter.builder()
+                .outer(outerFilter)
+                .inner(innerFilter)
+                .build();
+    }
+}

elide-datastore/elide-datastore-aggregation/src/test/java/com/yahoo/elide/datastores/aggregation/queryengines/sql/query/SubqueryFilterSplitterTest.java

@@ -0,0 +1,157 @@
+/*
+ * Copyright 2021, Yahoo Inc.
+ * Licensed under the Apache License, Version 2.0
+ * See LICENSE file in project root for terms.
+ */
+
+package com.yahoo.elide.datastores.aggregation.queryengines.sql.query;
+
+import static com.yahoo.elide.core.dictionary.EntityDictionary.NO_VERSION;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.mockito.Mockito.mock;
+import com.yahoo.elide.core.dictionary.EntityDictionary;
+import com.yahoo.elide.core.filter.dialect.ParseException;
+import com.yahoo.elide.core.filter.dialect.RSQLFilterDialect;
+import com.yahoo.elide.core.filter.expression.FilterExpression;
+import com.yahoo.elide.core.type.ClassType;
+import com.yahoo.elide.core.type.Type;
+import com.yahoo.elide.datastores.aggregation.example.Country;
+import com.yahoo.elide.datastores.aggregation.example.Player;
+import com.yahoo.elide.datastores.aggregation.example.PlayerRanking;
+import com.yahoo.elide.datastores.aggregation.example.PlayerStats;
+import com.yahoo.elide.datastores.aggregation.example.SubCountry;
+import com.yahoo.elide.datastores.aggregation.metadata.MetaDataStore;
+import com.yahoo.elide.datastores.aggregation.queryengines.sql.ConnectionDetails;
+import com.yahoo.elide.datastores.aggregation.queryengines.sql.SQLQueryEngine;
+import com.yahoo.elide.datastores.aggregation.queryengines.sql.dialects.SQLDialectFactory;
+import com.yahoo.elide.datastores.aggregation.queryengines.sql.metadata.SQLReferenceTable;
+import org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Set;
+import javax.sql.DataSource;
+
+public class SubqueryFilterSplitterTest {
+
+    private MetaDataStore metaDataStore;
+    private SQLReferenceTable lookupTable;
+    private RSQLFilterDialect dialect;
+    private static final Type<PlayerStats> PLAYER_STATS_TYPE = ClassType.of(PlayerStats.class);
+
+    public SubqueryFilterSplitterTest() {
+
+        Set<Type<?>> models = new HashSet<>();
+        models.add(ClassType.of(PlayerStats.class));
+        models.add(ClassType.of(Country.class));
+        models.add(ClassType.of(SubCountry.class));
+        models.add(ClassType.of(Player.class));
+        models.add(ClassType.of(PlayerRanking.class));
+
+        EntityDictionary dictionary = new EntityDictionary(new HashMap<>());
+
+        models.stream().forEach(dictionary::bindEntity);
+
+        metaDataStore = new MetaDataStore(models, true);
+        metaDataStore.populateEntityDictionary(dictionary);
+
+        DataSource mockDataSource = mock(DataSource.class);
+        //The query engine populates the metadata store with actual tables.
+        new SQLQueryEngine(metaDataStore, new ConnectionDetails(mockDataSource,
+                SQLDialectFactory.getDefaultDialect()));
+
+        lookupTable = new SQLReferenceTable(metaDataStore);
+        dialect = new RSQLFilterDialect(dictionary);
+    }
+
+
+    @Test
+    public void testSinglePredicateNoJoin() throws Exception {
+        FilterExpression expression = parse("overallRating=='Foo'");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertNull(splitExpressions.getOuter());
+        assertEquals(expression, splitExpressions.getInner());
+    }
+
+    @Test
+    public void testSinglePredicateWithJoin() throws Exception {
+        FilterExpression expression = parse("countryUnSeats>3");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertNull(splitExpressions.getInner());
+        assertEquals(expression, splitExpressions.getOuter());
+    }
+
+    @Test
+    public void testSplitByAnd() throws Exception {
+        FilterExpression expression = parse("countryUnSeats>3;overallRating=='Foo'");
+        FilterExpression expectedOuter = parse("countryUnSeats>3");
+        FilterExpression expectedInner = parse("overallRating=='Foo'");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertEquals(expectedOuter, splitExpressions.getOuter());
+        assertEquals(expectedInner, splitExpressions.getInner());
+    }
+
+    @Test
+    public void testSplitByOr() throws Exception {
+        FilterExpression expression = parse("countryUnSeats>3,overallRating=='Foo'");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertEquals(expression, splitExpressions.getOuter());
+        assertNull(splitExpressions.getInner());
+    }
+
+    @Test
+    public void testCompoundSplitByAnd() throws Exception {
+        FilterExpression expression = parse(
+                "(countryUnSeats>3,overallRating=='Foo');(overallRating=='Bar',overallRating=='Blah')");
+
+        FilterExpression expectedOuter = parse("countryUnSeats>3,overallRating=='Foo'");
+        FilterExpression expectedInner = parse("overallRating=='Bar',overallRating=='Blah'");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertEquals(expectedOuter, splitExpressions.getOuter());
+        assertEquals(expectedInner, splitExpressions.getInner());
+    }
+
+    @Test
+    public void testCompoundSplitByOr() throws Exception {
+        FilterExpression expression = parse(
+                "(countryUnSeats>3;overallRating=='Foo'),(overallRating=='Bar';overallRating=='Blah')");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertEquals(expression, splitExpressions.getOuter());
+        assertNull(splitExpressions.getInner());
+    }
+
+    @Test
+    public void testAllOrsWithNoJoins() throws Exception {
+        FilterExpression expression = parse(
+                "(overallRating=='Foobar',overallRating=='Foo'),(overallRating=='Bar',overallRating=='Blah')");
+
+        SubqueryFilterSplitter.SplitFilter splitExpressions =
+                SubqueryFilterSplitter.splitFilter(lookupTable, metaDataStore, expression);
+
+        assertEquals(expression, splitExpressions.getInner());
+        assertNull(splitExpressions.getOuter());
+    }
+
+    private FilterExpression parse(String filter) throws ParseException {
+        return dialect.parse(PLAYER_STATS_TYPE, new HashSet<>(), filter, NO_VERSION);
+    }
+}