Repository: mathesar-foundation/mathesar

Issue:
Implement a more complete variant of the partial summarization algorithm
Algorithm specified in: https://github.com/centerofci/mathesar/issues/1866

Current variant is implemented in #1900. Its shortcoming are that only the first summarization will receive interesting defaults and only if there isn't a prior "select subset of columns" transformation in the pipeline.

A more complete variant should solve these shortcomings.


Discussion:

Resolving patch:
diff --git a/db/queries/base.py b/db/queries/base.py
--- a/db/queries/base.py
+++ b/db/queries/base.py
@@ -34,6 +34,71 @@ def __init__(
         self.name = name
         self.metadata = metadata if metadata else get_empty_metadata()
 
+    def get_input_aliases(self, ix_of_transform):
+        """
+        Each transformation in a DBQuery has its own list of input aliases; this returns it.
+        """
+        initial_aliases = self.initial_aliases
+        if ix_of_transform == 0:
+            return initial_aliases
+        input_aliases = initial_aliases
+        previous_transforms = self.transformations[:ix_of_transform]
+        for transform in previous_transforms:
+            output_aliases = transform.get_output_aliases(input_aliases)
+            input_aliases = output_aliases
+        return input_aliases
+
+    def get_initial_column_by_input_alias(self, ix_of_transform, input_alias):
+        """
+        Retraces the chain of input aliases until it gets to an initial column.
+
+        Returns None if the alias does not originate from an initial column in a way that would
+        preserve a unique constraint. E.g. if it is generated by an aggregation.
+        """
+        initial_col_alias = \
+            self._get_initial_alias_by_input_alias(ix_of_transform, input_alias)
+        if initial_col_alias is None:
+            return None
+        initial_column = \
+            self._get_initial_column_by_initial_column_alias(initial_col_alias)
+        return initial_column
+
+    def _get_initial_alias_by_input_alias(self, ix_of_transform, input_alias):
+        if ix_of_transform == 0:
+            return input_alias
+        transforms = self.transformations[:ix_of_transform]
+        initial_aliases = self.initial_aliases
+        input_aliases = initial_aliases
+        uc_mappings_for_each_transform = [
+            transform.get_unique_constraint_mappings(input_aliases)
+            for transform in transforms
+        ]
+        for uc_mappings in reversed(uc_mappings_for_each_transform):
+            for uc_mapping in uc_mappings:
+                if uc_mapping.output_alias == input_alias:
+                    input_alias = uc_mapping.input_alias
+                    if input_alias is None:
+                        return None
+                    break
+        initial_alias = input_alias
+        return initial_alias
+
+    def _get_initial_column_by_initial_column_alias(self, alias):
+        """
+        Looks up an initial column by initial column alias; no recursive logic.
+        """
+        for initial_column in self.initial_columns:
+            if initial_column.alias == alias:
+                return initial_column
+
+    @property
+    def initial_aliases(self):
+        return [
+            initial_column.alias
+            for initial_column
+            in self.initial_columns
+        ]
+
     # mirrors a method in db.records.operations.select
     def get_records(self, **kwargs):
         # NOTE how through this method you can perform a second batch of
diff --git a/db/transforms/base.py b/db/transforms/base.py
--- a/db/transforms/base.py
+++ b/db/transforms/base.py
@@ -10,13 +10,41 @@
 from db.records.operations import group, relevance, sort as rec_sort
 
 
+class UniqueConstraintMapping:
+    """
+    A unique constraint mapping describes how a transform in a query maps a given input alias to an
+    output alias, in the context of unique constraints.
+
+    When a unique constraint on the input alias would be carried over to the output alias, both
+    `input_alias` and `output_alias` are set. Otherwise, only the `output_alias` is set and
+    `input_alias` is None. This signifies that `output_alias` cannot inherit a unique constraint.
+
+    You can start with an initial column (initial alias) that is unique constrained. As it flows
+    through the query's transformation pipeline it might be transformed and result in a new alias.
+    Some of these new aliases will preserve/inherit that unique constraint, while others will not.
+
+    For example, an output alias generated by Summarize transform's aggregation couldn't inherit a
+    unique constraint.
+
+    All UniqueConstraintMappings for all transforms in a query describe following things:
+        - tells you what the input and output aliases are for a given transform;
+        - whether a given alias is linked to a given initial column (initial-column-linked);
+        - whether a given initial-column-linked alias is unique-constrained when that initial column
+        is unique-constrained.
+    """
+
+    def __init__(self, input_alias, output_alias):
+        self.input_alias = input_alias
+        self.output_alias = output_alias
+
+
 class Transform(ABC):
     type = None
     spec = None
 
     def __init__(
-            self,
-            spec,
+        self,
+        spec,
     ):
         if self.type is None:
             raise ValueError(
@@ -38,6 +66,7 @@ def __eq__(self, other):
             and self.__dict__ == other.__dict__
         )
 
+    # TODO refactor to use `get_unique_constraint_mappings`
     @property
     def map_of_output_alias_to_input_alias(self):
         """
@@ -50,6 +79,28 @@ def map_of_output_alias_to_input_alias(self):
         """
         return dict()
 
+    def get_output_aliases(self, input_aliases):
+        uc_mappings = self.get_unique_constraint_mappings(input_aliases)
+        return [
+            uc_mapping.output_alias
+            for uc_mapping
+            in uc_mappings
+        ]
+
+    def get_unique_constraint_mappings(self, input_aliases):
+        """
+        By default, each input alias maps to an identically named output alias, and its unique
+        constraint, if any, is carried over.
+        """
+        return [
+            UniqueConstraintMapping(
+                input_alias,
+                input_alias,
+            )
+            for input_alias
+            in input_aliases
+        ]
+
 
 class Filter(Transform):
     type = "filter"
@@ -233,6 +284,28 @@ def _get_grouping_column(col_spec):
         )
         return _to_non_executable(executable)
 
