ES Agg pushdown v2

pg_es_fdw/__init__.py

@@ -319,6 +319,12 @@ def _handle_aggregation_response(self, query, response, aggs, group_clauses):
             result = {}
 
             for agg_name in aggs:
+                if agg_name == "count.*":
+                    # COUNT(*) is a special case, since it doesn't have a
+                    # corresponding aggregation primitice in ES
+                    result[agg_name] = response["hits"]["total"]["value"]
+                    continue
+
                 result[agg_name] = response["aggregations"][agg_name]["value"]
             yield result
         else:
@@ -331,6 +337,12 @@ def _handle_aggregation_response(self, query, response, aggs, group_clauses):
 
                     if aggs is not None:
                         for agg_name in aggs:
+                            if agg_name == "count.*":
+                                # In general case with GROUP BY clauses COUNT(*)
+                                # is taken from the bucket's doc_count field
+                                result[agg_name] = bucket["doc_count"]
+                                continue
+
                             result[agg_name] = bucket[agg_name]["value"]
 
                     yield result

pg_es_fdw/_es_query.py

@@ -18,6 +18,7 @@
     "min": "min",
     "sum": "sum",
     "count": "value_count",
+    "count.*": None  # not mapped to a particular function
 }
 
 
@@ -43,7 +44,7 @@ def _base_qual_to_es(col, op, value, column_map=None):
         return {"bool": {"must_not": {"term": {col: value}}}}
 
     if op == "~~":
-        return {"match": {col: value.replace("%", "*")}}
+        return {"wildcard": {col: value.replace("%", "*")}}
 
     # For unknown operators, get everything
     return {"match_all": {}}
@@ -82,6 +83,18 @@ def quals_to_es(
     """Convert a list of Multicorn quals to an ElasticSearch query"""
     ignore_columns = ignore_columns or []
 
+    query = {
+        "query": {
+            "bool": {
+                "must": [
+                    _qual_to_es(q, column_map)
+                    for q in quals
+                    if q.field_name not in ignore_columns
+                ]
+            }
+        }
+    }
+
     # Aggregation/grouping queries
     if aggs is not None:
         aggs_query = {
@@ -91,10 +104,18 @@ def quals_to_es(
                 }
             }
             for agg_name, agg_props in aggs.items()
+            if agg_name != "count.*"
         }
 
         if group_clauses is None:
-            return {"aggs": aggs_query}
+            if "count.*" in aggs:
+                # There is no particular COUNT(*) equivalent in ES, instead
+                # for plain aggregations (e.g. no grouping statements), we need
+                # to enable the track_total_hits option in order to get an
+                # accuate number of matched docs.
+                query["track_total_hits"] = True
+
+            query["aggs"] = aggs_query
 
     if group_clauses is not None:
         group_query = {
@@ -111,17 +132,6 @@ def quals_to_es(
         if aggs is not None:
             group_query["group_buckets"]["aggregations"] = aggs_query
 
-        return {"aggs": group_query}
+        query["aggs"] = group_query
 
-    # Regular query
-    return {
-        "query": {
-            "bool": {
-                "must": [
-                    _qual_to_es(q, column_map)
-                    for q in quals
-                    if q.field_name not in ignore_columns
-                ]
-            }
-        }
-    }
+    return query