pytrilogy 0.0.2.49__py3-none-any.whl → 0.0.2.51__py3-none-any.whl

trilogy/core/processing/node_generators/select_merge_node.py

@@ -13,6 +13,9 @@ from trilogy.core.models import (
     LooseConceptList,
     WhereClause,
 )
+from trilogy.core.processing.node_generators.select_helpers.datasource_injection import (
+    get_union_sources,
+)
 from trilogy.core.processing.nodes import (
     ConstantNode,
     GroupNode,
@@ -20,9 +23,6 @@ from trilogy.core.processing.nodes import (
     SelectNode,
     StrategyNode,
 )
-from trilogy.core.processing.nodes.base_node import (
-    concept_list_to_grain,
-)
 from trilogy.core.processing.utility import padding
 
 LOGGER_PREFIX = "[GEN_ROOT_MERGE_NODE]"
@@ -35,44 +35,75 @@ def extract_address(node: str):
 def get_graph_partial_nodes(
     g: nx.DiGraph, conditions: WhereClause | None
 ) -> dict[str, list[str]]:
-    datasources: dict[str, Datasource] = nx.get_node_attributes(g, "datasource")
+    datasources: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
+        g, "datasource"
+    )
     partial: dict[str, list[str]] = {}
     for node in g.nodes:
         if node in datasources:
             ds = datasources[node]
-            partial[node] = [concept_to_node(c) for c in ds.partial_concepts]
-            if ds.non_partial_for and conditions == ds.non_partial_for:
+            if not isinstance(ds, list):
+                if ds.non_partial_for and conditions == ds.non_partial_for:
+                    partial[node] = []
+                    continue
+                partial[node] = [concept_to_node(c) for c in ds.partial_concepts]
+                ds = [ds]
+            # assume union sources have no partial
+            else:
                 partial[node] = []
 
     return partial
 
 
 def get_graph_grain_length(g: nx.DiGraph) -> dict[str, int]:
-    datasources: dict[str, Datasource] = nx.get_node_attributes(g, "datasource")
-    partial: dict[str, int] = {}
+    datasources: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
+        g, "datasource"
+    )
+    grain_length: dict[str, int] = {}
     for node in g.nodes:
         if node in datasources:
-            partial[node] = len(datasources[node].grain.components)
-    return partial
+            lookup = datasources[node]
+            if not isinstance(lookup, list):
+                lookup = [lookup]
+            assert isinstance(lookup, list)
+            grain_length[node] = sum(len(x.grain.components) for x in lookup)
+    return grain_length
 
 
 def create_pruned_concept_graph(
     g: nx.DiGraph,
     all_concepts: List[Concept],
+    datasources: list[Datasource],
     accept_partial: bool = False,
     conditions: WhereClause | None = None,
+    depth: int = 0,
 ) -> nx.DiGraph:
     orig_g = g
     g = g.copy()
+
+    union_options = get_union_sources(datasources, all_concepts)
+    for ds_list in union_options:
+        node_address = "ds~" + "-".join([x.name for x in ds_list])
+        common: set[Concept] = set.intersection(
+            *[set(x.output_concepts) for x in ds_list]
+        )
+        g.add_node(node_address, datasource=ds_list)
+        for c in common:
+            g.add_edge(node_address, concept_to_node(c))
+
     target_addresses = set([c.address for c in all_concepts])
     concepts: dict[str, Concept] = nx.get_node_attributes(orig_g, "concept")
-    datasources: dict[str, Datasource] = nx.get_node_attributes(orig_g, "datasource")
+    datasource_map: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
+        orig_g, "datasource"
+    )
     relevant_concepts_pre = {
         n: x.address
         for n in g.nodes()
         # filter out synonyms
         if (x := concepts.get(n, None)) and x.address in target_addresses
     }
+    # from trilogy.hooks.graph_hook import GraphHook
+    # GraphHook().query_graph_built(g)
     relevant_concepts: list[str] = list(relevant_concepts_pre.keys())
     relevent_datasets: list[str] = []
     if not accept_partial:
