pytrilogy 0.0.1.102__py3-none-any.whl

trilogy/core/processing/utility.py

@@ -0,0 +1,278 @@
+from typing import List, Tuple, Dict, Set
+import networkx as nx
+from trilogy.core.models import (
+    Datasource,
+    JoinType,
+    BaseJoin,
+    Concept,
+    QueryDatasource,
+    LooseConceptList,
+)
+
+from trilogy.core.enums import Purpose, Granularity
+from trilogy.core.constants import CONSTANT_DATASET
+from enum import Enum
+from trilogy.utility import unique
+from collections import defaultdict
+from logging import Logger
+from pydantic import BaseModel
+
+
+class NodeType(Enum):
+    CONCEPT = 1
+    NODE = 2
+
+
+class PathInfo(BaseModel):
+    paths: Dict[str, List[str]]
+    datasource: Datasource
+    reduced_concepts: Set[str]
+    concept_subgraphs: List[List[Concept]]
+
+
+def concept_to_relevant_joins(concepts: list[Concept]) -> List[Concept]:
+    addresses = LooseConceptList(concepts=concepts)
+    sub_props = LooseConceptList(
+        concepts=[
+            x for x in concepts if x.keys and all([key in addresses for key in x.keys])
+        ]
+    )
+    final = [c for c in concepts if c 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[Concept]
+) -> 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()
+
+        # 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 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 resolve_join_order(joins: List[BaseJoin]) -> List[BaseJoin]:
+    available_aliases: set[str] = set()
+    final_joins_pre = [*joins]
+    final_joins = []
+    while final_joins_pre:
+        new_final_joins_pre: List[BaseJoin] = []
+        for join in final_joins_pre:
+            if not available_aliases:
+                final_joins.append(join)
+                available_aliases.add(join.left_datasource.identifier)
+                available_aliases.add(join.right_datasource.identifier)
+            elif join.left_datasource.identifier in available_aliases:
+                # we don't need to join twice
+                # so whatever join we found first, works
+                if join.right_datasource.identifier in available_aliases:
+                    continue
+                final_joins.append(join)
+                available_aliases.add(join.left_datasource.identifier)
+                available_aliases.add(join.right_datasource.identifier)
+            else:
+                new_final_joins_pre.append(join)
+        if len(new_final_joins_pre) == len(final_joins_pre):
+            remaining = [
+                join.left_datasource.identifier for join in new_final_joins_pre
+            ]
+            remaining_right = [
+                join.right_datasource.identifier for join in new_final_joins_pre
+            ]
+            raise SyntaxError(
+                f"did not find any new joins, available {available_aliases} remaining is {remaining + remaining_right} "
+            )
+        final_joins_pre = new_final_joins_pre
+    return final_joins
+
+
+def get_node_joins(
+    datasources: List[QueryDatasource],
+    grain: List[Concept],
+    # concepts:List[Concept],
+) -> List[BaseJoin]:
+    graph = nx.Graph()
+    concepts: List[Concept] = []
+    for datasource in datasources:
+        graph.add_node(datasource.identifier, type=NodeType.NODE)
+        for concept in datasource.output_concepts:
+            # we don't need to join on a concept if all of the keys exist in the grain
+            # if concept.keys and all([x in grain for x in concept.keys]):
+            #     continue
+            concepts.append(concept)
+            graph.add_node(concept.address, type=NodeType.CONCEPT)
+            graph.add_edge(datasource.identifier, concept.address)
+
+    # add edges for every constant to every datasource
+    for datasource in datasources:
+        for concept in datasource.output_concepts:
+            if concept.granularity == Granularity.SINGLE_ROW:
+                for node in graph.nodes:
+                    if graph.nodes[node]["type"] == NodeType.NODE:
+                        graph.add_edge(node, concept.address)
+
+    joins: defaultdict[str, set] = defaultdict(set)
+    identifier_map = {x.identifier: x for x in datasources}
+
+    node_list = sorted(
+        [x for x in graph.nodes if graph.nodes[x]["type"] == NodeType.NODE],
+        # sort so that anything with a partial match on the target is later
+        key=lambda x: len(
+            [x for x in identifier_map[x].partial_concepts if x in grain]
+        ),
+    )
+    for left in node_list:
+        # the constant dataset is a special case
+        # and can never be on the left of a join
+        if left == CONSTANT_DATASET:
+            continue
+
+        for cnode in graph.neighbors(left):
+            if graph.nodes[cnode]["type"] == NodeType.CONCEPT:
+                for right in graph.neighbors(cnode):
+                    # skip concepts
+                    if graph.nodes[right]["type"] == NodeType.CONCEPT:
+                        continue
+                    if left == right:
+                        continue
+                    identifier = [left, right]
+                    joins["-".join(identifier)].add(cnode)
+
+    final_joins_pre: List[BaseJoin] = []
+
+    for key, join_concepts in joins.items():
+        left, right = key.split("-")
+        local_concepts: List[Concept] = unique(
+            [c for c in concepts if c.address in join_concepts], "address"
+        )
+        if all([c.granularity == Granularity.SINGLE_ROW for c in local_concepts]):
+            # for the constant join, make it a full outer join on 1=1
+            join_type = JoinType.FULL
+            local_concepts = []
+        elif any(
+            [
+                c.address in [x.address for x in identifier_map[left].partial_concepts]
+                for c in local_concepts
+            ]
+        ):
+            join_type = JoinType.FULL
+            local_concepts = [
+                c for c in local_concepts if c.granularity != Granularity.SINGLE_ROW
+            ]
+        else:
+            join_type = JoinType.LEFT_OUTER
+            # remove any constants if other join keys exist
+            local_concepts = [
+                c for c in local_concepts if c.granularity != Granularity.SINGLE_ROW
+            ]
+
+            # if concept.keys and all([x in grain for x in concept.keys]):
+            #     continue
+        final_joins_pre.append(
+            BaseJoin(
+                left_datasource=identifier_map[left],
+                right_datasource=identifier_map[right],
+                join_type=join_type,
+                concepts=concept_to_relevant_joins(local_concepts),
+            )
+        )
+    final_joins = resolve_join_order(final_joins_pre)
+
+    # this is extra validation
+    non_single_row_ds = [x for x in datasources if not x.grain.abstract]
+    if len(non_single_row_ds) > 1:
+        for x in datasources:
+            if x.grain.abstract:
+                continue
+            found = False
+            for join in final_joins:
+                if (
+                    join.left_datasource.identifier == x.identifier
+                    or join.right_datasource.identifier == x.identifier
+                ):
+                    found = True
+            if not found:
+                raise SyntaxError(
+                    f"Could not find join for {x.identifier} with output {[c.address for c in x.output_concepts]}, all {[z.identifier for z in datasources]}"
+                )
+    single_row = [x for x in datasources if x.grain.abstract]
+    for x in single_row:
+        for join in final_joins:
+            found = False
+            for join in final_joins:
+                if (
+                    join.left_datasource.identifier == x.identifier
+                    or join.right_datasource.identifier == x.identifier
+                ):
+                    found = True
+            if not found:
+                raise SyntaxError(
+                    f"Could not find join for {x.identifier} with output {[c.address for c in x.output_concepts]}, all {[z.identifier for z in datasources]}"
+                )
+    return final_joins
+
+
+def get_disconnected_components(
+    concept_map: Dict[str, Set[Concept]]
+) -> 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:
+            # TODO: determine if this is the right way to handle things
+            # if concept.derivation in (PurposeLineage.FILTER, PurposeLineage.WINDOW):
+            #     if isinstance(concept.lineage, FilterItem):
+            #         graph.add_node(concept.lineage.content.address, type=NodeType.CONCEPT)
+            #         graph.add_edge(datasource, concept.lineage.content.address)
+            #     if isinstance(concept.lineage, WindowItem):
+            #         graph.add_node(concept.lineage.content.address, type=NodeType.CONCEPT)
+            #         graph.add_edge(datasource, concept.lineage.content.address)
+            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

