pytrilogy 0.0.1.102__py3-none-any.whl

trilogy/core/processing/nodes/base_node.py

@@ -0,0 +1,251 @@
+from typing import List, Optional, Sequence
+from collections import defaultdict
+
+from trilogy.core.models import (
+    Grain,
+    QueryDatasource,
+    SourceType,
+    Concept,
+    Environment,
+    Conditional,
+    UnnestJoin,
+    Datasource,
+    Comparison,
+    Parenthetical,
+    LooseConceptList,
+)
+from trilogy.core.enums import Purpose, JoinType, PurposeLineage, Granularity
+from trilogy.utility import unique
+from dataclasses import dataclass
+
+
+def concept_list_to_grain(
+    inputs: List[Concept], parent_sources: Sequence[QueryDatasource | Datasource]
+) -> Grain:
+    candidates = [
+        c
+        for c in inputs
+        if c.purpose == Purpose.KEY and c.granularity != Granularity.SINGLE_ROW
+    ]
+    for x in inputs:
+        if x.granularity == Granularity.SINGLE_ROW:
+            continue
+        if x.purpose == Purpose.PROPERTY and not any(
+            [key in candidates for key in (x.keys or [])]
+        ):
+            candidates.append(x)
+        elif x.purpose == Purpose.CONSTANT:
+            candidates.append(x)
+        elif x.purpose == Purpose.METRIC:
+            # metrics that were previously calculated must be included in grain
+            if any([x in parent.output_concepts for parent in parent_sources]):
+                candidates.append(x)
+
+    return Grain(components=candidates)
+
+
+def resolve_concept_map(
+    inputs: List[QueryDatasource],
+    targets: List[Concept],
+    inherited_inputs: List[Concept],
+    full_joins: List[Concept] | None = None,
+) -> dict[str, set[Datasource | QueryDatasource | UnnestJoin]]:
+    targets = targets or []
+    concept_map: dict[str, set[Datasource | QueryDatasource | UnnestJoin]] = (
+        defaultdict(set)
+    )
+    full_addresses = {c.address for c in full_joins} if full_joins else set()
+    for input in inputs:
+        for concept in input.output_concepts:
+            if concept.address not in input.non_partial_concept_addresses:
+                continue
+            if concept.address not in [t.address for t in inherited_inputs]:
+                continue
+            if concept.address in full_addresses:
+                concept_map[concept.address].add(input)
+            elif concept.address not in concept_map:
+                concept_map[concept.address].add(input)
+
+    # second loop, include partials
+    for input in inputs:
+        for concept in input.output_concepts:
+            if concept.address not in [t.address for t in inherited_inputs]:
+                continue
+            if len(concept_map.get(concept.address, [])) == 0:
+                concept_map[concept.address].add(input)
+    # this adds our new derived metrics, which are not created in this CTE
+    for target in targets:
+        if target not in inherited_inputs:
+            # an empty source means it is defined in this CTE
+            concept_map[target.address] = set()
+    return concept_map
+
+
+def get_all_parent_partial(all_concepts: List[Concept], parents: List["StrategyNode"]):
+    return [
+        c
+        for c in all_concepts
+        if len([c.address in [x.address for x in p.partial_concepts] for p in parents])
+        >= 1
+        and all([c.address in [x.address for x in p.partial_concepts] for p in parents])
+    ]
+
+
+class StrategyNode:
+    source_type = SourceType.ABSTRACT
+
+    def __init__(
+        self,
+        input_concepts: List[Concept],
+        output_concepts: List[Concept],
+        environment: Environment,
+        g,
+        whole_grain: bool = False,
+        parents: List["StrategyNode"] | None = None,
+        partial_concepts: List[Concept] | None = None,
+        depth: int = 0,
+        conditions: Conditional | Comparison | Parenthetical | None = None,
+        force_group: bool | None = None,
+        grain: Optional[Grain] = None,
+    ):
+        self.input_concepts: List[Concept] = (
+            unique(input_concepts, "address") if input_concepts else []
+        )
+        self.input_lcl = LooseConceptList(concepts=self.input_concepts)
+        self.output_concepts: List[Concept] = unique(output_concepts, "address")
+        self.output_lcl = LooseConceptList(concepts=self.output_concepts)
+
+        self.environment = environment
+        self.g = g
+        self.whole_grain = whole_grain
+        self.parents = parents or []
+        self.resolution_cache: Optional[QueryDatasource] = None
+        self.partial_concepts = partial_concepts or get_all_parent_partial(
+            self.output_concepts, self.parents
+        )
+        self.partial_lcl = LooseConceptList(concepts=self.partial_concepts)
+        self.depth = depth
+        self.conditions = conditions
+        self.grain = grain
+        self.force_group = force_group
+        self.tainted = False
+        for parent in self.parents:
+            if not parent:
+                raise SyntaxError("Unresolvable parent")
+
+    def add_output_concept(self, concept: Concept):
+        self.output_concepts.append(concept)
+        self.output_lcl = LooseConceptList(concepts=self.output_concepts)
+        self.rebuild_cache()
+
+    @property
+    def logging_prefix(self) -> str:
+        return "\t" * self.depth
+
+    @property
+    def all_concepts(self) -> list[Concept]:
+        return [*self.output_concepts]
+
+    @property
+    def all_used_concepts(self) -> list[Concept]:
+        return [*self.input_concepts]
+
+    def __repr__(self):
+        concepts = self.all_concepts
+        contents = ",".join(sorted([c.address for c in concepts]))
+        return f"{self.__class__.__name__}<{contents}>"
+
+    def _resolve(self) -> QueryDatasource:
+        parent_sources = [p.resolve() for p in self.parents]
+
+        # if conditional:
+        #     for condition in conditions[1:]:
+        #         conditional += condition
+        grain = Grain(components=self.output_concepts)
+        source_map = resolve_concept_map(
+            parent_sources, self.output_concepts, self.input_concepts
+        )
+        return QueryDatasource(
+            input_concepts=self.input_concepts,
+            output_concepts=self.output_concepts,
+            datasources=parent_sources,
+            source_type=self.source_type,
+            source_map=source_map,
+            joins=[],
+            grain=grain,
+            condition=self.conditions,
+            partial_concepts=self.partial_concepts,
+            force_group=self.force_group,
+        )
+
+    def rebuild_cache(self) -> QueryDatasource:
+        self.tainted = True
+        if not self.resolution_cache:
+            return self.resolve()
+        self.resolution_cache = None
+        return self.resolve()
+
+    def resolve(self) -> QueryDatasource:
+        if self.resolution_cache:
+            return self.resolution_cache
+        qds = self._resolve()
+        self.resolution_cache = qds
+        return qds
+
+
+@dataclass
+class NodeJoin:
+    left_node: StrategyNode
+    right_node: StrategyNode
+    concepts: List[Concept]
+    join_type: JoinType
+    filter_to_mutual: bool = False
+
+    def __post_init__(self):
+        final_concepts = []
+        for concept in self.concepts:
+            include = True
+            for ds in [self.left_node, self.right_node]:
+                if concept.address not in [c.address for c in ds.all_concepts]:
+                    if self.filter_to_mutual:
+                        include = False
+                    else:
+                        raise SyntaxError(
+                            f"Invalid join, missing {concept} on {str(ds)}, have"
+                            f" {[c.address for c in ds.all_concepts]}"
+                        )
+            if include:
+                final_concepts.append(concept)
+        if not final_concepts and self.concepts:
+            # if one datasource only has constants
+            # we can join on 1=1
+            for ds in [self.left_node, self.right_node]:
+                if all(
+                    [c.derivation == PurposeLineage.CONSTANT for c in ds.all_concepts]
+                ):
+                    self.concepts = []
+                    return
+
+            left_keys = [c.address for c in self.left_node.all_concepts]
+            right_keys = [c.address for c in self.right_node.all_concepts]
+            match_concepts = [c.address for c in self.concepts]
+            raise SyntaxError(
+                "No mutual join keys found between"
+                f" {self.left_node} and"
+                f" {self.right_node}, left_keys {left_keys},"
+                f" right_keys {right_keys},"
+                f" provided join concepts {match_concepts}"
+            )
+        self.concepts = final_concepts
+
+    @property
+    def unique_id(self) -> str:
+        nodes = sorted([self.left_node, self.right_node], key=lambda x: str(x))
+        return str(nodes) + self.join_type.value
+
+    def __str__(self):
+        return (
+            f"{self.join_type.value} JOIN {self.left_node} and"
+            f" {self.right_node} on"
+            f" {','.join([str(k) for k in self.concepts])}"
+        )