@@ -81,13 +112,13 @@ def create_pruned_concept_graph(
         to_remove = []
         for edge in g.edges:
             if (
-                edge[0] in datasources
+                edge[0] in datasource_map
                 and (pnodes := partial.get(edge[0], []))
                 and edge[1] in pnodes
             ):
                 to_remove.append(edge)
             if (
-                edge[1] in datasources
+                edge[1] in datasource_map
                 and (pnodes := partial.get(edge[1], []))
                 and edge[0] in pnodes
             ):
@@ -128,15 +159,26 @@ def create_pruned_concept_graph(
 
     subgraphs = list(nx.connected_components(g.to_undirected()))
     if not subgraphs:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} cannot resolve root graph - no subgraphs after node prune"
+        )
         return None
     if subgraphs and len(subgraphs) != 1:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} cannot resolve root graph - subgraphs are split - have {len(subgraphs)} from {subgraphs}"
+        )
         return None
     # add back any relevant edges that might have been partially filtered
     relevant = set(relevant_concepts + relevent_datasets)
     for edge in orig_g.edges():
         if edge[0] in relevant and edge[1] in relevant:
             g.add_edge(edge[0], edge[1])
-
+    # if we have no ds nodes at all, for non constant, we can't find it
+    if not any([n.startswith("ds~") for n in g.nodes]):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} cannot resolve root graph - No datasource nodes found"
+        )
+        return None
     return g
 
 
