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

trilogy/core/processing/node_generators/node_merge_node.py

@@ -1,16 +1,16 @@
 from typing import List, Optional
 
-from trilogy.core.models import Concept, Environment, Conditional, WhereClause
-from trilogy.core.processing.nodes import MergeNode, History, StrategyNode
 import networkx as nx
-from trilogy.core.graph_models import concept_to_node
+from networkx.algorithms import approximation as ax
+
 from trilogy.constants import logger
-from trilogy.utility import unique
+from trilogy.core.enums import PurposeLineage
 from trilogy.core.exceptions import AmbiguousRelationshipResolutionException
+from trilogy.core.graph_models import concept_to_node
+from trilogy.core.models import Concept, Conditional, Environment, WhereClause
+from trilogy.core.processing.nodes import History, MergeNode, StrategyNode
 from trilogy.core.processing.utility import padding
-from networkx.algorithms import approximation as ax
-from trilogy.core.enums import PurposeLineage
-
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[GEN_MERGE_NODE]"
 AMBIGUITY_CHECK_LIMIT = 20
@@ -194,7 +194,6 @@ def resolve_weak_components(
     filter_downstream: bool = True,
     accept_partial: bool = False,
 ) -> list[list[Concept]] | None:
-
     break_flag = False
     found = []
     search_graph = environment_graph.copy()
@@ -355,7 +354,6 @@ def gen_merge_node(
     conditions: Conditional | None = None,
     search_conditions: WhereClause | None = None,
 ) -> Optional[MergeNode]:
-
     if search_conditions:
         all_concepts = unique(all_concepts + search_conditions.row_arguments, "address")
     for filter_downstream in [True, False]:

trilogy/core/processing/node_generators/rowset_node.py

@@ -1,19 +1,19 @@
+from typing import List
+
+from trilogy.constants import logger
+from trilogy.core.enums import PurposeLineage
 from trilogy.core.models import (
     Concept,
     Environment,
-    SelectStatement,
+    MultiSelectStatement,
     RowsetDerivationStatement,
     RowsetItem,
-    MultiSelectStatement,
+    SelectStatement,
     WhereClause,
 )
-from trilogy.core.processing.nodes import MergeNode, History, StrategyNode
-from typing import List
-
-from trilogy.core.enums import PurposeLineage
-from trilogy.constants import logger
-from trilogy.core.processing.utility import padding, concept_to_relevant_joins
+from trilogy.core.processing.nodes import History, MergeNode, StrategyNode
 from trilogy.core.processing.nodes.base_node import concept_list_to_grain
+from trilogy.core.processing.utility import concept_to_relevant_joins, padding
 
 LOGGER_PREFIX = "[GEN_ROWSET_NODE]"
 
