pytrilogy 0.0.1.102__py3-none-any.whl

trilogy/core/processing/node_generators/select_node.py

@@ -0,0 +1,328 @@
+from itertools import combinations
+from typing import List, Optional
+
+from trilogy.core.enums import PurposeLineage
+from trilogy.core.models import Concept, Environment, Grain, LooseConceptList
+from trilogy.core.processing.nodes import (
+    StrategyNode,
+    SelectNode,
+    MergeNode,
+    GroupNode,
+    ConstantNode,
+)
+from trilogy.core.exceptions import NoDatasourceException
+import networkx as nx
+from trilogy.core.graph_models import concept_to_node, datasource_to_node
+from trilogy.constants import logger
+from trilogy.core.processing.utility import padding
+
+LOGGER_PREFIX = "[GEN_SELECT_NODE]"
+
+
+def gen_select_node_from_table(
+    target_concept: Concept,
+    all_concepts: List[Concept],
+    g: nx.DiGraph,
+    environment: Environment,
+    depth: int,
+    target_grain: Grain,
+    accept_partial: bool = False,
+) -> Optional[StrategyNode]:
+    # if we have only constants
+    # we don't need a table
+    # so verify nothing, select node will render
+    all_lcl = LooseConceptList(concepts=all_concepts)
+    if all([c.derivation == PurposeLineage.CONSTANT for c in all_concepts]):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} All concepts {[x.address for x in all_concepts]} are constants, returning constant node"
+        )
+        return ConstantNode(
+            output_concepts=all_concepts,
+            input_concepts=[],
+            environment=environment,
+            g=g,
+            parents=[],
+            depth=depth,
+            # no partial for constants
+            partial_concepts=[],
+            force_group=False,
+        )
+    candidates: dict[str, StrategyNode] = {}
+    scores: dict[str, int] = {}
+    # otherwise, we need to look for a table
+    for datasource in environment.datasources.values():
+        all_found = True
+        for raw_concept in all_concepts:
+            # look for connection to abstract grain
+            req_concept = raw_concept.with_default_grain()
+            # if we don't have a concept in the graph
+            # exit early
+            if concept_to_node(req_concept) not in g.nodes:
+                raise ValueError(concept_to_node(req_concept))
+            try:
+                path = nx.shortest_path(
+                    g,
+                    source=datasource_to_node(datasource),
+                    target=concept_to_node(req_concept),
+                )
+            except nx.NodeNotFound as e:
+                # just to provide better error
+                ncandidates = [
+                    datasource_to_node(datasource),
+                    concept_to_node(req_concept),
+                ]
+                for ncandidate in ncandidates:
+                    try:
+                        g.nodes[ncandidate]
+                    except KeyError:
+                        raise SyntaxError(
+                            "Could not find node for {}".format(ncandidate)
+                        )
+                raise e
+            except nx.exception.NetworkXNoPath:
+                all_found = False
+                break
+            # 2023-10-18 - more strict condition then below
+            # 2023-10-20 - we need this to get through abstract concepts
+            # but we may want to add a filter to avoid using this for anything with lineage
+            # if len(path) != 2:
+            #     all_found = False
+            #     break
+            if len([p for p in path if g.nodes[p]["type"] == "datasource"]) != 1:
+                all_found = False
+                break
+            for node in path:
+                if g.nodes[node]["type"] == "datasource":
+                    continue
+                if g.nodes[node]["concept"].address == raw_concept.address:
+                    continue
+                all_found = False
+                break
+
+        if all_found:
+            partial_concepts = [
+                c.concept
+                for c in datasource.columns
+                if not c.is_complete and c.concept in all_lcl
+            ]
+            partial_lcl = LooseConceptList(concepts=partial_concepts)
+            if not accept_partial and target_concept in partial_lcl:
+                continue
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} target grain is {str(target_grain)}"
+            )
+            if target_grain and target_grain.issubset(datasource.grain):
+
+                if all([x in all_lcl for x in target_grain.components]):
+                    force_group = False
+                # if we are not returning the grain
+                # we have to group
+                else:
+                    logger.info(
+                        f"{padding(depth)}{LOGGER_PREFIX} not all grain components are in output {str(all_lcl)}, group to actual grain"
+                    )
+                    force_group = True
+            elif all([x in all_lcl for x in datasource.grain.components]):
+                logger.info(
+                    f"{padding(depth)}{LOGGER_PREFIX} query output includes all grain components, no reason to group further"
+                )
+                force_group = False
+            else:
+                logger.info(
+                    f"{padding(depth)}{LOGGER_PREFIX} target grain is not subset of datasource grain {datasource.grain}, required to group"
+                )
+                force_group = True
+
+            bcandidate: StrategyNode = SelectNode(
+                input_concepts=[c.concept for c in datasource.columns],
+                output_concepts=all_concepts,
+                environment=environment,
+                g=g,
+                parents=[],
+                depth=depth,
+                partial_concepts=[c for c in all_concepts if c in partial_lcl],
+                accept_partial=accept_partial,
+                datasource=datasource,
+                grain=Grain(components=all_concepts),
+            )
+            # we need to ntest the group node one further
+            if force_group is True:
+                candidate: StrategyNode = GroupNode(
+                    output_concepts=all_concepts,
+                    input_concepts=all_concepts,
+                    environment=environment,
+                    g=g,
+                    parents=[bcandidate],
+                    depth=depth,
+                    partial_concepts=bcandidate.partial_concepts,
+                )
+            else:
+                candidate = bcandidate
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} found select node with {datasource.identifier}, returning {candidate.output_lcl}"
+            )
+            candidates[datasource.identifier] = candidate
+            scores[datasource.identifier] = -len(partial_concepts)
+    if not candidates:
+        return None
+    final = max(candidates, key=lambda x: scores[x])
+    return candidates[final]
+
+
+def gen_select_node(
+    concept: Concept,
+    local_optional: List[Concept],
+    environment: Environment,
+    g,
+    depth: int,
+    accept_partial: bool = False,
+    fail_if_not_found: bool = True,
+    accept_partial_optional: bool = True,
+    target_grain: Grain | None = None,
+) -> StrategyNode | None:
+    all_concepts = [concept] + local_optional
+    all_lcl = LooseConceptList(concepts=all_concepts)
+    materialized_lcl = LooseConceptList(
+        concepts=[
+            x
+            for x in all_concepts
+            if x.address in [z.address for z in environment.materialized_concepts]
+            or x.derivation == PurposeLineage.CONSTANT
+        ]
+    )
+    if not target_grain:
+        target_grain = Grain()
+        for ac in all_concepts:
+            target_grain += ac.grain
+    if materialized_lcl != all_lcl:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} Skipping select node generation for {concept.address} "
+            f" as it + optional (looking for all {all_lcl}) includes non-materialized concepts {all_lcl.difference(materialized_lcl)} vs materialized: {materialized_lcl}"
+        )
+        if fail_if_not_found:
+            raise NoDatasourceException(f"No datasource exists for {concept}")
+        return None
+
+    ds: StrategyNode | None = None
+
+    # attempt to select all concepts from table
+    ds = gen_select_node_from_table(
+        concept,
+        [concept] + local_optional,
+        g=g,
+        environment=environment,
+        depth=depth,
+        accept_partial=accept_partial,
+        target_grain=target_grain,
+    )
+    if ds:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} Found select node with all target concepts, force group is {ds.force_group}, target grain {target_grain}"
+        )
+        return ds
+    # if we cannot find a match
+    parents: List[StrategyNode] = []
+    found: List[Concept] = []
+    logger.info(
+        f"{padding(depth)}{LOGGER_PREFIX} looking for multiple sources that can satisfy"
+    )
+    all_found = False
+    for x in reversed(range(1, len(local_optional) + 1)):
+        if all_found:
+            break
+        for combo in combinations(local_optional, x):
+            if all_found:
+                break
+            # filter to just the original ones we need to get
+            local_combo = [x for x in combo if x not in found]
+            # skip if nothing new in this combo
+            if not local_combo:
+                continue
+            # include core concept as join
+            all_concepts = [concept, *local_combo]
+
+            ds = gen_select_node_from_table(
+                concept,
+                all_concepts,
+                g=g,
+                environment=environment,
+                depth=depth + 1,
+                accept_partial=accept_partial,
+                target_grain=Grain(components=all_concepts),
+            )
+            if ds:
+                logger.info(
+                    f"{padding(depth)}{LOGGER_PREFIX} found a source with {[x.address for x in all_concepts]}"
+                )
+                parents.append(ds)
+                found += [x for x in ds.output_concepts if x != concept]
+                if {x.address for x in found} == {c.address for c in local_optional}:
+                    logger.info(
+                        f"{padding(depth)}{LOGGER_PREFIX} found all optional {[c.address for c in local_optional]}"
+                    )
+                    all_found = True
+    if parents and (all_found or accept_partial_optional):
+        if all_found:
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} found all optional {[c.address for c in local_optional]} via joins"
+            )
+        else:
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} found some optional {[c.address for c in found]}, returning"
+            )
+        all_partial = [
+            c
+            for c in all_concepts
+            if all(
+                [c.address in [x.address for x in p.partial_concepts] for p in parents]
+            )
+        ]
+        force_group = False
+        for candidate in parents:
+            if candidate.grain and not candidate.grain.issubset(target_grain):
+                force_group = True
+        if len(parents) == 1:
+            candidate = parents[0]
+        else:
+            candidate = MergeNode(
+                output_concepts=[concept] + found,
+                input_concepts=[concept] + found,
+                environment=environment,
+                g=g,
+                parents=parents,
+                depth=depth,
+                partial_concepts=all_partial,
+                grain=sum([x.grain for x in parents if x.grain], Grain()),
+            )
+        candidate.depth += 1
+        source_grain = candidate.grain
+        if force_group:
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} datasource grain {source_grain} does not match target grain {target_grain} for select, adding group node"
+            )
+            return GroupNode(
+                output_concepts=candidate.output_concepts,
+                input_concepts=candidate.output_concepts,
+                environment=environment,
+                g=g,
+                parents=[candidate],
+                depth=depth,
+                partial_concepts=candidate.partial_concepts,
+            )
+        return candidate
+
+    if not accept_partial_optional:
+        return None
+    ds = gen_select_node_from_table(
+        concept,
+        [concept],
+        g=g,
+        environment=environment,
+        depth=depth,
+        accept_partial=accept_partial,
+        target_grain=Grain(components=[concept]),
+    )
+
+    if not ds and fail_if_not_found:
+        raise NoDatasourceException(f"No datasource exists for {concept}")
+    return ds