@@ -190,6 +232,54 @@ def resolve_subgraphs(
     return pruned_subgraphs
 
 
+def create_datasource_node(
+    datasource: Datasource,
+    all_concepts: List[Concept],
+    accept_partial: bool,
+    environment: Environment,
+    depth: int,
+    conditions: WhereClause | None = None,
+) -> tuple[StrategyNode, bool]:
+    target_grain = Grain.from_concepts(all_concepts)
+    force_group = False
+    if not datasource.grain.issubset(target_grain):
+        force_group = True
+    partial_concepts = [
+        c.concept
+        for c in datasource.columns
+        if not c.is_complete and c.concept.address in all_concepts
+    ]
+    partial_lcl = LooseConceptList(concepts=partial_concepts)
+    nullable_concepts = [
+        c.concept
+        for c in datasource.columns
+        if c.is_nullable and c.concept.address in all_concepts
+    ]
+    nullable_lcl = LooseConceptList(concepts=nullable_concepts)
+    partial_is_full = conditions and (conditions == datasource.non_partial_for)
+    return (
+        SelectNode(
+            input_concepts=[c.concept for c in datasource.columns],
+            output_concepts=all_concepts,
+            environment=environment,
+            parents=[],
+            depth=depth,
+            partial_concepts=(
+                [] if partial_is_full else [c for c in all_concepts if c in partial_lcl]
+            ),
+            nullable_concepts=[c for c in all_concepts if c in nullable_lcl],
+            accept_partial=accept_partial,
+            datasource=datasource,
+            grain=Grain.from_concepts(all_concepts),
+            conditions=datasource.where.conditional if datasource.where else None,
+            preexisting_conditions=(
+                conditions.conditional if partial_is_full and conditions else None
+            ),
+        ),
+        force_group,
+    )
+
+
 def create_select_node(
     ds_name: str,
     subgraph: list[str],
@@ -199,12 +289,11 @@ def create_select_node(
     depth: int,
     conditions: WhereClause | None = None,
 ) -> StrategyNode:
-    ds_name = ds_name.split("~")[1]
+
     all_concepts = [
         environment.concepts[extract_address(c)] for c in subgraph if c.startswith("c~")
     ]
 
-    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"
@@ -213,7 +302,6 @@ def create_select_node(
             output_concepts=all_concepts,
             input_concepts=[],
             environment=environment,
-            g=g,
             parents=[],
             depth=depth,
             # no partial for constants
@@ -221,41 +309,44 @@ def create_select_node(
             force_group=False,
         )
 
-    datasource = environment.datasources[ds_name]
-    target_grain = Grain(components=all_concepts)
-    force_group = False
-    if not datasource.grain.issubset(target_grain):
-        force_group = True
-    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)
-    nullable_concepts = [
-        c.concept for c in datasource.columns if c.is_nullable and c.concept in all_lcl
-    ]
-    nullable_lcl = LooseConceptList(concepts=nullable_concepts)
-    partial_is_full = conditions and (conditions == datasource.non_partial_for)
-    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=(
-            [] if partial_is_full else [c for c in all_concepts if c in partial_lcl]
-        ),
-        nullable_concepts=[c for c in all_concepts if c in nullable_lcl],
-        accept_partial=accept_partial,
-        datasource=datasource,
-        grain=Grain(components=all_concepts),
-        conditions=datasource.where.conditional if datasource.where else None,
-        preexisting_conditions=(
-            conditions.conditional if partial_is_full and conditions else None
-        ),
-    )
+    datasource: dict[str, Datasource | list[Datasource]] = nx.get_node_attributes(
+        g, "datasource"
+    )[ds_name]
+    if isinstance(datasource, Datasource):
+        bcandidate, force_group = create_datasource_node(
+            datasource,
+            all_concepts,
+            accept_partial,
+            environment,
+            depth,
+            conditions=conditions,
+        )
+
+    elif isinstance(datasource, list):
+        from trilogy.core.processing.nodes.union_node import UnionNode
+
+        force_group = False
+        parents = []
+        for x in datasource:
+            subnode, fg = create_datasource_node(
+                x,
+                all_concepts,
+                accept_partial,
+                environment,
+                depth,
+                conditions=conditions,
+            )
+            parents.append(subnode)
+            force_group = force_group or fg
+        bcandidate = UnionNode(
+            output_concepts=all_concepts,
+            input_concepts=all_concepts,
+            environment=environment,
+            parents=parents,
+            depth=depth,
+        )
+    else:
+        raise ValueError(f"Unknown datasource type {datasource}")
 
     # we need to nest the group node one further
     if force_group is True:
@@ -263,14 +354,11 @@ def create_select_node(
             output_concepts=all_concepts,
             input_concepts=all_concepts,
             environment=environment,
-            g=g,
             parents=[bcandidate],
             depth=depth,
             partial_concepts=bcandidate.partial_concepts,
             nullable_concepts=bcandidate.nullable_concepts,
-            preexisting_conditions=(
-                conditions.conditional if partial_is_full and conditions else None
-            ),
+            preexisting_conditions=bcandidate.preexisting_conditions,
         )
     else:
         candidate = bcandidate
@@ -292,7 +380,6 @@ def gen_select_merge_node(
             output_concepts=constants,
             input_concepts=[],
             environment=environment,
-            g=g,
             parents=[],
             depth=depth,
             partial_concepts=[],
@@ -300,7 +387,12 @@ def gen_select_merge_node(
         )
     for attempt in [False, True]:
         pruned_concept_graph = create_pruned_concept_graph(
-            g, non_constant, attempt, conditions
+            g,
+            non_constant,
+            accept_partial=attempt,
+            conditions=conditions,
+            datasources=list(environment.datasources.values()),
+            depth=depth,
         )
         if pruned_concept_graph:
             logger.info(
@@ -309,9 +401,7 @@ def gen_select_merge_node(
             break
 
     if not pruned_concept_graph:
-        logger.info(
-            f"{padding(depth)}{LOGGER_PREFIX} no covering graph found {attempt}"
-        )
+        logger.info(f"{padding(depth)}{LOGGER_PREFIX} no covering graph found.")
         return None
 
     sub_nodes = resolve_subgraphs(pruned_concept_graph, conditions)
@@ -321,7 +411,7 @@ def gen_select_merge_node(
         create_select_node(
             k,
             subgraph,
-            g=g,
+            g=pruned_concept_graph,
             accept_partial=accept_partial,
             environment=environment,
             depth=depth,
@@ -338,7 +428,6 @@ def gen_select_merge_node(
                 output_concepts=constants,
                 input_concepts=[],
                 environment=environment,
-                g=g,
                 parents=[],
                 depth=depth,
                 partial_concepts=[],
@@ -361,18 +450,16 @@ def gen_select_merge_node(
         output_concepts=all_concepts,
         input_concepts=non_constant,
         environment=environment,
-        g=g,
         depth=depth,
         parents=parents,
         preexisting_conditions=preexisting_conditions,
     )
-    target_grain = concept_list_to_grain(all_concepts, [])
+    target_grain = Grain.from_concepts(all_concepts)
     if not base.resolve().grain.issubset(target_grain):
         return GroupNode(
             output_concepts=all_concepts,
             input_concepts=all_concepts,
             environment=environment,
-            g=g,
             parents=[base],
             depth=depth,
             preexisting_conditions=preexisting_conditions,

trilogy/core/processing/node_generators/union_node.py

@@ -70,6 +70,5 @@ def gen_union_node(
         input_concepts=[concept] + local_optional,
         output_concepts=[concept] + local_optional,
         environment=environment,
-        g=g,
         parents=parents,
     )

trilogy/core/processing/node_generators/unnest_node.py

@@ -46,7 +46,6 @@ def gen_unnest_node(
         input_concepts=arguments + non_equivalent_optional,
         output_concepts=[concept] + local_optional,
         environment=environment,
-        g=g,
         parents=([parent] if (arguments or local_optional) else []),
     )
     # we need to sometimes nest an unnest node,
@@ -56,7 +55,6 @@ def gen_unnest_node(
         input_concepts=base.output_concepts,
         output_concepts=base.output_concepts,
         environment=environment,
-        g=g,
         parents=[base],
         preexisting_conditions=conditions.conditional if conditions else None,
     )

trilogy/core/processing/node_generators/window_node.py

@@ -86,7 +86,6 @@ def gen_window_node(
         input_concepts=parent_concepts + targets + non_equivalent_optional,
         output_concepts=[concept] + parent_concepts + local_optional,
         environment=environment,
-        g=g,
         parents=[
             parent_node,
         ],
@@ -98,7 +97,6 @@ def gen_window_node(
         input_concepts=[concept] + local_optional,
         output_concepts=[concept] + local_optional,
         environment=environment,
-        g=g,
         parents=[_window_node],
         preexisting_conditions=conditions.conditional if conditions else None,
     )

trilogy/core/processing/nodes/base_node.py

@@ -1,12 +1,10 @@
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import List, Optional, Sequence
+from typing import List, Optional
 
 from trilogy.core.enums import (
     BooleanOperator,
-    Granularity,
     JoinType,
-    Purpose,
     PurposeLineage,
 )
 from trilogy.core.models import (
@@ -26,31 +24,6 @@ from trilogy.core.models import (
 from trilogy.utility import unique
 
 
-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 | Datasource],
     targets: List[Concept],
@@ -156,7 +129,6 @@ class StrategyNode:
         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,
@@ -178,7 +150,6 @@ class StrategyNode:
         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
@@ -353,11 +324,7 @@ class StrategyNode:
             p.resolve() for p in self.parents
         ]
 
-        grain = (
-            self.grain
-            if self.grain
-            else concept_list_to_grain(self.output_concepts, [])
-        )
+        grain = self.grain if self.grain else Grain.from_concepts(self.output_concepts)
         source_map = resolve_concept_map(
             parent_sources,
             targets=self.output_concepts,
@@ -399,7 +366,6 @@ class StrategyNode:
             input_concepts=list(self.input_concepts),
             output_concepts=list(self.output_concepts),
             environment=self.environment,
-            g=self.g,
             whole_grain=self.whole_grain,
             parents=list(self.parents),
             partial_concepts=list(self.partial_concepts),

trilogy/core/processing/nodes/filter_node.py

@@ -27,7 +27,6 @@ class FilterNode(StrategyNode):
         input_concepts: List[Concept],
         output_concepts: List[Concept],
         environment,
-        g,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
         depth: int = 0,
@@ -41,7 +40,6 @@ class FilterNode(StrategyNode):
         super().__init__(
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -59,7 +57,6 @@ class FilterNode(StrategyNode):
             input_concepts=list(self.input_concepts),
             output_concepts=list(self.output_concepts),
             environment=self.environment,
-            g=self.g,
             whole_grain=self.whole_grain,
             parents=self.parents,
             depth=self.depth,

trilogy/core/processing/nodes/group_node.py

@@ -15,10 +15,10 @@ from trilogy.core.models import (
 )
 from trilogy.core.processing.nodes.base_node import (
     StrategyNode,
-    concept_list_to_grain,
     resolve_concept_map,
 )
 from trilogy.core.processing.utility import find_nullable_concepts, is_scalar_condition
+from trilogy.parsing.common import concepts_to_grain_concepts
 from trilogy.utility import unique
 
 LOGGER_PREFIX = "[CONCEPT DETAIL - GROUP NODE]"
@@ -32,7 +32,6 @@ class GroupNode(StrategyNode):
         output_concepts: List[Concept],
         input_concepts: List[Concept],
         environment: Environment,
-        g,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
         depth: int = 0,
@@ -48,7 +47,6 @@ class GroupNode(StrategyNode):
             input_concepts=input_concepts,
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -66,19 +64,27 @@ class GroupNode(StrategyNode):
             p.resolve() for p in self.parents
         ]
 
-        grain = self.grain or concept_list_to_grain(self.output_concepts, [])
+        target_grain = self.grain or Grain.from_concepts(
+            concepts_to_grain_concepts(
+                self.output_concepts, environment=self.environment
+            )
+        )
         comp_grain = Grain()
         for source in parent_sources:
             comp_grain += source.grain
-
+        comp_grain = Grain.from_concepts(
+            concepts_to_grain_concepts(
+                comp_grain.components, environment=self.environment
+            )
+        )
         # dynamically select if we need to group
         # because sometimes, we are already at required grain
-        if comp_grain == grain and self.force_group is not True:
+        if comp_grain == target_grain and self.force_group is not True:
             # if there is no group by, and inputs equal outputs
             # return the parent
             logger.info(
                 f"{self.logging_prefix}{LOGGER_PREFIX} Grain of group by equals output"
-                f" grains {comp_grain} and {grain}"
+                f" grains {comp_grain} and {target_grain}"
             )
             if (
                 len(parent_sources) == 1
@@ -96,10 +102,11 @@ class GroupNode(StrategyNode):
             source_type = SourceType.SELECT
         else:
             logger.info(
-                f"{self.logging_prefix}{LOGGER_PREFIX} Group node has different grain than parents; forcing group"
-                f" upstream grains {[str(source.grain) for source in parent_sources]}"
+                f"{self.logging_prefix}{LOGGER_PREFIX} Group node has different grain than parents; group is required."
+                f" Upstream grains {[str(source.grain) for source in parent_sources]}"
                 f" with final grain {comp_grain} vs"
-                f" target grain {grain}"
+                f" target grain {target_grain}"
+                f" delta: {comp_grain - target_grain}"
             )
             for parent in self.parents:
                 logger.info(
@@ -136,7 +143,7 @@ class GroupNode(StrategyNode):
             source_type=source_type,
             source_map=source_map,
             joins=[],
-            grain=grain,
+            grain=target_grain,
             partial_concepts=self.partial_concepts,
             nullable_concepts=nullable_concepts,
             hidden_concepts=self.hidden_concepts,
@@ -165,7 +172,7 @@ class GroupNode(StrategyNode):
                 source_type=SourceType.SELECT,
                 source_map=source_map,
                 joins=[],
-                grain=grain,
+                grain=target_grain,
                 nullable_concepts=base.nullable_concepts,
                 partial_concepts=self.partial_concepts,
                 condition=self.conditions,
@@ -178,7 +185,6 @@ class GroupNode(StrategyNode):
             input_concepts=list(self.input_concepts),
             output_concepts=list(self.output_concepts),
             environment=self.environment,
-            g=self.g,
             whole_grain=self.whole_grain,
             parents=self.parents,
             depth=self.depth,

trilogy/core/processing/nodes/merge_node.py

@@ -58,7 +58,7 @@ def deduplicate_nodes(
                 og = merged[k1]
                 subset_to = merged[k2]
                 logger.info(
-                    f"{logging_prefix}{LOGGER_PREFIX} extraneous parent node that is subset of another parent node {og.grain.issubset(subset_to.grain)} {og.grain.set} {subset_to.grain.set}"
+                    f"{logging_prefix}{LOGGER_PREFIX} extraneous parent node that is subset of another parent node {og.grain.issubset(subset_to.grain)} {og.grain.components} {subset_to.grain.components}"
                 )
                 merged = {k: v for k, v in merged.items() if k != k1}
                 removed.add(k1)
@@ -103,7 +103,6 @@ class MergeNode(StrategyNode):
         input_concepts: List[Concept],
         output_concepts: List[Concept],
         environment,
-        g,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
         node_joins: List[NodeJoin] | None = None,
@@ -124,7 +123,6 @@ class MergeNode(StrategyNode):
             input_concepts=input_concepts,
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -199,7 +197,7 @@ class MergeNode(StrategyNode):
     ) -> List[BaseJoin | UnnestJoin]:
         # only finally, join between them for unique values
         dataset_list: List[QueryDatasource | Datasource] = sorted(
-            final_datasets, key=lambda x: -len(x.grain.components_copy)
+            final_datasets, key=lambda x: -len(x.grain.components)
         )
 
         logger.info(
@@ -364,7 +362,6 @@ class MergeNode(StrategyNode):
             input_concepts=list(self.input_concepts),
             output_concepts=list(self.output_concepts),
             environment=self.environment,
-            g=self.g,
             whole_grain=self.whole_grain,
             parents=self.parents,
             depth=self.depth,