@@ -37,7 +37,8 @@ def gen_rowset_node(
     lineage: RowsetItem = concept.lineage
     rowset: RowsetDerivationStatement = lineage.rowset
     select: SelectStatement | MultiSelectStatement = lineage.rowset.select
-    node = get_query_node(environment, select, graph=g, history=history)
+
+    node = get_query_node(environment, select)
 
     if not node:
         logger.info(
@@ -93,15 +94,22 @@ def gen_rowset_node(
         logger.info(
             f"{padding(depth)}{LOGGER_PREFIX} no enrichment required for rowset node as all optional found or no optional; exiting early."
         )
-        # node.set_preexisting_conditions(conditions.conditional if conditions else None)
         return node
-
-    possible_joins = concept_to_relevant_joins(node.output_concepts)
+    possible_joins = concept_to_relevant_joins(
+        [x for x in node.output_concepts if x.derivation != PurposeLineage.ROWSET]
+    )
+    logger.info({x.address: x.keys for x in possible_joins})
     if not possible_joins:
         logger.info(
             f"{padding(depth)}{LOGGER_PREFIX} no possible joins for rowset node to get {[x.address for x in local_optional]}; have {[x.address for x in node.output_concepts]}"
         )
         return node
+    if any(x.derivation == PurposeLineage.ROWSET for x in possible_joins):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} cannot enrich rowset node with rowset concepts; exiting early"
+        )
+        return node
+    logger.info([x.address for x in possible_joins + local_optional])
     enrich_node: MergeNode = source_concepts(  # this fetches the parent + join keys
         # to then connect to the rest of the query
         mandatory_list=possible_joins + local_optional,
@@ -109,15 +117,28 @@ def gen_rowset_node(
         g=g,
         depth=depth + 1,
         conditions=conditions,
+        history=history,
     )
     if not enrich_node:
         logger.info(
             f"{padding(depth)}{LOGGER_PREFIX} Cannot generate rowset enrichment node for {concept} with optional {local_optional}, returning just rowset node"
         )
         return node
+
+    non_hidden = [
+        x for x in node.output_concepts if x.address not in node.hidden_concepts
+    ]
+    for x in possible_joins:
+        if x.address in node.hidden_concepts:
+            node.unhide_output_concepts([x])
+    non_hidden_enrich = [
+        x
+        for x in enrich_node.output_concepts
+        if x.address not in enrich_node.hidden_concepts
+    ]
     return MergeNode(
-        input_concepts=enrich_node.output_concepts + node.output_concepts,
-        output_concepts=node.output_concepts + local_optional,
+        input_concepts=non_hidden + non_hidden_enrich,
+        output_concepts=non_hidden + local_optional,
         environment=environment,
         g=g,
         depth=depth,
@@ -125,6 +146,6 @@ def gen_rowset_node(
             node,
             enrich_node,
         ],
-        partial_concepts=node.partial_concepts,
+        partial_concepts=node.partial_concepts + enrich_node.partial_concepts,
         preexisting_conditions=conditions.conditional if conditions else None,
     )

trilogy/core/processing/node_generators/select_merge_node.py

@@ -1,28 +1,29 @@
 from typing import List, Optional
 
+import networkx as nx
+
+from trilogy.constants import logger
+from trilogy.core.enums import PurposeLineage
+from trilogy.core.graph_models import concept_to_node
 from trilogy.core.models import (
     Concept,
+    Datasource,
     Environment,
     Grain,
-    Datasource,
-    WhereClause,
     LooseConceptList,
+    WhereClause,
 )
 from trilogy.core.processing.nodes import (
-    MergeNode,
-    StrategyNode,
-    GroupNode,
     ConstantNode,
+    GroupNode,
+    MergeNode,
     SelectNode,
+    StrategyNode,
 )
-import networkx as nx
-from trilogy.core.graph_models import concept_to_node
-from trilogy.constants import logger
-from trilogy.core.processing.utility import padding
-from trilogy.core.enums import PurposeLineage
 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]"
 

trilogy/core/processing/node_generators/select_node.py

@@ -1,19 +1,19 @@
+from trilogy.constants import logger
 from trilogy.core.enums import PurposeLineage
+from trilogy.core.exceptions import NoDatasourceException
 from trilogy.core.models import (
     Concept,
     Environment,
     LooseConceptList,
     WhereClause,
 )
+from trilogy.core.processing.node_generators.select_merge_node import (
+    gen_select_merge_node,
+)
 from trilogy.core.processing.nodes import (
     StrategyNode,
 )
-from trilogy.core.exceptions import NoDatasourceException
-from trilogy.constants import logger
 from trilogy.core.processing.utility import padding
-from trilogy.core.processing.node_generators.select_merge_node import (
-    gen_select_merge_node,
-)
 
 LOGGER_PREFIX = "[GEN_SELECT_NODE]"
 

trilogy/core/processing/node_generators/union_node.py

@@ -0,0 +1,75 @@
+from typing import List
+
+from trilogy.constants import logger
+from trilogy.core.enums import FunctionType, Purpose
+from trilogy.core.models import Concept, Function, WhereClause
+from trilogy.core.processing.nodes import History, StrategyNode, UnionNode
+from trilogy.core.processing.utility import padding
+
+LOGGER_PREFIX = "[GEN_UNION_NODE]"
+
+
+def is_union(c: Concept):
+    return isinstance(c.lineage, Function) and c.lineage.operator == FunctionType.UNION
+
+
+def gen_union_node(
+    concept: Concept,
+    local_optional: List[Concept],
+    environment,
+    g,
+    depth: int,
+    source_concepts,
+    history: History | None = None,
+    conditions: WhereClause | None = None,
+) -> StrategyNode | None:
+    all_unions = [x for x in local_optional if is_union(x)] + [concept]
+
+    parents = []
+    keys = [x for x in all_unions if x.purpose == Purpose.KEY]
+    base = keys.pop()
+    remaining = [x for x in all_unions if x.address != base.address]
+    arguments = []
+    if isinstance(base.lineage, Function):
+        arguments = base.lineage.concept_arguments
+    for arg in arguments:
+        relevant_parents: list[Concept] = []
+        for other_union in remaining:
+            assert other_union.lineage
+            potential_parents = [
+                z for z in other_union.lineage.arguments if isinstance(z, Concept)
+            ]
+            relevant_parents += [
+                x for x in potential_parents if x.keys and arg.address in x.keys
+            ]
+        logger.info(
+            f"For parent arg {arg.address}, including additional union inputs {[c.address for c in relevant_parents]}"
+        )
+        parent: StrategyNode = source_concepts(
+            mandatory_list=[arg] + relevant_parents,
+            environment=environment,
+            g=g,
+            depth=depth + 1,
+            history=history,
+            conditions=conditions,
+        )
+        parent.hide_output_concepts(parent.output_concepts)
+        # parent.remove_output_concepts(parent.output_concepts)
+        parent.add_output_concept(concept)
+        for x in remaining:
+            parent.add_output_concept(x)
+
+        parents.append(parent)
+        if not parent:
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} could not find union node parents"
+            )
+            return None
+
+    return UnionNode(
+        input_concepts=[concept] + local_optional,
+        output_concepts=[concept] + local_optional,
+        environment=environment,
+        g=g,
+        parents=parents,
+    )