trilogy/core/processing/node_generators/unnest_node.py

@@ -0,0 +1,37 @@
+from typing import List
+
+
+from trilogy.core.models import Concept, Function
+from trilogy.core.processing.nodes import UnnestNode, History
+
+
+def gen_unnest_node(
+    concept: Concept,
+    local_optional: List[Concept],
+    environment,
+    g,
+    depth: int,
+    source_concepts,
+    history: History | None = None,
+) -> UnnestNode | None:
+    arguments = []
+    if isinstance(concept.lineage, Function):
+        arguments = concept.lineage.concept_arguments
+    if arguments or local_optional:
+        parent = source_concepts(
+            mandatory_list=arguments + local_optional,
+            environment=environment,
+            g=g,
+            depth=depth + 1,
+            history=history,
+        )
+        if not parent:
+            return None
+    return UnnestNode(
+        unnest_concept=concept,
+        input_concepts=arguments + local_optional,
+        output_concepts=[concept] + local_optional,
+        environment=environment,
+        g=g,
+        parents=([parent] if (arguments or local_optional) else []),
+    )

trilogy/core/processing/node_generators/window_node.py

@@ -0,0 +1,85 @@
+from typing import List
+
+
+from trilogy.core.models import Concept, WindowItem, Environment
+from trilogy.utility import unique
+from trilogy.core.processing.nodes import (
+    WindowNode,
+)
+from trilogy.core.processing.nodes import MergeNode, History
+
+from trilogy.constants import logger
+from trilogy.core.processing.utility import padding, create_log_lambda
+from trilogy.core.processing.node_generators.common import (
+    gen_enrichment_node,
+    concept_to_relevant_joins,
+)
+
+LOGGER_PREFIX = "[GEN_WINDOW_NODE]"
+
+
+def resolve_window_parent_concepts(concept: Concept) -> List[Concept]:
+    if not isinstance(concept.lineage, WindowItem):
+        raise ValueError
+    base = [concept.lineage.content]
+    if concept.lineage.over:
+        base += concept.lineage.over
+    if concept.lineage.order_by:
+        for item in concept.lineage.order_by:
+            base += [item.expr.output]
+    return unique(base, "address")
+
+
+def gen_window_node(
+    concept: Concept,
+    local_optional: list[Concept],
+    environment: Environment,
+    g,
+    depth: int,
+    source_concepts,
+    history: History | None = None,
+) -> WindowNode | MergeNode | None:
+    parent_concepts = resolve_window_parent_concepts(concept)
+
+    parent_node = source_concepts(
+        mandatory_list=parent_concepts,
+        environment=environment,
+        g=g,
+        depth=depth + 1,
+        history=history,
+    )
+    if not parent_node:
+        logger.info(f"{padding(depth)}{LOGGER_PREFIX} window node parents unresolvable")
+        return None
+    _window_node = WindowNode(
+        input_concepts=parent_concepts,
+        output_concepts=[concept] + parent_concepts,
+        environment=environment,
+        g=g,
+        parents=[
+            parent_node,
+        ],
+    )
+    window_node = MergeNode(
+        parents=[_window_node],
+        environment=environment,
+        g=g,
+        input_concepts=_window_node.input_concepts,
+        output_concepts=_window_node.output_concepts,
+        grain=_window_node.grain,
+        force_group=False,
+    )
+    if not local_optional:
+        return window_node
+    logger.info(f"{padding(depth)}{LOGGER_PREFIX} group node requires enrichment")
+    return gen_enrichment_node(
+        window_node,
+        join_keys=concept_to_relevant_joins(parent_concepts),
+        local_optional=local_optional,
+        environment=environment,
+        g=g,
+        depth=depth,
+        source_concepts=source_concepts,
+        log_lambda=create_log_lambda(LOGGER_PREFIX, depth, logger),
+        history=history,
+    )