+    def get_unique_constraint_mappings(self, _):
+        mappings_that_carry_uniqueness_over = [
+            UniqueConstraintMapping(
+                input_alias=col_spec['input_alias'],
+                output_alias=col_spec['output_alias'],
+            )
+            for col_spec
+            in self._grouping_col_specs
+        ]
+        mappings_that_dont_carry_uniqueness_over = [
+            UniqueConstraintMapping(
+                input_alias=None,
+                output_alias=col_spec['output_alias'],
+            )
+            for col_spec
+            in self._aggregation_col_specs
+        ]
+        return (
+            mappings_that_carry_uniqueness_over
+            + mappings_that_dont_carry_uniqueness_over
+        )
+
     def get_new_with_aliases_added_to_group_by(self, aliases):
         def get_col_spec_from_alias(alias):
             return dict(
@@ -323,7 +396,7 @@ class SelectSubsetOfColumns(Transform):
     type = "select"
 
     def apply_to_relation(self, relation):
-        columns_to_select = self.spec
+        columns_to_select = self._columns_to_select
         if columns_to_select:
             processed_columns_to_select = [
                 _make_sure_column_expression(column)
@@ -335,6 +408,21 @@ def apply_to_relation(self, relation):
         else:
             return relation
 
+    def get_unique_constraint_mappings(self, _):
+        columns_to_select = self._columns_to_select
+        return [
+            UniqueConstraintMapping(
+                column_to_select,
+                column_to_select,
+            )
+            for column_to_select
+            in columns_to_select
+        ]
+
+    @property
+    def _columns_to_select(self):
+        return self.spec
+
 
 def _make_sure_column_expression(input):
     if isinstance(input, str):
diff --git a/db/transforms/operations/finish_specifying.py b/db/transforms/operations/finish_specifying.py
--- a/db/transforms/operations/finish_specifying.py
+++ b/db/transforms/operations/finish_specifying.py
@@ -1,6 +1,6 @@
 from db.tables.operations import select as tables_select
 from db.tables.operations.select import get_joinable_tables
-from db.transforms.base import Summarize, SelectSubsetOfColumns
+from db.transforms.base import Summarize
 from db.columns.operations.select import get_column_from_oid_and_attnum
 
 
@@ -11,13 +11,13 @@ def finish_specifying_summarize_transform(
     Adds missing input aliases to the summarize transform.
 
     Will find input aliases that are not mentioned in the summarize_transform and will add each
-    of them to its group-by set and/or aggregate-on set.
+    of them to its group-by list and/or aggregate-on list.
 
     If the user selected input alias (summarize's base grouping column) is not unique-constrained,
-    will put the unmentioned input aliases in the aggregation set.
+    will put the unmentioned input aliases in the aggregation list.
 
     If the user selected input alias (summarize's base grouping column) is unique-constrained,
-    then it might put at least some input aliases in the grouping set, depending on what
+    then it might put at least some input aliases in the grouping list, depending on what
     _should_group_by returns.
     """
     summarize_transform = db_query.transformations[ix_of_summarize_transform]
@@ -45,29 +45,28 @@ def _split_missing_input_aliases_into_group_and_agg_lists(
     aliases_to_be_added_to_group_by = []
     aliases_to_be_added_to_agg_on = []
     summarize_transform = db_query.transformations[ix_of_summarize_transform]
-    initial_columns_not_in_summarize = \
-        _get_initial_columns_not_in_summarize(db_query, summarize_transform)
-    if not initial_columns_not_in_summarize:
+    missing_input_aliases = \
+        _get_missing_input_aliases(db_query, ix_of_summarize_transform)
+    if not missing_input_aliases:
         # If all input aliases for summarize transform are in the transform's group-by or
-        # aggregate-on sets, there's nothing to do.
+        # aggregate-on lists, there's nothing to do.
         return aliases_to_be_added_to_group_by, aliases_to_be_added_to_agg_on
     # A summarize transform has a base_grouping_column (which is an alias) around
     # which our suggestions will be based.
     base_grouping_alias = summarize_transform.base_grouping_column
     base_grouping_initial_column = \
-        _get_initial_column_by_alias(
-            db_query.initial_columns, base_grouping_alias
+        db_query.get_initial_column_by_input_alias(
+            ix_of_summarize_transform, base_grouping_alias,
         )
-    # We'll always want user-selected alias (base_grouping_column) in the "group-by set";
-    if base_grouping_initial_column in initial_columns_not_in_summarize:
-        aliases_to_be_added_to_group_by.append(base_grouping_alias)
-        initial_columns_not_in_summarize.remove(base_grouping_initial_column)
+    _make_sure_base_grouping_column_will_be_in_group_by_list(
+        base_grouping_alias,
+        missing_input_aliases,
+        aliases_to_be_added_to_group_by,
+    )
     # Most of logic in the rest of method is around whether or not we can add some of the other
-    # input aliases to the "group-by set"; otherwise we'll put them in "aggregate-on set".
+    # missing input aliases to the "group-by list"; otherwise we'll put them in "aggregate-on list".
     can_add_other_aliases_to_group_by = (
-        _is_first_alias_generating_transform(
-            db_query, ix_of_summarize_transform
-        )
+        base_grouping_initial_column is not None
         and _is_initial_column_unique_constrained(
             base_grouping_initial_column, engine, metadata
         )
@@ -79,62 +78,60 @@ def _split_missing_input_aliases_into_group_and_agg_lists(
             )
         oid_of_base_grouping_initial_column = \
             _get_oid_of_initial_column(base_grouping_initial_column)
-        for initial_column in initial_columns_not_in_summarize:
-            if _should_group_by(
-                _get_oid_of_initial_column(initial_column),
-                oid_of_base_grouping_initial_column,
-                oids_of_joinable_tables_with_single_results,
+        for input_alias in missing_input_aliases:
+            initial_column = \
+                db_query.get_initial_column_by_input_alias(
+                    ix_of_summarize_transform, input_alias
+                )
+            if (
+                initial_column is not None
+                and _should_group_by(
+                    _get_oid_of_initial_column(initial_column),
+                    oid_of_base_grouping_initial_column,
+                    oids_of_joinable_tables_with_single_results,
+                )
             ):
-                alias_set_to_add_to = aliases_to_be_added_to_group_by
+                alias_list_to_add_to = aliases_to_be_added_to_group_by
             else:
-                alias_set_to_add_to = aliases_to_be_added_to_agg_on
-            alias_set_to_add_to.append(initial_column.alias)
+                alias_list_to_add_to = aliases_to_be_added_to_agg_on
+            alias_list_to_add_to.append(input_alias)
     else:
-        aliases_to_be_added_to_agg_on = list(
-            initial_column.alias
-            for initial_column
-            in initial_columns_not_in_summarize
-        )
+        aliases_to_be_added_to_agg_on = list(missing_input_aliases)
     return aliases_to_be_added_to_group_by, aliases_to_be_added_to_agg_on
 
 
-def _is_first_alias_generating_transform(db_query, ix_of_summarize_transform):
+def _make_sure_base_grouping_column_will_be_in_group_by_list(
+    base_grouping_alias,
+    missing_input_aliases,
+    aliases_to_be_added_to_group_by,
+):
     """
-    Checks if the transform is the first alias-generating transform. An alias-generating transform
-    means that it itroduces new aliases (columns) to the transform pipeline. We want to know when
-    a given alias-generating transform is the first in the pipeline, because then we can consider
-    its input aliases to be fully described by initial columns, which can be a useful
-    simplification.
+    We'll always want base_grouping_column in the "group-by list";
     """
-    prior_transforms = db_query.transformations[:ix_of_summarize_transform]
-    for prior_transform in prior_transforms:
-        alias_generating_transforms = {Summarize, SelectSubsetOfColumns}
-        is_alias_generating = \
-            type(prior_transform) in alias_generating_transforms
-        if is_alias_generating:
-            return False
-    return True
-
-
-def _get_initial_columns_not_in_summarize(db_query, summarize_transform):
-    initial_columns = db_query.initial_columns
+    if base_grouping_alias in missing_input_aliases:
+        aliases_to_be_added_to_group_by.append(base_grouping_alias)
+        missing_input_aliases.remove(base_grouping_alias)
+
+
+def _get_missing_input_aliases(db_query, ix_of_summarize_transform):
+    """
+    Missing input aliases are those that are not mentioned in a summarize transform's spec. It's
+    the input aliases that we'll be automatically adding into the said spec, so that it can be
+    considered fully specified.
+    """
+    summarize_transform = db_query.transformations[ix_of_summarize_transform]
     group_by_aliases = summarize_transform.grouping_input_aliases
     agg_on_aliases = summarize_transform.aggregation_input_aliases
     aliases_in_summarize = group_by_aliases + agg_on_aliases
+    input_aliases = db_query.get_input_aliases(ix_of_summarize_transform)
     return [
-        initial_column
-        for initial_column in
-        initial_columns
-        if initial_column.alias not in aliases_in_summarize
+        input_alias
+        for input_alias
+        in input_aliases
+        if input_alias not in aliases_in_summarize
     ]
 
 
-def _get_initial_column_by_alias(initial_columns, alias):
-    for initial_column in initial_columns:
-        if initial_column.alias == alias:
-            return initial_column
-
-
 def _should_group_by(
     oid_of_initial_column,
     oid_of_base_grouping_group_by_col,
@@ -186,6 +183,9 @@ def _is_sa_column_unique_constrained(sa_column):
 
 
 def _get_oid_of_initial_column(initial_column):
+    """
+    Using this just because reloid is an obscure term (in our codebase).
+    """
     return initial_column.reloid
 
 


Test patch:
diff --git a/db/tests/transforms/test_finish_specifying_summarization.py b/db/tests/transforms/test_finish_specifying_summarization.py
--- a/db/tests/transforms/test_finish_specifying_summarization.py
+++ b/db/tests/transforms/test_finish_specifying_summarization.py
@@ -243,34 +243,29 @@ def initial_columns(academics_ids):
         ],
         [
             empty_summarize,
-            full_summarize_no_defaults,
+            Summarize(
+                dict(
+                    base_grouping_column=gen_alias.academics.id,
+                    grouping_expressions=[
+                        _gen_grouping_expr(gen_alias.academics.id),
+                        _gen_grouping_expr(gen_alias.academics.name),
+                    ],
+                    aggregation_expressions=[]
+                )
+            ),
             [
                 # should prevent summarization from providing good defaults
                 SelectSubsetOfColumns(
                     [
+                        gen_alias.academics.id,
                         gen_alias.academics.name,
-                        gen_alias.universities.name,
                     ]
                 ),
             ],
             [],
         ],
         [
-            empty_summarize,
-            full_summarize,
-            [],
-            [
-                # should not affect summarization
-                SelectSubsetOfColumns(
-                    [
-                        gen_alias.academics.name,
-                        gen_alias.universities.name,
-                    ]
-                ),
-            ],
-        ],
-        [
-            # partly empty summarization
+            # partial summarization
             Summarize(
                 dict(
                     base_grouping_column=gen_alias.academics.id,
@@ -285,7 +280,7 @@ def initial_columns(academics_ids):
             [],
         ],
         [
-            # partly summarization
+            # partial summarization
             Summarize(
                 dict(
                     base_grouping_column=gen_alias.academics.id,
@@ -314,6 +309,49 @@ def initial_columns(academics_ids):
             [],
             [],
         ],
+        [
+            # partial summarization
+            # like full_summarize, but `gen_alias.academics.name` grouping expr is
+            # after `gen_alias.universities.name`
+            Summarize(
+                dict(
+                    base_grouping_column=gen_alias.academics.id + group_output_alias_suffix,
+                    grouping_expressions=[],
+                    aggregation_expressions=[]
+                )
+            ),
+            Summarize(
+                dict(
+                    base_grouping_column=gen_alias.academics.id + group_output_alias_suffix,
+                    grouping_expressions=[
+                        _gen_grouping_expr(gen_alias.academics.id + group_output_alias_suffix),
+                        _gen_grouping_expr(gen_alias.academics.name + group_output_alias_suffix),
+                    ],
+                    aggregation_expressions=[
+                        _gen_agg_expr(gen_alias.universities.name + agg_output_alias_suffix),
+                        _gen_agg_expr(gen_alias.articles.title + agg_output_alias_suffix),
+                    ]
+                )
+            ),
+            [
+                # a full summarization that groups on 2/3 of columns that would be grouped on by
+                # default.
+                Summarize(
+                    dict(
+                        base_grouping_column=gen_alias.academics.id,
+                        grouping_expressions=[
+                            _gen_grouping_expr(gen_alias.academics.id),
+                            _gen_grouping_expr(gen_alias.academics.name),
+                        ],
+                        aggregation_expressions=[
+                            _gen_agg_expr(gen_alias.universities.name),
+                            _gen_agg_expr(gen_alias.articles.title),
+                        ]
+                    )
+                ),
+            ],
+            [],
+        ],
     ]
 )
 def test_some_transforms(