Add tests for upperrel qual pushdown

splitgraph · Jan 14, 2022 · ec5fd86 · ec5fd86
1 parent e3aab81
commit ec5fd86
Show file tree

Hide file tree

Showing 6 changed files with 314 additions and 7 deletions.
diff --git a/...in_combinations/account_join_sub_aggs.yml → ...in_combinations/account_join_sub_aggs.yml b/...in_combinations/account_join_sub_aggs.yml → ...in_combinations/account_join_sub_aggs.yml
diff --git a/...ggregations_bare/account_count_by_age.yml → ...ggregations_bare/account_count_by_age.yml b/...ggregations_bare/account_count_by_age.yml → ...ggregations_bare/account_count_by_age.yml
diff --git a/...tion_pushdown/test_grouping_and_aggregations_filtering/min_balance_state_age_filtered.yml b/...tion_pushdown/test_grouping_and_aggregations_filtering/min_balance_state_age_filtered.yml
@@ -0,0 +1,180 @@
+- !!python/tuple
+  - AK
+  - 25
+  - 31138.0
+- !!python/tuple
+  - AK
+  - 35
+  - 37707.0
+- !!python/tuple
+  - AR
+  - 25
+  - 23545.0
+- !!python/tuple
+  - AZ
+  - 35
+  - 48868.0
+- !!python/tuple
+  - CO
+  - 25
+  - 40891.0
+- !!python/tuple
+  - CT
+  - 25
+  - 33474.0
+- !!python/tuple
+  - DC
+  - 25
+  - 11329.0
+- !!python/tuple
+  - DE
+  - 25
+  - 2643.0
+- !!python/tuple
+  - FL
+  - 35
+  - 6302.0
+- !!python/tuple
+  - GA
+  - 35
+  - 16782.0
+- !!python/tuple
+  - IA
+  - 35
+  - 14485.0
+- !!python/tuple
+  - ID
+  - 25
+  - 36776.0
+- !!python/tuple
+  - ID
+  - 35
+  - 20455.0
+- !!python/tuple
+  - KS
+  - 35
+  - 32852.0
+- !!python/tuple
+  - KY
+  - 25
+  - 35596.0
+- !!python/tuple
+  - KY
+  - 35
+  - 26581.0
+- !!python/tuple
+  - LA
+  - 25
+  - 13265.0
+- !!python/tuple
+  - MA
+  - 35
+  - 43370.0
+- !!python/tuple
+  - MD
+  - 25
+  - 47657.0
+- !!python/tuple
+  - MD
+  - 35
+  - 2028.0
+- !!python/tuple
+  - ME
+  - 25
+  - 21153.0
+- !!python/tuple
+  - ME
+  - 35
+  - 22752.0
+- !!python/tuple
+  - MI
+  - 25
+  - 31198.0
+- !!python/tuple
+  - MN
+  - 25
+  - 27583.0
+- !!python/tuple
+  - MN
+  - 35
+  - 28044.0
+- !!python/tuple
+  - MO
+  - 35
+  - 39063.0
+- !!python/tuple
+  - MS
+  - 25
+  - 49567.0
+- !!python/tuple
+  - MS
+  - 35
+  - 34708.0
+- !!python/tuple
+  - NE
+  - 25
+  - 46091.0
+- !!python/tuple
+  - NH
+  - 35
+  - 28647.0
+- !!python/tuple
+  - NM
+  - 25
+  - 23860.0
+- !!python/tuple
+  - NV
+  - 25
+  - 17127.0
+- !!python/tuple
+  - OK
+  - 25
+  - 8036.0
+- !!python/tuple
+  - OK
+  - 35
+  - 22575.0
+- !!python/tuple
+  - OR
+  - 25
+  - 13024.0
+- !!python/tuple
+  - OR
+  - 35
+  - 1481.0
+- !!python/tuple
+  - SC
+  - 35
+  - 5848.0
+- !!python/tuple
+  - SD
+  - 35
+  - 31613.0
+- !!python/tuple
+  - TN
+  - 35
+  - 11402.0
+- !!python/tuple
+  - TX
+  - 35
+  - 31055.0
+- !!python/tuple
+  - UT
+  - 25
+  - 38695.0
+- !!python/tuple
+  - UT
+  - 35
+  - 11422.0
+- !!python/tuple
+  - VA
+  - 25
+  - 45973.0
+- !!python/tuple
+  - WI
+  - 35
+  - 4994.0
+- !!python/tuple
+  - WY
+  - 35
+  - 23285.0
diff --git a/...clauses_only/account_genders_and_ages.yml → ...ping_clauses/account_genders_and_ages.yml b/...clauses_only/account_genders_and_ages.yml → ...ping_clauses/account_genders_and_ages.yml
diff --git a/..._gropuing_clauses_only/account_states.yml → ...imple_grouping_clauses/account_states.yml b/..._gropuing_clauses_only/account_states.yml → ...imple_grouping_clauses/account_states.yml
diff --git a/...aph/commands/test_aggregation_pushdown.py → .../commands/test_es_aggregation_pushdown.py b/...aph/commands/test_aggregation_pushdown.py → .../commands/test_es_aggregation_pushdown.py
@@ -26,11 +26,10 @@ def _extract_queries_from_explain(result):
 
 
 _bare_sequential_scan = {"query": {"bool": {"must": []}}}