trilogy/core/processing/nodes/filter_node.py

@@ -0,0 +1,49 @@
+from typing import List
+
+
+from trilogy.core.models import (
+    SourceType,
+    Concept,
+    Conditional,
+    Comparison,
+    Parenthetical,
+)
+from trilogy.core.processing.nodes.base_node import StrategyNode
+
+
+class FilterNode(StrategyNode):
+    """Filter nodes represent a restriction operation
+    on a concept that creates a new derived concept.
+
+    They should only output a concept and it's filtered
+    version, but will have parents that provide all required
+    filtering keys as inputs.
+    """
+
+    source_type = SourceType.FILTER
+
+    def __init__(
+        self,
+        input_concepts: List[Concept],
+        output_concepts: List[Concept],
+        environment,
+        g,
+        whole_grain: bool = False,
+        parents: List["StrategyNode"] | None = None,
+        depth: int = 0,
+        conditions: Conditional | Comparison | Parenthetical | None = None,
+        partial_concepts: List[Concept] | None = None,
+        force_group: bool = False,
+    ):
+        super().__init__(
+            output_concepts=output_concepts,
+            environment=environment,
+            g=g,
+            whole_grain=whole_grain,
+            parents=parents,
+            depth=depth,
+            input_concepts=input_concepts,
+            conditions=conditions,
+            partial_concepts=partial_concepts,
+            force_group=force_group,
+        )