trilogy/core/processing/node_generators/unnest_node.py

@@ -1,10 +1,9 @@
 from typing import List
 
-
+from trilogy.constants import logger
 from trilogy.core.models import Concept, Function, WhereClause
-from trilogy.core.processing.nodes import UnnestNode, History, StrategyNode
+from trilogy.core.processing.nodes import History, StrategyNode, UnnestNode
 from trilogy.core.processing.utility import padding
-from trilogy.constants import logger
 
 LOGGER_PREFIX = "[GEN_UNNEST_NODE]"
 

trilogy/core/processing/node_generators/window_node.py

@@ -1,11 +1,10 @@
 from typing import List
 
-
-from trilogy.core.models import Concept, WindowItem, Environment, WhereClause
-from trilogy.utility import unique
-from trilogy.core.processing.nodes import WindowNode, StrategyNode, History
 from trilogy.constants import logger
+from trilogy.core.models import Concept, Environment, WhereClause, WindowItem
+from trilogy.core.processing.nodes import History, StrategyNode, WindowNode
 from trilogy.core.processing.utility import padding
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[GEN_WINDOW_NODE]"
 

trilogy/core/processing/nodes/__init__.py

@@ -1,12 +1,15 @@
+from pydantic import BaseModel, ConfigDict, Field
+
+from trilogy.core.models import Concept, Environment, WhereClause
+
+from .base_node import NodeJoin, StrategyNode
 from .filter_node import FilterNode
 from .group_node import GroupNode
 from .merge_node import MergeNode
-from .select_node_v2 import SelectNode, ConstantNode
-from .window_node import WindowNode
-from .base_node import StrategyNode, NodeJoin
+from .select_node_v2 import ConstantNode, SelectNode
+from .union_node import UnionNode
 from .unnest_node import UnnestNode
-from pydantic import BaseModel, Field, ConfigDict
-from trilogy.core.models import Concept, Environment, WhereClause
+from .window_node import WindowNode
 
 
 class History(BaseModel):
@@ -160,5 +163,6 @@ __all__ = [
     "NodeJoin",
     "ConstantNode",
     "UnnestNode",
+    "UnionNode",
     "History",
 ]

trilogy/core/processing/nodes/base_node.py

@@ -1,24 +1,29 @@
-from typing import List, Optional, Sequence
 from collections import defaultdict
+from dataclasses import dataclass
+from typing import List, Optional, Sequence
 
+from trilogy.core.enums import (
+    BooleanOperator,
+    Granularity,
+    JoinType,
+    Purpose,
+    PurposeLineage,
+)
 from trilogy.core.models import (
-    Grain,
-    QueryDatasource,
-    SourceType,
+    Comparison,
     Concept,
-    Environment,
+    ConceptPair,
     Conditional,
-    UnnestJoin,
     Datasource,
-    Comparison,
-    Parenthetical,
+    Environment,
+    Grain,
     LooseConceptList,
-    ConceptPair,
+    Parenthetical,
+    QueryDatasource,
+    SourceType,
+    UnnestJoin,
 )
-from trilogy.core.enums import Purpose, JoinType, PurposeLineage, Granularity
 from trilogy.utility import unique
-from dataclasses import dataclass
-from trilogy.core.enums import BooleanOperator
 
 
 def concept_list_to_grain(
@@ -67,7 +72,6 @@ def resolve_concept_map(
             ]:
                 continue
             if concept.address in full_addresses:
-
                 concept_map[concept.address].add(input)
             elif concept.address not in concept_map:
                 # equi_targets = [x for x in targets if concept.address in x.pseudonyms or x.address in concept.pseudonyms]
@@ -207,8 +211,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()
@@ -284,6 +302,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)
@@ -295,6 +321,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