trilogy/core/query_processor.py

@@ -0,0 +1,331 @@
+from typing import List, Optional, Set, Union, Dict
+
+from trilogy.core.env_processor import generate_graph
+from trilogy.core.graph_models import ReferenceGraph
+from trilogy.core.constants import CONSTANT_DATASET
+from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
+from trilogy.constants import CONFIG, DEFAULT_NAMESPACE
+from trilogy.core.models import (
+    Environment,
+    PersistStatement,
+    SelectStatement,
+    MultiSelectStatement,
+    CTE,
+    Join,
+    UnnestJoin,
+    JoinKey,
+    MaterializedDataset,
+    ProcessedQuery,
+    ProcessedQueryPersist,
+    QueryDatasource,
+    Datasource,
+    BaseJoin,
+    InstantiatedUnnestJoin,
+)
+
+from trilogy.utility import unique
+from collections import defaultdict
+from trilogy.hooks.base_hook import BaseHook
+from trilogy.constants import logger
+from random import shuffle
+from trilogy.core.ergonomics import CTE_NAMES
+from math import ceil
+
+LOGGER_PREFIX = "[QUERY BUILD]"
+
+
+def base_join_to_join(
+    base_join: BaseJoin | UnnestJoin, ctes: List[CTE]
+) -> Join | InstantiatedUnnestJoin:
+    """This function converts joins at the datasource level
+    to joins at the CTE level"""
+    if isinstance(base_join, UnnestJoin):
+        return InstantiatedUnnestJoin(concept=base_join.concept, alias=base_join.alias)
+    if base_join.left_datasource.identifier == base_join.right_datasource.identifier:
+        raise ValueError(f"Joining on same datasource {base_join}")
+    left_ctes = [
+        cte
+        for cte in ctes
+        if (cte.source.full_name == base_join.left_datasource.full_name)
+    ]
+    if not left_ctes:
+        left_ctes = [
+            cte
+            for cte in ctes
+            if (
+                cte.source.datasources[0].full_name
+                == base_join.left_datasource.full_name
+            )
+        ]
+    left_cte = left_ctes[0]
+    right_ctes = [
+        cte
+        for cte in ctes
+        if (cte.source.full_name == base_join.right_datasource.full_name)
+    ]
+    if not right_ctes:
+        right_ctes = [
+            cte
+            for cte in ctes
+            if (
+                cte.source.datasources[0].full_name
+                == base_join.right_datasource.full_name
+            )
+        ]
+    right_cte = right_ctes[0]
+    return Join(
+        left_cte=left_cte,
+        right_cte=right_cte,
+        joinkeys=[JoinKey(concept=concept) for concept in base_join.concepts],
+        jointype=base_join.join_type,
+    )
+
+
+def generate_source_map(
+    query_datasource: QueryDatasource, all_new_ctes: List[CTE]
+) -> Dict[str, str | list[str]]:
+    source_map: Dict[str, list[str]] = defaultdict(list)
+    # now populate anything derived in this level
+    for qdk, qdv in query_datasource.source_map.items():
+        if (
+            qdk not in source_map
+            and len(qdv) == 1
+            and isinstance(list(qdv)[0], UnnestJoin)
+        ):
+            source_map[qdk] = []
+
+        else:
+            for cte in all_new_ctes:
+                output_address = [
+                    x.address
+                    for x in cte.output_columns
+                    if x.address not in [z.address for z in cte.partial_concepts]
+                ]
+                if qdk in output_address:
+                    source_map[qdk].append(cte.name)
+            # now do a pass that accepts partials
+            # TODO: move this into a second loop by first creationg all sub sourcdes
+            # then loop through this
+            for cte in all_new_ctes:
+                output_address = [x.address for x in cte.output_columns]
+                if qdk in output_address:
+                    if qdk not in source_map:
+                        source_map[qdk] = [cte.name]
+        if qdk not in source_map and not qdv:
+            # set source to empty, as it must be derived in this element
+            source_map[qdk] = []
+        if qdk not in source_map:
+            raise ValueError(
+                f"Missing {qdk} in {source_map}, source map {query_datasource.source_map.keys()} "
+            )
+    return {k: "" if not v else v for k, v in source_map.items()}
+
+
+def datasource_to_query_datasource(datasource: Datasource) -> QueryDatasource:
+    sub_select: Dict[str, Set[Union[Datasource, QueryDatasource, UnnestJoin]]] = {
+        **{c.address: {datasource} for c in datasource.concepts},
+    }
+    concepts = [c for c in datasource.concepts]
+    concepts = unique(concepts, "address")
+    return QueryDatasource(
+        output_concepts=concepts,
+        input_concepts=concepts,
+        source_map=sub_select,
+        grain=datasource.grain,
+        datasources=[datasource],
+        joins=[],
+        partial_concepts=[x.concept for x in datasource.columns if not x.is_complete],
+    )
+
+
+def generate_cte_name(full_name: str, name_map: dict[str, str]) -> str:
+    if CONFIG.human_identifiers:
+        if full_name in name_map:
+            return name_map[full_name]
+        suffix = ""
+        idx = len(name_map)
+        if idx >= len(CTE_NAMES):
+            int = ceil(idx / len(CTE_NAMES))
+            suffix = f"_{int}"
+        valid = [x for x in CTE_NAMES if x + suffix not in name_map.values()]
+        shuffle(valid)
+        lookup = valid[0]
+        new_name = f"{lookup}{suffix}"
+        name_map[full_name] = new_name
+        return new_name
+    else:
+        return full_name.replace("<", "").replace(">", "").replace(",", "_")
+
+
+def datasource_to_ctes(
+    query_datasource: QueryDatasource, name_map: dict[str, str]
+) -> List[CTE]:
+    output: List[CTE] = []
+    parents: list[CTE] = []
+    if len(query_datasource.datasources) > 1 or any(
+        [isinstance(x, QueryDatasource) for x in query_datasource.datasources]
+    ):
+        all_new_ctes: List[CTE] = []
+        for datasource in query_datasource.datasources:
+            if isinstance(datasource, QueryDatasource):
+                sub_datasource = datasource
+            else:
+                sub_datasource = datasource_to_query_datasource(datasource)
+
+            sub_cte = datasource_to_ctes(sub_datasource, name_map)
+            parents += sub_cte
+            all_new_ctes += sub_cte
+        source_map = generate_source_map(query_datasource, all_new_ctes)
+    else:
+        # source is the first datasource of the query datasource
+        source = query_datasource.datasources[0]
+        # this is required to ensure that constant datasets
+        # render properly on initial access; since they have
+        # no actual source
+        if source.full_name == DEFAULT_NAMESPACE + "_" + CONSTANT_DATASET:
+            source_map = {k: "" for k in query_datasource.source_map}
+        else:
+            source_map = {
+                k: "" if not v else source.full_name
+                for k, v in query_datasource.source_map.items()
+            }
+    human_id = generate_cte_name(query_datasource.full_name, name_map)
+    cte = CTE(
+        name=human_id,
+        source=query_datasource,
+        # output columns are what are selected/grouped by
+        output_columns=[
+            c.with_grain(query_datasource.grain)
+            for c in query_datasource.output_concepts
+        ],
+        source_map=source_map,
+        # related columns include all referenced columns, such as filtering
+        joins=[
+            x
+            for x in [
+                base_join_to_join(join, parents) for join in query_datasource.joins
+            ]
+            if x
+        ],
+        grain=query_datasource.grain,
+        group_to_grain=query_datasource.group_required,
+        # we restrict parent_ctes to one level
+        # as this set is used as the base for rendering the query
+        parent_ctes=parents,
+        condition=query_datasource.condition,
+        partial_concepts=query_datasource.partial_concepts,
+        join_derived_concepts=query_datasource.join_derived_concepts,
+    )
+    if cte.grain != query_datasource.grain:
+        raise ValueError("Grain was corrupted in CTE generation")
+    for x in cte.output_columns:
+        if x.address not in cte.source_map:
+            raise ValueError(
+                f"Missing {x.address} in {cte.source_map}, source map {cte.source.source_map.keys()} "
+            )
+
+    output.append(cte)
+    return output
+
+
+def get_query_datasources(
+    environment: Environment,
+    statement: SelectStatement | MultiSelectStatement,
+    graph: Optional[ReferenceGraph] = None,
+    hooks: Optional[List[BaseHook]] = None,
+) -> QueryDatasource:
+    graph = graph or generate_graph(environment)
+    logger.info(
+        f"{LOGGER_PREFIX} getting source datasource for query with output {[str(c) for c in statement.output_components]}"
+    )
+    if not statement.output_components:
+        raise ValueError(f"Statement has no output components {statement}")
+    ds = source_query_concepts(
+        statement.output_components, environment=environment, g=graph
+    )
+    if hooks:
+        for hook in hooks:
+            hook.process_root_strategy_node(ds)
+    final_qds = ds.resolve()
+    return final_qds
+
+
+def flatten_ctes(input: CTE) -> list[CTE]:
+    output = [input]
+    for cte in input.parent_ctes:
+        output += flatten_ctes(cte)
+    return output
+
+
+def process_auto(
+    environment: Environment,
+    statement: PersistStatement | SelectStatement,
+    hooks: List[BaseHook] | None = None,
+):
+    if isinstance(statement, PersistStatement):
+        return process_persist(environment, statement, hooks)
+    elif isinstance(statement, SelectStatement):
+        return process_query(environment, statement, hooks)
+    raise ValueError(f"Do not know how to process {type(statement)}")
+
+
+def process_persist(
+    environment: Environment,
+    statement: PersistStatement,
+    hooks: List[BaseHook] | None = None,
+) -> ProcessedQueryPersist:
+    select = process_query(
+        environment=environment, statement=statement.select, hooks=hooks
+    )
+
+    # build our object to return
+    arg_dict = {k: v for k, v in select.__dict__.items()}
+    return ProcessedQueryPersist(
+        **arg_dict,
+        output_to=MaterializedDataset(address=statement.address),
+        datasource=statement.datasource,
+    )
+
+
+def process_query(
+    environment: Environment,
+    statement: SelectStatement | MultiSelectStatement,
+    hooks: List[BaseHook] | None = None,
+) -> ProcessedQuery:
+    hooks = hooks or []
+    graph = generate_graph(environment)
+    root_datasource = get_query_datasources(
+        environment=environment, graph=graph, statement=statement, hooks=hooks
+    )
+    for hook in hooks:
+        hook.process_root_datasource(root_datasource)
+    # this should always return 1 - TODO, refactor
+    root_cte = datasource_to_ctes(root_datasource, environment.cte_name_map)[0]
+    for hook in hooks:
+        hook.process_root_cte(root_cte)
+    raw_ctes: List[CTE] = list(reversed(flatten_ctes(root_cte)))
+    seen = dict()
+    # we can have duplicate CTEs at this point
+    # so merge them together
+    for cte in raw_ctes:
+        if cte.name not in seen:
+            seen[cte.name] = cte
+        else:
+            # merge them up
+            seen[cte.name] = seen[cte.name] + cte
+    for cte in raw_ctes:
+        cte.parent_ctes = [seen[x.name] for x in cte.parent_ctes]
+    final_ctes: List[CTE] = list(seen.values())
+
+    return ProcessedQuery(
+        order_by=statement.order_by,
+        grain=statement.grain,
+        limit=statement.limit,
+        where_clause=statement.where_clause,
+        output_columns=statement.output_components,
+        ctes=final_ctes,
+        base=root_cte,
+        # we no longer do any joins at final level, this should always happen in parent CTEs
+        joins=[],
+        hidden_columns=[x for x in statement.hidden_components],
+    )