Cast-extend aggregation functions

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