trilogy/core/processing/nodes/filter_node.py

@@ -1,13 +1,12 @@
 from typing import List
 
-
 from trilogy.core.models import (
-    SourceType,
+    Comparison,
     Concept,
     Conditional,
-    Comparison,
-    Parenthetical,
     Grain,
+    Parenthetical,
+    SourceType,
 )
 from trilogy.core.processing.nodes.base_node import StrategyNode
 

trilogy/core/processing/nodes/group_node.py

@@ -2,25 +2,24 @@ from typing import List, Optional
 
 from trilogy.constants import logger
 from trilogy.core.models import (
-    Grain,
-    QueryDatasource,
-    Datasource,
-    SourceType,
+    Comparison,
     Concept,
+    Conditional,
+    Datasource,
     Environment,
+    Grain,
     LooseConceptList,
-    Conditional,
-    Comparison,
     Parenthetical,
+    QueryDatasource,
+    SourceType,
 )
 from trilogy.core.processing.nodes.base_node import (
     StrategyNode,
-    resolve_concept_map,
     concept_list_to_grain,
+    resolve_concept_map,
 )
+from trilogy.core.processing.utility import find_nullable_concepts, is_scalar_condition
 from trilogy.utility import unique
-from trilogy.core.processing.utility import is_scalar_condition
-from trilogy.core.processing.utility import find_nullable_concepts
 
 LOGGER_PREFIX = "[CONCEPT DETAIL - GROUP NODE]"
 
@@ -96,7 +95,6 @@ class GroupNode(StrategyNode):
             # 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 grain than parents; forcing group"
                 f" upstream grains {[str(source.grain) for source in parent_sources]}"
@@ -107,9 +105,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
@@ -148,7 +146,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,

trilogy/core/processing/nodes/merge_node.py

@@ -1,28 +1,27 @@
 from typing import List, Optional, Tuple
 
-
 from trilogy.constants import logger
 from trilogy.core.models import (
     BaseJoin,
+    Comparison,
+    Concept,
+    Conditional,
+    Datasource,
+    Environment,
     Grain,
     JoinType,
+    Parenthetical,
     QueryDatasource,
-    Datasource,
     SourceType,
-    Concept,
     UnnestJoin,
-    Conditional,
-    Comparison,
-    Parenthetical,
-    Environment,
 )
-from trilogy.utility import unique
 from trilogy.core.processing.nodes.base_node import (
+    NodeJoin,
     StrategyNode,
     resolve_concept_map,
-    NodeJoin,
 )
-from trilogy.core.processing.utility import get_node_joins, find_nullable_concepts
+from trilogy.core.processing.utility import find_nullable_concepts, get_node_joins
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[CONCEPT DETAIL - MERGE NODE]"
 
@@ -37,7 +36,8 @@ def deduplicate_nodes(
     set_map: dict[str, set[str]] = {}
     for k, v in merged.items():
         unique_outputs = [
-            environment.concepts[x.address].address
+            # the concept may be a in a different environment for a rowset.
+            (environment.concepts.get(x.address) or x).address
             for x in v.output_concepts
             if x not in v.partial_concepts
         ]
@@ -330,8 +330,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 +340,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"),

trilogy/core/processing/nodes/select_node_v2.py

@@ -1,25 +1,23 @@
 from typing import List, Optional
 
-
 from trilogy.constants import logger
 from trilogy.core.constants import CONSTANT_DATASET
 from trilogy.core.enums import Purpose, PurposeLineage
 from trilogy.core.models import (
+    Comparison,
+    Concept,
+    Conditional,
+    Datasource,
+    Environment,
     Function,
     Grain,
+    Parenthetical,
     QueryDatasource,
     SourceType,
-    Concept,
-    Environment,
     UnnestJoin,
-    Datasource,
-    Conditional,
-    Comparison,
-    Parenthetical,
 )
-from trilogy.utility import unique
 from trilogy.core.processing.nodes.base_node import StrategyNode, resolve_concept_map
-
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[CONCEPT DETAIL - SELECT NODE]"
 
@@ -69,6 +67,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:
@@ -98,6 +101,7 @@ class SelectNode(StrategyNode):
                 PurposeLineage.BASIC,
                 PurposeLineage.ROWSET,
                 PurposeLineage.BASIC,
+                PurposeLineage.UNION,
             ):
                 source_map[x.address] = set()