-_bare_filtering_query = {"query": {"bool": {"must": [{"range": {"age": {"gt": 30}}}]}}}
 
 
 @pytest.mark.mounting
-def test_elasticsearch_aggregation_functions_only(local_engine_empty):
+def test_simple_aggregation_functions(local_engine_empty):
     _mount_elasticsearch()
 
     query = """
@@ -62,7 +61,38 @@ def test_elasticsearch_aggregation_functions_only(local_engine_empty):
 
 
 @pytest.mark.mounting
-def test_elasticsearch_gropuing_clauses_only(snapshot, local_engine_empty):
+def test_simple_aggregation_functions_filtering(local_engine_empty):
+    _mount_elasticsearch()
+    query = """
+    SELECT avg(age), max(balance)
+    FROM es.account
+    WHERE balance > 20000 AND age < 30
+    """
+
+    # Ensure query is going to be aggregated on the foreign server
+    result = get_engine().run_sql("EXPLAIN " + query)
+    assert _extract_queries_from_explain(result)[0] == {
+        "query": {
+            "bool": {
+                "must": [{"range": {"balance": {"gt": "20000"}}}, {"range": {"age": {"lt": 30}}}]
+            }
+        },
+        "aggs": {
+            "avg.age": {"avg": {"field": "age"}},
+            "max.balance": {"max": {"field": "balance"}},
+        },
+    }
+
+    # Ensure results are correct
+    result = get_engine().run_sql(query)
+    assert len(result) == 1
+
+    # Assert aggregation result
+    assert result[0] == (Decimal("24.439862542955325"), 49795.0)
+
+
+@pytest.mark.mounting
+def test_simple_grouping_clauses(snapshot, local_engine_empty):
     _mount_elasticsearch()
 
     # Single column grouping
@@ -115,7 +145,50 @@ def test_elasticsearch_gropuing_clauses_only(snapshot, local_engine_empty):
 
 
 @pytest.mark.mounting
-def test_elasticsearch_grouping_and_aggregations_bare(snapshot, local_engine_empty):
+def test_simple_grouping_clauses_filtering(snapshot, local_engine_empty):
+    _mount_elasticsearch()
+
+    # Single column grouping
+    query = "SELECT state, gender FROM es.account WHERE state IN ('TX', 'WA', 'CO') GROUP BY state, gender"
+
+    # Ensure grouping is going to be pushed down
+    result = get_engine().run_sql("EXPLAIN " + query)
+    assert _extract_queries_from_explain(result)[0] == {
+        "query": {
+            "bool": {
+                "must": [
+                    {
+                        "bool": {
+                            "should": [
+                                {"term": {"state": "TX"}},
+                                {"term": {"state": "WA"}},
+                                {"term": {"state": "CO"}},
+                            ]
+                        }
+                    }
+                ]
+            }
+        },
+        "aggs": {
+            "group_buckets": {
+                "composite": {
+                    "sources": [
+                        {"state": {"terms": {"field": "state"}}},
+                        {"gender": {"terms": {"field": "gender"}}},
+                    ],
+                    "size": 5,
+                }
+            }
+        },
+    }
+
+    # Ensure results are correct
+    result = get_engine().run_sql(query)
+    assert result == [("CO", "F"), ("CO", "M"), ("TX", "F"), ("TX", "M"), ("WA", "F"), ("WA", "M")]
+
+
+@pytest.mark.mounting
+def test_grouping_and_aggregations_bare(snapshot, local_engine_empty):
     _mount_elasticsearch()
 
     # Aggregations functions and grouping bare combination
@@ -174,7 +247,61 @@ def test_elasticsearch_grouping_and_aggregations_bare(snapshot, local_engine_emp
 
 
 @pytest.mark.mounting
-def test_elasticsearch_agg_subquery_pushdown(local_engine_empty):
+def test_grouping_and_aggregations_filtering(snapshot, local_engine_empty):
+    _mount_elasticsearch()
+
+    # Aggregations functions and grouping bare combination
+    query = """
+    SELECT state, age, min(balance)
+    FROM es.account
+    WHERE gender = 'M' AND age = ANY(ARRAY[25, 35])
+    GROUP BY state, age
+    """
+
+    # Ensure query is going to be pushed down
+    result = get_engine().run_sql("EXPLAIN " + query)
+    assert _extract_queries_from_explain(result)[0] == {
+        "query": {
+            "bool": {
+                "must": [
+                    {
+                        "bool": {
+                            "should": [
+                                {"term": {"age": 25}},
+                                {"term": {"age": 35}},
+                            ]
+                        }
+                    },
+                    {"term": {"gender": "M"}},
+                ]
+            }
+        },
+        "aggs": {
+            "group_buckets": {
+                "composite": {
+                    "sources": [
+                        {"state": {"terms": {"field": "state"}}},
+                        {"age": {"terms": {"field": "age"}}},
+                    ],
+                    "size": 5,
+                },
+                "aggregations": {
+                    "min.balance": {"min": {"field": "balance"}},
+                },
+            }
+        },
+    }
+
+    # Ensure results are correct
+    result = get_engine().run_sql(query)
+    assert len(result) == 45
+
+    # Assert aggregation result
+    snapshot.assert_match(yaml.dump(result), "min_balance_state_age_filtered.yml")
+
+
+@pytest.mark.mounting
+def test_agg_subquery_pushdown(local_engine_empty):
     """
     Most of the magic in these examples is coming from PG, not our Multicorn code
     (i.e. discarding redundant targets from subqueries).
@@ -284,7 +411,7 @@ def test_elasticsearch_agg_subquery_pushdown(local_engine_empty):
 
 
 @pytest.mark.mounting
-def test_elasticsearch_aggregations_join_combinations(snapshot, local_engine_empty):
+def test_aggregations_join_combinations(snapshot, local_engine_empty):
     # Sub-aggregations in a join are pushed down
     query = """
     SELECT t1.*, t2.min FROM (
@@ -352,7 +479,7 @@ def test_elasticsearch_aggregations_join_combinations(snapshot, local_engine_emp
 
 
 @pytest.mark.mounting
-def test_elasticsearch_not_pushed_down(local_engine_empty):
+def test_not_pushed_down(local_engine_empty):
     _mount_elasticsearch()
 
     # COUNT STAR is not going to be pushed down