trilogy/core/processing/nodes/group_node.py

@@ -0,0 +1,110 @@
+from typing import List, Optional
+
+from trilogy.constants import logger
+from trilogy.core.models import (
+    Grain,
+    QueryDatasource,
+    SourceType,
+    Concept,
+    Environment,
+    LooseConceptList,
+)
+from trilogy.core.processing.nodes.base_node import (
+    StrategyNode,
+    resolve_concept_map,
+    concept_list_to_grain,
+)
+
+
+LOGGER_PREFIX = "[CONCEPT DETAIL - GROUP NODE]"
+
+
+class GroupNode(StrategyNode):
+    source_type = SourceType.GROUP
+
+    def __init__(
+        self,
+        output_concepts: List[Concept],
+        input_concepts: List[Concept],
+        environment: Environment,
+        g,
+        whole_grain: bool = False,
+        parents: List["StrategyNode"] | None = None,
+        depth: int = 0,
+        partial_concepts: Optional[List[Concept]] = None,
+    ):
+        super().__init__(
+            input_concepts=input_concepts,
+            output_concepts=output_concepts,
+            environment=environment,
+            g=g,
+            whole_grain=whole_grain,
+            parents=parents,
+            depth=depth,
+            partial_concepts=partial_concepts,
+        )
+
+    def _resolve(self) -> QueryDatasource:
+        parent_sources: list[QueryDatasource] = [p.resolve() for p in self.parents]
+
+        grain = concept_list_to_grain(self.output_concepts, [])
+        comp_grain = Grain()
+        for source in parent_sources:
+            comp_grain += source.grain
+
+        # dynamically select if we need to group
+        # because sometimes, we are already at required grain
+        if comp_grain == grain and self.output_lcl == self.input_lcl:
+            # if there is no group by, and inputs equal outputs
+            # return the parent
+            logger.info(
+                f"{self.logging_prefix}{LOGGER_PREFIX} Output of group by node equals input of group by node"
+                f" {self.output_lcl}"
+                f" grains {comp_grain} and {grain}"
+            )
+            if (
+                len(parent_sources) == 1
+                and LooseConceptList(concepts=parent_sources[0].output_concepts)
+                == self.output_lcl
+            ):
+                logger.info(
+                    f"{self.logging_prefix}{LOGGER_PREFIX} No group by required, returning parent node"
+                )
+                return parent_sources[0]
+            # otherwise if no group by, just treat it as a select
+            source_type = SourceType.SELECT
+        else:
+
+            logger.info(
+                f"{self.logging_prefix}{LOGGER_PREFIX} Group node has different output than input, forcing group"
+                f" {self.input_lcl}"
+                " vs"
+                f" {self.output_lcl}"
+                " and"
+                f" upstream grains {[str(source.grain) for source in parent_sources]}"
+                " vs"
+                f" target grain {grain}"
+            )
+            for parent in parent_sources:
+                logger.info(
+                    f"{self.logging_prefix}{LOGGER_PREFIX} Parent node"
+                    f" {[c.address for c in parent.output_concepts]}"
+                    " grain"
+                    f" {parent.grain}"
+                )
+            source_type = SourceType.GROUP
+        return QueryDatasource(
+            input_concepts=self.input_concepts,
+            output_concepts=self.output_concepts,
+            datasources=parent_sources,
+            source_type=source_type,
+            source_map=resolve_concept_map(
+                parent_sources,
+                targets=self.output_concepts,
+                inherited_inputs=self.input_concepts,
+            ),
+            joins=[],
+            grain=grain,
+            partial_concepts=self.partial_concepts,
+            condition=self.conditions,
+        )