Cast-extend aggregation functions (#381)

### What

This PR makes the introspection of aggregation functions take implicit
casts into account.

This enables aggregation functions on domain types, which would
typically only be defined on the base type.

It also enables certain other aggregations for scalar types that weren't
defined before, most notably aggregation functions for `float4` on
CockroachDB, which becomes accessible through implicit casts to
`float8`.

### How

The cast-extension for aggregation functions introspection code is based
on the same for comparison operators.

We also change our notion of "occurring scalar types" to include the
return types aggregation functions. Previously, aggregation functions
could be filtered out if their return type didn't feature in the set of
occurring scalar types, which was unnecessarily strict.

Author: plcplc

changelog.md

@@ -4,6 +4,8 @@
 
 ### Added
 
+- Make aggregation functions available through implicit casts.
+  ([#381](https://github.com/hasura/ndc-postgres/pull/380))
 - Support for introspecting domain types.
   ([#380](https://github.com/hasura/ndc-postgres/pull/380))
 

crates/configuration/src/version3/mod.rs

@@ -151,7 +151,11 @@ pub async fn introspect(
     .instrument(info_span!("Decode introspection result"))
     .await?;
 
-    let scalar_types = occurring_scalar_types(&tables, &args.metadata.native_queries);
+    let scalar_types = occurring_scalar_types(
+        &tables,
+        &args.metadata.native_queries,
+        &args.metadata.aggregate_functions,
+    );
 
     let relevant_comparison_operators =
         filter_comparison_operators(&scalar_types, comparison_operators);
@@ -178,6 +182,7 @@ pub async fn introspect(
 pub fn occurring_scalar_types(
     tables: &metadata::TablesInfo,
     native_queries: &metadata::NativeQueries,
+    aggregate_functions: &metadata::AggregateFunctions,
 ) -> BTreeSet<metadata::ScalarType> {
     let tables_column_types = tables
         .0
@@ -191,6 +196,10 @@ pub fn occurring_scalar_types(
         .0
         .values()
         .flat_map(|v| v.arguments.values().map(|c| &c.r#type));
+    let aggregate_functions_result_types = aggregate_functions
+        .0
+        .values()
+        .flat_map(|x| x.values().map(|agg_fn| agg_fn.return_type.clone()));
 
     tables_column_types
         .chain(native_queries_column_types)
@@ -199,6 +208,7 @@ pub fn occurring_scalar_types(
             metadata::Type::ScalarType(ref t) => Some(t.clone()), // only keep scalar types
             metadata::Type::ArrayType(_) | metadata::Type::CompositeType(_) => None,
         })
+        .chain(aggregate_functions_result_types)
         .collect::<BTreeSet<metadata::ScalarType>>()
 }
 
@@ -240,14 +250,6 @@ fn filter_aggregate_functions(
             .0
             .into_iter()
             .filter(|(typ, _)| scalar_types.contains(typ))
-            .map(|(typ, ops)| {
-                (
-                    typ,
-                    ops.into_iter()
-                        .filter(|(_, op)| scalar_types.contains(&op.return_type))
-                        .collect(),
-                )
-            })
             .collect(),
     )
 }

crates/configuration/src/version3/version3.sql

@@ -354,6 +354,46 @@ WITH
                             -- the purpose of selecting preferred implicit casts.
   ),
 
+  implicit_casts AS
+  (
+    SELECT
+      t_from.type_name as from_type,
+      t_to.type_name as to_type
+    FROM
+      pg_cast
+    INNER JOIN
+      scalar_types
+      AS t_from
+      ON (t_from.type_id = pg_cast.castsource)
+    INNER JOIN
+      scalar_types
+      AS t_to
+      ON (t_to.type_id = pg_cast.casttarget)
+    WHERE
+      pg_cast.castcontext = 'i'
+      AND t_from.type_name != t_to.type_name
+
+      -- This is a good candidate for a configurable option.
+      AND (t_from.type_name, t_to.type_name) NOT IN
+        (
+          -- Ignore other casts that are unlikely to ever be relevant
+          ('bytea', 'geography'),
+          ('bytea', 'geometry'),
+          ('geography', 'bytea'),
+          ('geometry', 'bytea'),
+          ('geometry', 'text'),
+          ('text', 'geometry')
+        )
+    UNION
+    -- Any domain type may be implicitly cast to its base type, even though these casts
+    -- are not declared in `pg_cast`.
+    SELECT
+      domain_types.type_name as from_type,
+      domain_types.base_type as to_type
+    FROM
+      domain_types
+  ),
+
   -- Aggregate functions are recorded across 'pg_proc' and 'pg_aggregate', see
   -- https://www.postgresql.org/docs/current/catalog-pg-proc.html and
   -- https://www.postgresql.org/docs/current/catalog-pg-aggregate.html for
@@ -372,6 +412,13 @@ WITH
     FROM
       pg_catalog.pg_proc AS proc
 
+    INNER JOIN
+      -- Until the schema is made part of our model of types we only consider
+      -- types defined in the public schema.
+      unqualified_schemas_for_types_and_procedures
+      AS q
+      ON (q.schema_id = proc.pronamespace)
+
     -- fetch the argument type name, discarding any unsupported types
     INNER JOIN scalar_types AS arg_type
       ON (arg_type.type_id = proc.proargtypes[0])
@@ -393,6 +440,65 @@ WITH
       AND aggregate.aggnumdirectargs = 0
 
   ),
+  aggregates_cast_extended AS
+  (
+    WITH
+      type_combinations AS
+    (
+      SELECT
+        agg.proc_name AS proc_name,
+        agg.return_type AS return_type,
+        cast1.from_type AS argument_type,
+        true AS argument_casted
+      FROM
+        aggregates
+        AS agg
+      INNER JOIN
+        implicit_casts
+        AS cast1
+        ON (cast1.to_type = agg.argument_type)
+      UNION
+      SELECT
+        agg.proc_name AS proc_name,
+        agg.return_type AS return_type,
+        agg.argument_type AS argument_type,
+        false AS argument_casted
+      FROM
+        aggregates
+        AS agg
+    ),
+
+    preferred_combinations AS
+    (
+      SELECT
+        *,
+        -- CockroachDB does not observe ORDER BY of nested expressions,
+        -- So we cannot use the DISTINCT ON idiom to remove duplicates.
+        -- Therefore we resort to filtering by ordered ROW_NUMBER().
+        ROW_NUMBER()
+          OVER
+          (
+            PARTITION BY
+              proc_name, argument_type
+            ORDER BY
+              -- Prefer uncast argument.
+              NOT argument_casted DESC,
+              -- Arbitrary desperation: Lexical ordering
+              return_type ASC
+          )
+          AS row_number
+      FROM
+        type_combinations
+    )
+    SELECT
+      proc_name,
+      argument_type,
+      return_type
+    FROM
+      preferred_combinations
+    WHERE
+      row_number = 1
+  ),
 
   -- Comparison procedures are any entries in 'pg_proc' that happen to be
   -- binary functions that return booleans. We also require, for the sake of
@@ -508,46 +614,6 @@ WITH
     SELECT * FROM comparison_procedures
   ),
 
-  implicit_casts AS
-  (
-    SELECT
-      t_from.type_name as from_type,
-      t_to.type_name as to_type
-    FROM
-      pg_cast
-    INNER JOIN
-      scalar_types
-      AS t_from
-      ON (t_from.type_id = pg_cast.castsource)
-    INNER JOIN
-      scalar_types
-      AS t_to
-      ON (t_to.type_id = pg_cast.casttarget)
-    WHERE
-      pg_cast.castcontext = 'i'
-      AND t_from.type_name != t_to.type_name
-
-      -- This is a good candidate for a configurable option.
-      AND (t_from.type_name, t_to.type_name) NOT IN
-        (
-          -- Ignore other casts that are unlikely to ever be relevant
-          ('bytea', 'geography'),
-          ('bytea', 'geometry'),
-          ('geography', 'bytea'),
-          ('geometry', 'bytea'),
-          ('geometry', 'text'),
-          ('text', 'geometry')
-        )
-    UNION
-    -- Any domain type may be implicitly cast to its base type, even though these casts
-    -- are not declared in `pg_cast`.
-    SELECT
-      domain_types.type_name as from_type,
-      domain_types.base_type as to_type
-    FROM
-      domain_types
-  ),
-
   -- Some comparison operators are not defined explicitly for every type they would be
   -- valid for, relying instead on implicit casts to extend the types they can apply to.
   --
@@ -1103,7 +1169,7 @@ FROM
           agg.argument_type,
           agg.return_type
         FROM
-          aggregates AS agg
+          aggregates_cast_extended AS agg
         ORDER BY argument_type, proc_name, return_type
       ) AS agg
       GROUP BY agg.argument_type

crates/connectors/ndc-postgres/src/schema.rs

@@ -21,60 +21,64 @@ pub async fn get_schema(
 ) -> Result<models::SchemaResponse, connector::SchemaError> {
     let metadata = &config.metadata;
     let mut scalar_types: BTreeMap<String, models::ScalarType> =
-        configuration::occurring_scalar_types(&metadata.tables, &metadata.native_queries)
-            .iter()
-            .map(|scalar_type| {
-                (
-                    scalar_type.0.clone(),
-                    models::ScalarType {
-                        aggregate_functions: metadata
-                            .aggregate_functions
-                            .0
-                            .get(scalar_type)
-                            .unwrap_or(&BTreeMap::new())
-                            .iter()
-                            .map(|(function_name, function_definition)| {
-                                (
-                                    function_name.clone(),
-                                    models::AggregateFunctionDefinition {
-                                        result_type: models::Type::Named {
-                                            name: function_definition.return_type.0.clone(),
-                                        },
+        configuration::occurring_scalar_types(
+            &metadata.tables,
+            &metadata.native_queries,
+            &metadata.aggregate_functions,
+        )
+        .iter()
+        .map(|scalar_type| {
+            (
+                scalar_type.0.clone(),
+                models::ScalarType {
+                    aggregate_functions: metadata
+                        .aggregate_functions
+                        .0
+                        .get(scalar_type)
+                        .unwrap_or(&BTreeMap::new())
+                        .iter()
+                        .map(|(function_name, function_definition)| {
+                            (
+                                function_name.clone(),
+                                models::AggregateFunctionDefinition {
+                                    result_type: models::Type::Named {
+                                        name: function_definition.return_type.0.clone(),
                                     },
-                                )
-                            })
-                            .collect(),
-                        comparison_operators: metadata
-                            .comparison_operators
-                            .0
-                            .get(scalar_type)
-                            .unwrap_or(&BTreeMap::new())
-                            .iter()
-                            .map(|(op_name, op_def)| {
-                                (
-                                    op_name.clone(),
-                                    match op_def.operator_kind {
-                                        metadata::OperatorKind::Equal => {
-                                            models::ComparisonOperatorDefinition::Equal
+                                },
+                            )
+                        })
+                        .collect(),
+                    comparison_operators: metadata
+                        .comparison_operators
+                        .0
+                        .get(scalar_type)
+                        .unwrap_or(&BTreeMap::new())
+                        .iter()
+                        .map(|(op_name, op_def)| {
+                            (
+                                op_name.clone(),
+                                match op_def.operator_kind {
+                                    metadata::OperatorKind::Equal => {
+                                        models::ComparisonOperatorDefinition::Equal
+                                    }
+                                    metadata::OperatorKind::In => {
+                                        models::ComparisonOperatorDefinition::In
+                                    }
+                                    metadata::OperatorKind::Custom => {
+                                        models::ComparisonOperatorDefinition::Custom {
+                                            argument_type: models::Type::Named {
+                                                name: op_def.argument_type.0.clone(),
+                                            },
                                         }
-                                        metadata::OperatorKind::In => {
-                                            models::ComparisonOperatorDefinition::In
-                                        }
-                                        metadata::OperatorKind::Custom => {
-                                            models::ComparisonOperatorDefinition::Custom {
-                                                argument_type: models::Type::Named {
-                                                    name: op_def.argument_type.0.clone(),
-                                                },
-                                            }
-                                        }
-                                    },
-                                )
-                            })
-                            .collect(),
-                    },
-                )
-            })
-            .collect();
+                                    }
+                                },
+                            )
+                        })
+                        .collect(),
+                },
+            )
+        })
+        .collect();
 
     let collections_by_identifier: BTreeMap<(&str, &str), &str> = metadata
         .tables