pytrilogy 0.0.2.48__py3-none-any.whl → 0.0.2.50__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,
@@ -35,38 +38,66 @@ 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,
 ) -> 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()
@@ -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
             ):
@@ -136,7 +167,9 @@ def create_pruned_concept_graph(
     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]):
+        return None
     return g
 
 
@@ -190,6 +223,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(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.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(components=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 +280,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 +293,6 @@ def create_select_node(
             output_concepts=all_concepts,
             input_concepts=[],
             environment=environment,
-            g=g,
             parents=[],
             depth=depth,
             # no partial for constants
@@ -221,41 +300,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 +345,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 +371,6 @@ def gen_select_merge_node(
             output_concepts=constants,
             input_concepts=[],
             environment=environment,
-            g=g,
             parents=[],
             depth=depth,
             partial_concepts=[],
@@ -300,7 +378,11 @@ 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()),
         )
         if pruned_concept_graph:
             logger.info(
@@ -321,7 +403,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 +420,6 @@ def gen_select_merge_node(
                 output_concepts=constants,
                 input_concepts=[],
                 environment=environment,
-                g=g,
                 parents=[],
                 depth=depth,
                 partial_concepts=[],
@@ -361,7 +442,6 @@ 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,
@@ -372,7 +452,6 @@ def gen_select_merge_node(
             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

@@ -156,7 +156,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 +177,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
@@ -211,8 +209,22 @@ class StrategyNode:
                 operator=BooleanOperator.AND,
             )
         self.validate_parents()
+        self.validate_inputs()
         self.log = True
 
+    def validate_inputs(self):
+        if not self.parents:
+            return
+        non_hidden = set()
+        for x in self.parents:
+            for z in x.usable_outputs:
+                non_hidden.add(z.address)
+        if not all([x.address in non_hidden for x in self.input_concepts]):
+            missing = [x for x in self.input_concepts if x.address not in non_hidden]
+            raise ValueError(
+                f"Invalid input concepts; {missing} are missing non-hidden parent nodes"
+            )
+
     def add_parents(self, parents: list["StrategyNode"]):
         self.parents += parents
         self.validate_parents()
@@ -288,6 +300,14 @@ class StrategyNode:
             self.rebuild_cache()
         return self
 
+    def unhide_output_concepts(self, concepts: List[Concept], rebuild: bool = True):
+        self.hidden_concepts = [
+            x for x in self.hidden_concepts if x.address not in concepts
+        ]
+        if rebuild:
+            self.rebuild_cache()
+        return self
+
     def remove_output_concepts(self, concepts: List[Concept], rebuild: bool = True):
         for x in concepts:
             self.hidden_concepts.append(x)
@@ -299,6 +319,12 @@ class StrategyNode:
             self.rebuild_cache()
         return self
 
+    @property
+    def usable_outputs(self) -> list[Concept]:
+        return [
+            x for x in self.output_concepts if x.address not in self.hidden_concepts
+        ]
+
     @property
     def logging_prefix(self) -> str:
         return "\t" * self.depth
@@ -371,7 +397,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

@@ -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,
@@ -105,9 +103,9 @@ class GroupNode(StrategyNode):
                 logger.info(
                     f"{self.logging_prefix}{LOGGER_PREFIX} Parent node"
                     f" {[c.address for c in parent.output_concepts[:2]]}... has"
-                    " grain"
+                    " set node grain"
                     f" {parent.grain}"
-                    f" resolved grain {parent.resolve().grain}"
+                    f" and resolved grain {parent.resolve().grain}"
                     f" {type(parent)}"
                 )
             source_type = SourceType.GROUP
@@ -146,7 +144,13 @@ class GroupNode(StrategyNode):
         # inject an additional CTE
         if self.conditions and not is_scalar_condition(self.conditions):
             base.condition = None
-            base.output_concepts = self.output_concepts + self.conditions.row_arguments
+            base.output_concepts = unique(
+                base.output_concepts + self.conditions.row_arguments, "address"
+            )
+            # re-visible any hidden concepts
+            base.hidden_concepts = [
+                x for x in base.hidden_concepts if x not in base.output_concepts
+            ]
             source_map = resolve_concept_map(
                 [base],
                 targets=self.output_concepts,
@@ -172,7 +176,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

@@ -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,
@@ -330,8 +328,9 @@ class MergeNode(StrategyNode):
             force_group = None
 
         qd_joins: List[BaseJoin | UnnestJoin] = [*joins]
+
         source_map = resolve_concept_map(
-            list(merged.values()),
+            final_datasets,
             targets=self.output_concepts,
             inherited_inputs=self.input_concepts + self.existence_concepts,
             full_joins=full_join_concepts,
@@ -339,6 +338,7 @@ class MergeNode(StrategyNode):
         nullable_concepts = find_nullable_concepts(
             source_map=source_map, joins=joins, datasources=final_datasets
         )
+
         qds = QueryDatasource(
             input_concepts=unique(self.input_concepts, "address"),
             output_concepts=unique(self.output_concepts, "address"),
@@ -362,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,

trilogy/core/processing/nodes/select_node_v2.py

@@ -34,7 +34,6 @@ class SelectNode(StrategyNode):
         input_concepts: List[Concept],
         output_concepts: List[Concept],
         environment: Environment,
-        g,
         datasource: Datasource | None = None,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
@@ -52,7 +51,6 @@ class SelectNode(StrategyNode):
             input_concepts=input_concepts,
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -67,6 +65,11 @@ class SelectNode(StrategyNode):
         self.accept_partial = accept_partial
         self.datasource = datasource
 
+    def validate_inputs(self):
+        # we do not need to validate inputs for a select node
+        # as it will be a root
+        return
+
     def resolve_from_provided_datasource(
         self,
     ) -> QueryDatasource:
@@ -192,7 +195,6 @@ class SelectNode(StrategyNode):
             input_concepts=list(self.input_concepts),
             output_concepts=list(self.output_concepts),
             environment=self.environment,
-            g=self.g,
             datasource=self.datasource,
             depth=self.depth,
             parents=self.parents,
@@ -216,7 +218,6 @@ class ConstantNode(SelectNode):
             input_concepts=list(self.input_concepts),
             output_concepts=list(self.output_concepts),
             environment=self.environment,
-            g=self.g,
             datasource=self.datasource,
             depth=self.depth,
             partial_concepts=list(self.partial_concepts),

trilogy/core/processing/nodes/union_node.py

@@ -18,7 +18,6 @@ class UnionNode(StrategyNode):
         input_concepts: List[Concept],
         output_concepts: List[Concept],
         environment,
-        g,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
         depth: int = 0,
@@ -27,7 +26,6 @@ class UnionNode(StrategyNode):
             input_concepts=input_concepts,
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -43,7 +41,6 @@ class UnionNode(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/unnest_node.py

@@ -23,7 +23,6 @@ class UnnestNode(StrategyNode):
         input_concepts: List[Concept],
         output_concepts: List[Concept],
         environment,
-        g,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
         depth: int = 0,
@@ -32,7 +31,6 @@ class UnnestNode(StrategyNode):
             input_concepts=input_concepts,
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -62,7 +60,6 @@ class UnnestNode(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/window_node.py

@@ -12,7 +12,6 @@ class WindowNode(StrategyNode):
         input_concepts: List[Concept],
         output_concepts: List[Concept],
         environment,
-        g,
         whole_grain: bool = False,
         parents: List["StrategyNode"] | None = None,
         depth: int = 0,
@@ -21,7 +20,6 @@ class WindowNode(StrategyNode):
             input_concepts=input_concepts,
             output_concepts=output_concepts,
             environment=environment,
-            g=g,
             whole_grain=whole_grain,
             parents=parents,
             depth=depth,
@@ -36,7 +34,6 @@ class WindowNode(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,