pytrilogy 0.3.148__cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

trilogy/core/processing/utility.py

@@ -0,0 +1,823 @@
+from dataclasses import dataclass
+from datetime import date, datetime
+from enum import Enum
+from logging import Logger
+from typing import Any, Dict, List, Set, Tuple
+
+import networkx as nx
+
+from trilogy.constants import MagicConstants
+from trilogy.core.enums import (
+    BooleanOperator,
+    DatePart,
+    Derivation,
+    FunctionClass,
+    Granularity,
+    JoinType,
+    Modifier,
+    Purpose,
+)
+from trilogy.core.models.build import (
+    BuildAggregateWrapper,
+    BuildCaseElse,
+    BuildCaseWhen,
+    BuildComparison,
+    BuildConcept,
+    BuildConditional,
+    BuildDatasource,
+    BuildFilterItem,
+    BuildFunction,
+    BuildGrain,
+    BuildParenthetical,
+    BuildSubselectComparison,
+    BuildWindowItem,
+    LooseBuildConceptList,
+)
+from trilogy.core.models.build_environment import BuildEnvironment
+from trilogy.core.models.core import (
+    ArrayType,
+    DataType,
+    ListWrapper,
+    MapType,
+    MapWrapper,
+    NumericType,
+    TraitDataType,
+    TupleWrapper,
+)
+from trilogy.core.models.execute import (
+    CTE,
+    BaseJoin,
+    ConceptPair,
+    QueryDatasource,
+    UnionCTE,
+    UnnestJoin,
+)
+from trilogy.core.statements.author import MultiSelectStatement, SelectStatement
+from trilogy.core.statements.execute import ProcessedQuery
+from trilogy.utility import unique
+
+AGGREGATE_TYPES = (BuildAggregateWrapper,)
+SUBSELECT_TYPES = (BuildSubselectComparison,)
+COMPARISON_TYPES = (BuildComparison,)
+FUNCTION_TYPES = (BuildFunction,)
+PARENTHETICAL_TYPES = (BuildParenthetical,)
+CONDITIONAL_TYPES = (BuildConditional,)
+CONCEPT_TYPES = (BuildConcept,)
+WINDOW_TYPES = (BuildWindowItem,)
+
+
+class NodeType(Enum):
+    CONCEPT = 1
+    NODE = 2
+
+
+@dataclass
+class JoinOrderOutput:
+    right: str
+    type: JoinType
+    keys: dict[str, set[str]]
+    left: str | None = None
+
+    @property
+    def lefts(self):
+        return set(self.keys.keys())
+
+
+@dataclass
+class GroupRequiredResponse:
+    target: BuildGrain
+    upstream: BuildGrain
+    required: bool
+
+
+def find_all_connecting_concepts(g: nx.Graph, ds1: str, ds2: str) -> set[str]:
+    """Find all concepts that connect two datasources"""
+    concepts1 = set(g.neighbors(ds1))
+    concepts2 = set(g.neighbors(ds2))
+    return concepts1 & concepts2
+
+
+def get_connection_keys(
+    all_connections: dict[tuple[str, str], set[str]], left: str, right: str
+) -> set[str]:
+    """Get all concepts that connect two datasources"""
+    lookup = sorted([left, right])
+    key: tuple[str, str] = (lookup[0], lookup[1])
+    return all_connections.get(key, set())
+
+
+def get_join_type(
+    left: str,
+    right: str,
+    partials: dict[str, list[str]],
+    nullables: dict[str, list[str]],
+    all_connecting_keys: set[str],
+) -> JoinType:
+    left_is_partial = any(key in partials.get(left, []) for key in all_connecting_keys)
+    left_is_nullable = any(
+        key in nullables.get(left, []) for key in all_connecting_keys
+    )
+    right_is_partial = any(
+        key in partials.get(right, []) for key in all_connecting_keys
+    )
+    right_is_nullable = any(
+        key in nullables.get(right, []) for key in all_connecting_keys
+    )
+
+    left_complete = not left_is_partial and not left_is_nullable
+    right_complete = not right_is_partial and not right_is_nullable
+
+    if not left_complete and not right_complete:
+        join_type = JoinType.FULL
+    elif not left_complete and right_complete:
+        join_type = JoinType.RIGHT_OUTER
+    elif not right_complete and left_complete:
+        join_type = JoinType.LEFT_OUTER
+    else:
+        join_type = JoinType.INNER
+    return join_type
+
+
+def reduce_join_types(join_types: Set[JoinType]) -> JoinType:
+    final_join_type = JoinType.INNER
+    has_full = any([x == JoinType.FULL for x in join_types])
+
+    if has_full:
+        final_join_type = JoinType.FULL
+        return final_join_type
+    has_left = any([x == JoinType.LEFT_OUTER for x in join_types])
+    has_right = any([x == JoinType.RIGHT_OUTER for x in join_types])
+    if has_left and has_right:
+        final_join_type = JoinType.FULL
+    elif has_left:
+        final_join_type = JoinType.LEFT_OUTER
+    elif has_right:
+        final_join_type = JoinType.RIGHT_OUTER
+
+    return final_join_type
+
+
+def ensure_content_preservation(joins: list[JoinOrderOutput]):
+    # ensure that for a join, if we have prior joins that would
+    # introduce nulls, we are controlling for that
+    for idx, review_join in enumerate(joins):
+        predecessors = joins[:idx]
+        if review_join.type == JoinType.FULL:
+            continue
+        has_prior_left = False
+        has_prior_right = False
+        for pred in predecessors:
+            if (
+                pred.type in (JoinType.LEFT_OUTER, JoinType.FULL)
+                and pred.right in review_join.lefts
+            ):
+                has_prior_left = True
+            if pred.type in (JoinType.RIGHT_OUTER, JoinType.FULL) and any(
+                x in review_join.lefts for x in pred.lefts
+            ):
+                has_prior_right = True
+        if has_prior_left and has_prior_right:
+            target = JoinType.FULL
+        elif has_prior_left:
+            target = (
+                JoinType.LEFT_OUTER
+                if review_join.type != JoinType.RIGHT_OUTER
+                else JoinType.FULL
+            )
+        elif has_prior_right:
+            target = (
+                JoinType.RIGHT_OUTER
+                if review_join.type != JoinType.LEFT_OUTER
+                else JoinType.FULL
+            )
+        else:
+            target = review_join.type
+        if review_join.type != target:
+            review_join.type = target
+            continue
+
+
+def resolve_join_order_v2(
+    g: nx.Graph, partials: dict[str, list[str]], nullables: dict[str, list[str]]
+) -> list[JoinOrderOutput]:
+    datasources = [x for x in g.nodes if x.startswith("ds~")]
+    concepts = [x for x in g.nodes if x.startswith("c~")]
+
+    # Pre-compute all possible connections between datasources
+    all_connections: dict[tuple[str, str], set[str]] = {}
+    for i, ds1 in enumerate(datasources):
+        for ds2 in datasources[i + 1 :]:
+            connecting_concepts = find_all_connecting_concepts(g, ds1, ds2)
+            if connecting_concepts:
+                key = tuple(sorted([ds1, ds2]))
+                all_connections[key] = connecting_concepts
+
+    output: list[JoinOrderOutput] = []
+
+    # create our map of pivots, or common join concepts
+    pivot_map = {
+        concept: [x for x in g.neighbors(concept) if x in datasources]
+        for concept in concepts
+    }
+    pivots = list(
+        sorted(
+            [x for x in pivot_map if len(pivot_map[x]) > 1],
+            key=lambda x: (len(pivot_map[x]), len(x), x),
+        )
+    )
+    solo = [x for x in pivot_map if len(pivot_map[x]) == 1]
+    eligible_left: set[str] = set()
+
+    # while we have pivots, keep joining them in
+    while pivots:
+        next_pivots = [
+            x for x in pivots if any(y in eligible_left for y in pivot_map[x])
+        ]
+        if next_pivots:
+            root = next_pivots[0]
+            pivots = [x for x in pivots if x != root]
+        else:
+            root = pivots.pop(0)
+
+        # sort so less partials is last and eligible lefts are first
+        def score_key(x: str) -> tuple[int, int, str]:
+            base = 1
+            # if it's left, higher weight
+            if x in eligible_left:
+                base += 3
+            # if it has the concept as a partial, lower weight
+            if root in partials.get(x, []):
+                base -= 1
+            if root in nullables.get(x, []):
+                base -= 1
+            return (base, len(x), x)
+
+        # get remaining un-joined datasets
+        to_join = sorted(
+            [x for x in pivot_map[root] if x not in eligible_left], key=score_key
+        )
+        while to_join:
+            # need to sort this to ensure we join on the best match
+            # but check ALL left in case there are non-pivt keys to join on
+            base = sorted([x for x in eligible_left], key=score_key)
+            if not base:
+                new = to_join.pop()
+                eligible_left.add(new)
+                base = [new]
+            right = to_join.pop()
+            # we already joined it
+            # this could happen if the same pivot is shared with multiple DSes
+            if right in eligible_left:
+                continue
+
+            joinkeys: dict[str, set[str]] = {}
+            # sorting puts the best candidate last for pop
+            # so iterate over the reversed list
+            join_types = set()
+
+            for left_candidate in reversed(base):
+                # Get all concepts that connect these two datasources
+                all_connecting_keys = get_connection_keys(
+                    all_connections, left_candidate, right
+                )
+
+                if not all_connecting_keys:
+                    continue
+
+                # Check if we already have this exact set of keys
+                exists = False
+                for _, v in joinkeys.items():
+                    if v == all_connecting_keys:
+                        exists = True
+                if exists:
+                    continue
+
+                join_type = get_join_type(
+                    left_candidate, right, partials, nullables, all_connecting_keys
+                )
+                join_types.add(join_type)
+                joinkeys[left_candidate] = all_connecting_keys
+
+            final_join_type = reduce_join_types(join_types)
+
+            output.append(
+                JoinOrderOutput(
+                    right=right,
+                    type=final_join_type,
+                    keys=joinkeys,
+                )
+            )
+            eligible_left.add(right)
+
+    for concept in solo:
+        for ds in pivot_map[concept]:
+            # if we already have it, skip it
+            if ds in eligible_left:
+                continue
+            # if we haven't had ANY left datasources yet
+            # this needs to become it
+            if not eligible_left:
+                eligible_left.add(ds)
+                continue
+            # otherwise do a full outer join
+            # Try to find if there are any connecting keys with existing left tables
+            best_left = None
+            best_keys: set[str] = set()
+            for existing_left in eligible_left:
+                connecting_keys = get_connection_keys(
+                    all_connections, existing_left, ds
+                )
+                if connecting_keys and len(connecting_keys) > len(best_keys):
+                    best_left = existing_left
+                    best_keys = connecting_keys
+
+            if best_left and best_keys:
+                output.append(
+                    JoinOrderOutput(
+                        left=best_left,
+                        right=ds,
+                        type=JoinType.FULL,
+                        keys={best_left: best_keys},
+                    )
+                )
+            else:
+                output.append(
+                    JoinOrderOutput(
+                        # pick random one to be left
+                        left=list(eligible_left)[0],
+                        right=ds,
+                        type=JoinType.FULL,
+                        keys={},
+                    )
+                )
+            eligible_left.add(ds)
+
+    # only once we have all joins
+    # do we know if some inners need to be left outers
+    ensure_content_preservation(output)
+
+    return output
+
+
+def concept_to_relevant_joins(concepts: list[BuildConcept]) -> List[BuildConcept]:
+    sub_props = LooseBuildConceptList(
+        concepts=[
+            x for x in concepts if x.keys and all([key in concepts for key in x.keys])
+        ]
+    )
+    final = [c for c in concepts if c.address not in sub_props]
+    return unique(final, "address")
+
+
+def padding(x: int) -> str:
+    return "\t" * x
+
+
+def create_log_lambda(prefix: str, depth: int, logger: Logger):
+    pad = padding(depth)
+
+    def log_lambda(msg: str):
+        logger.info(f"{pad}{prefix} {msg}")
+
+    return log_lambda
+
+
+def calculate_graph_relevance(
+    g: nx.DiGraph, subset_nodes: set[str], concepts: set[BuildConcept]
+) -> int:
+    """Calculate the relevance of each node in a graph
+    Relevance is used to prune irrelevant nodes from the graph
+    """
+    relevance = 0
+    for node in g.nodes:
+
+        if node not in subset_nodes:
+            continue
+        if not g.nodes[node]["type"] == NodeType.CONCEPT:
+            continue
+        concept = [x for x in concepts if x.address == node].pop()
+        # debug granularity and derivation
+        # a single row concept can always be crossjoined
+        # therefore a graph with only single row concepts is always relevant
+        if concept.granularity == Granularity.SINGLE_ROW:
+            continue
+        if concept.derivation == Derivation.CONSTANT:
+            continue
+        # if it's an aggregate up to an arbitrary grain, it can be joined in later
+        # and can be ignored in subgraph
+        if concept.purpose == Purpose.METRIC:
+            if not concept.grain:
+                continue
+            if len(concept.grain.components) == 0:
+                continue
+        if concept.grain and len(concept.grain.components) > 0:
+            relevance += 1
+            continue
+        # Added 2023-10-18 since we seemed to be strangely dropping things
+        relevance += 1
+    return relevance
+
+
+def add_node_join_concept(
+    graph: nx.DiGraph,
+    concept: BuildConcept,
+    concept_map: dict[str, BuildConcept],
+    ds_node: str,
+    environment: BuildEnvironment,
+):
+    name = f"c~{concept.address}"
+    graph.add_node(name, type=NodeType.CONCEPT)
+    graph.add_edge(ds_node, name)
+    concept_map[name] = concept
+    for v_address in concept.pseudonyms:
+        v = environment.alias_origin_lookup.get(
+            v_address, environment.concepts[v_address]
+        )
+        if f"c~{v.address}" in graph.nodes:
+            continue
+        if v != concept.address:
+            add_node_join_concept(
+                graph=graph,
+                concept=v,
+                concept_map=concept_map,
+                ds_node=ds_node,
+                environment=environment,
+            )
+
+
+def resolve_instantiated_concept(
+    concept: BuildConcept, datasource: QueryDatasource | BuildDatasource
+) -> BuildConcept:
+    if concept.address in datasource.output_concepts:
+        return concept
+    for k in concept.pseudonyms:
+        if k in datasource.output_concepts:
+            return [x for x in datasource.output_concepts if x.address == k].pop()
+        if any(k in x.pseudonyms for x in datasource.output_concepts):
+            return [x for x in datasource.output_concepts if k in x.pseudonyms].pop()
+    raise SyntaxError(
+        f"Could not find {concept.address} in {datasource.identifier} output {[c.address for c in datasource.output_concepts]}, acceptable synonyms {concept.pseudonyms}"
+    )
+
+
+def reduce_concept_pairs(
+    input: list[ConceptPair], right_source: QueryDatasource | BuildDatasource
+) -> list[ConceptPair]:
+    left_keys = set()
+    right_keys = set()
+    for pair in input:
+        if pair.left.purpose == Purpose.KEY:
+            left_keys.add(pair.left.address)
+        if pair.right.purpose == Purpose.KEY:
+            right_keys.add(pair.right.address)
+    final: list[ConceptPair] = []
+    seen_right_keys = set()
+    for pair in input:
+        if pair.right.address in seen_right_keys:
+            continue
+        if (
+            pair.left.purpose == Purpose.PROPERTY
+            and pair.left.keys
+            and pair.left.keys.issubset(left_keys)
+        ):
+            continue
+        if (
+            pair.right.purpose == Purpose.PROPERTY
+            and pair.right.keys
+            and pair.right.keys.issubset(right_keys)
+        ):
+            continue
+
+        seen_right_keys.add(pair.right.address)
+        final.append(pair)
+    all_keys = set([x.right.address for x in final])
+    if right_source.grain.components and right_source.grain.components.issubset(
+        all_keys
+    ):
+        return [x for x in final if x.right.address in right_source.grain.components]
+
+    return final
+
+
+def get_modifiers(
+    concept: str,
+    join: JoinOrderOutput,
+    ds_node_map: dict[str, QueryDatasource | BuildDatasource],
+):
+    base = []
+
+    if join.right and concept in ds_node_map[join.right].nullable_concepts:
+        base.append(Modifier.NULLABLE)
+    if join.left and concept in ds_node_map[join.left].nullable_concepts:
+        base.append(Modifier.NULLABLE)
+    return list(set(base))
+
+
+def get_node_joins(
+    datasources: List[QueryDatasource | BuildDatasource],
+    environment: BuildEnvironment,
+    # concepts:List[Concept],
+) -> List[BaseJoin]:
+    graph = nx.Graph()
+    partials: dict[str, list[str]] = {}
+    nullables: dict[str, list[str]] = {}
+    ds_node_map: dict[str, QueryDatasource | BuildDatasource] = {}
+    concept_map: dict[str, BuildConcept] = {}
+    for datasource in datasources:
+        ds_node = f"ds~{datasource.identifier}"
+        ds_node_map[ds_node] = datasource
+        graph.add_node(ds_node, type=NodeType.NODE)
+        partials[ds_node] = [f"c~{c.address}" for c in datasource.partial_concepts]
+        nullables[ds_node] = [f"c~{c.address}" for c in datasource.nullable_concepts]
+        for concept in datasource.output_concepts:
+            if concept.address in datasource.hidden_concepts:
+                continue
+            add_node_join_concept(
+                graph=graph,
+                concept=concept,
+                concept_map=concept_map,
+                ds_node=ds_node,
+                environment=environment,
+            )
+
+    joins = resolve_join_order_v2(graph, partials=partials, nullables=nullables)
+    return [
+        BaseJoin(
+            left_datasource=ds_node_map[j.left] if j.left else None,
+            right_datasource=ds_node_map[j.right],
+            join_type=j.type,
+            # preserve empty field for maps
+            concepts=[] if not j.keys else None,
+            concept_pairs=reduce_concept_pairs(
+                [
+                    ConceptPair.model_construct(
+                        left=resolve_instantiated_concept(
+                            concept_map[concept], ds_node_map[k]
+                        ),
+                        right=resolve_instantiated_concept(
+                            concept_map[concept], ds_node_map[j.right]
+                        ),
+                        existing_datasource=ds_node_map[k],
+                        modifiers=get_modifiers(
+                            concept_map[concept].address, j, ds_node_map
+                        ),
+                    )
+                    for k, v in j.keys.items()
+                    for concept in v
+                ],
+                ds_node_map[j.right],
+            ),
+        )
+        for j in joins
+    ]
+
+
+def get_disconnected_components(
+    concept_map: Dict[str, Set[BuildConcept]],
+) -> Tuple[int, List]:
+    """Find if any of the datasources are not linked"""
+    import networkx as nx
+
+    graph = nx.Graph()
+    all_concepts = set()
+    for datasource, concepts in concept_map.items():
+        graph.add_node(datasource, type=NodeType.NODE)
+        for concept in concepts:
+            graph.add_node(concept.address, type=NodeType.CONCEPT)
+            graph.add_edge(datasource, concept.address)
+            all_concepts.add(concept)
+    sub_graphs = list(nx.connected_components(graph))
+    sub_graphs = [
+        x for x in sub_graphs if calculate_graph_relevance(graph, x, all_concepts) > 0
+    ]
+    return len(sub_graphs), sub_graphs
+
+
+def is_scalar_condition(
+    element: (
+        int
+        | str
+        | float
+        | date
+        | datetime
+        | list[Any]
+        | BuildConcept
+        | BuildWindowItem
+        | BuildFilterItem
+        | BuildConditional
+        | BuildComparison
+        | BuildParenthetical
+        | BuildFunction
+        | BuildAggregateWrapper
+        | BuildCaseWhen
+        | BuildCaseElse
+        | MagicConstants
+        | TraitDataType
+        | DataType
+        | MapWrapper[Any, Any]
+        | ArrayType
+        | MapType
+        | NumericType
+        | DatePart
+        | ListWrapper[Any]
+        | TupleWrapper[Any]
+    ),
+    materialized: set[str] | None = None,
+) -> bool:
+    if isinstance(element, PARENTHETICAL_TYPES):
+        return is_scalar_condition(element.content, materialized)
+    elif isinstance(element, SUBSELECT_TYPES):
+        return True
+    elif isinstance(element, COMPARISON_TYPES):
+        return is_scalar_condition(element.left, materialized) and is_scalar_condition(
+            element.right, materialized
+        )
+    elif isinstance(element, FUNCTION_TYPES):
+        if element.operator in FunctionClass.AGGREGATE_FUNCTIONS.value:
+            return False
+        return all([is_scalar_condition(x, materialized) for x in element.arguments])
+    elif isinstance(element, CONCEPT_TYPES):
+        if materialized and element.address in materialized:
+            return True
+        if element.lineage and isinstance(element.lineage, AGGREGATE_TYPES):
+            return is_scalar_condition(element.lineage, materialized)
+        if element.lineage and isinstance(element.lineage, FUNCTION_TYPES):
+            return is_scalar_condition(element.lineage, materialized)
+        return True
+    elif isinstance(element, AGGREGATE_TYPES):
+        return is_scalar_condition(element.function, materialized)
+    elif isinstance(element, CONDITIONAL_TYPES):
+        return is_scalar_condition(element.left, materialized) and is_scalar_condition(
+            element.right, materialized
+        )
+    elif isinstance(element, (BuildCaseWhen,)):
+        return is_scalar_condition(
+            element.comparison, materialized
+        ) and is_scalar_condition(element.expr, materialized)
+    elif isinstance(element, (BuildCaseElse,)):
+        return is_scalar_condition(element.expr, materialized)
+    elif isinstance(element, MagicConstants):
+        return True
+    return True
+
+
+CONDITION_TYPES = (
+    BuildSubselectComparison,
+    BuildComparison,
+    BuildConditional,
+    BuildParenthetical,
+)
+
+
+def decompose_condition(
+    conditional: BuildConditional | BuildComparison | BuildParenthetical,
+) -> list[
+    BuildSubselectComparison | BuildComparison | BuildConditional | BuildParenthetical
+]:
+    chunks: list[
+        BuildSubselectComparison
+        | BuildComparison
+        | BuildConditional
+        | BuildParenthetical
+    ] = []
+    if not isinstance(conditional, BuildConditional):
+        return [conditional]
+    if conditional.operator == BooleanOperator.AND:
+        if not (
+            isinstance(conditional.left, CONDITION_TYPES)
+            and isinstance(
+                conditional.right,
+                CONDITION_TYPES,
+            )
+        ):
+            chunks.append(conditional)
+        else:
+            for val in [conditional.left, conditional.right]:
+                if isinstance(val, BuildConditional):
+                    chunks.extend(decompose_condition(val))
+                else:
+                    chunks.append(val)
+    else:
+        chunks.append(conditional)
+    return chunks
+
+
+def find_nullable_concepts(
+    source_map: Dict[str, set[BuildDatasource | QueryDatasource | UnnestJoin]],
+    datasources: List[BuildDatasource | QueryDatasource],
+    joins: List[BaseJoin | UnnestJoin],
+) -> List[str]:
+    """give a set of datasources and joins, find the concepts
+    that may contain nulls in the output set
+    """
+    nullable_datasources = set()
+    datasource_map = {
+        x.identifier: x
+        for x in datasources
+        if isinstance(x, (BuildDatasource, QueryDatasource))
+    }
+    for join in joins:
+        is_on_nullable_condition = False
+        if not isinstance(join, BaseJoin):
+            continue
+        if not join.concept_pairs:
+            continue
+        for pair in join.concept_pairs:
+            if pair.right.address in [
+                y.address
+                for y in datasource_map[
+                    join.right_datasource.identifier
+                ].nullable_concepts
+            ]:
+                is_on_nullable_condition = True
+                break
+            left_check = (
+                join.left_datasource.identifier
+                if join.left_datasource is not None
+                else pair.existing_datasource.identifier
+            )
+            if pair.left.address in [
+                y.address for y in datasource_map[left_check].nullable_concepts
+            ]:
+                is_on_nullable_condition = True
+                break
+        if is_on_nullable_condition:
+            nullable_datasources.add(datasource_map[join.right_datasource.identifier])
+    final_nullable = set()
+
+    for k, v in source_map.items():
+        local_nullable = [
+            x for x in datasources if k in [v.address for v in x.nullable_concepts]
+        ]
+        nullable_matches = [
+            k in [v.address for v in x.nullable_concepts]
+            for x in datasources
+            if k in [z.address for z in x.output_concepts]
+        ]
+        if all(nullable_matches) and len(nullable_matches) > 0:
+            final_nullable.add(k)
+        all_ds = set([ds for ds in local_nullable]).union(nullable_datasources)
+        if nullable_datasources:
+            if set(v).issubset(all_ds):
+                final_nullable.add(k)
+    return list(sorted(final_nullable))
+
+
+def sort_select_output_processed(
+    cte: CTE | UnionCTE, query: SelectStatement | MultiSelectStatement | ProcessedQuery
+) -> CTE | UnionCTE:
+    if isinstance(query, ProcessedQuery):
+        targets = query.output_columns
+        hidden = query.hidden_columns
+    else:
+        targets = query.output_components
+        hidden = query.hidden_components
+
+    output_addresses = [c.address for c in targets]
+
+    mapping = {x.address: x for x in cte.output_columns}
+
+    new_output: list[BuildConcept] = []
+    for x in targets:
+        if x.address in mapping:
+            new_output.append(mapping[x.address])
+        for oc in cte.output_columns:
+            if x.address in oc.pseudonyms:
+                # create a wrapper BuildConcept to render the pseudonym under the original name
+                if any(x.address == y for y in mapping.keys()):
+                    continue
+                new_output.append(
+                    BuildConcept(
+                        name=x.name,
+                        canonical_name=x.name,
+                        namespace=x.namespace,
+                        pseudonyms={oc.address},
+                        datatype=oc.datatype,
+                        purpose=oc.purpose,
+                        grain=oc.grain,
+                        build_is_aggregate=oc.build_is_aggregate,
+                    )
+                )
+                break
+
+    for oc in cte.output_columns:
+        # add hidden back
+        if oc.address not in output_addresses:
+            new_output.append(oc)
+
+    cte.hidden_concepts = set(
+        [
+            c.address
+            for c in cte.output_columns
+            if (c.address not in targets or c.address in hidden)
+        ]
+    )
+    cte.output_columns = new_output
+    return cte
+
+
+def sort_select_output(
+    cte: CTE | UnionCTE, query: SelectStatement | MultiSelectStatement | ProcessedQuery
+) -> CTE | UnionCTE:
+
+    return sort_select_output_processed(cte, query)