trilogy/core/processing/nodes/__init__.py

@@ -0,0 +1,76 @@
+from .filter_node import FilterNode
+from .group_node import GroupNode
+from .merge_node import MergeNode
+from .select_node_v2 import SelectNode, StaticSelectNode, ConstantNode
+from .window_node import WindowNode
+from .base_node import StrategyNode, NodeJoin
+from .unnest_node import UnnestNode
+from pydantic import BaseModel, Field, ConfigDict
+from trilogy.core.models import Concept
+
+
+class History(BaseModel):
+    history: dict[str, StrategyNode | None] = Field(default_factory=dict)
+    started: set[str] = Field(default_factory=set)
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def _concepts_to_lookup(self, search: list[Concept], accept_partial: bool) -> str:
+        return "-".join([c.address for c in search]) + str(accept_partial)
+
+    def search_to_history(
+        self, search: list[Concept], accept_partial: bool, output: StrategyNode | None
+    ):
+        self.history[self._concepts_to_lookup(search, accept_partial)] = output
+
+    def get_history(
+        self,
+        search: list[Concept],
+        accept_partial: bool = False,
+    ) -> StrategyNode | None | bool:
+        return self.history.get(
+            self._concepts_to_lookup(
+                search,
+                accept_partial,
+            ),
+            False,
+        )
+
+    def log_start(
+        self,
+        search: list[Concept],
+        accept_partial: bool = False,
+    ):
+        self.started.add(
+            self._concepts_to_lookup(
+                search,
+                accept_partial,
+            )
+        )
+
+    def check_started(
+        self,
+        search: list[Concept],
+        accept_partial: bool = False,
+    ):
+        return (
+            self._concepts_to_lookup(
+                search,
+                accept_partial,
+            )
+            in self.started
+        )
+
+
+__all__ = [
+    "FilterNode",
+    "GroupNode",
+    "MergeNode",
+    "SelectNode",
+    "StaticSelectNode",
+    "WindowNode",
+    "StrategyNode",
+    "NodeJoin",
+    "ConstantNode",
+    "UnnestNode",
+    "History",
+]