pytrilogy 0.0.2.47__py3-none-any.whl → 0.0.2.48__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,6 +37,7 @@ 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)
 
     if not node:

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]

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]}"

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
         ]

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]"
 
@@ -98,6 +96,7 @@ class SelectNode(StrategyNode):
                 PurposeLineage.BASIC,
                 PurposeLineage.ROWSET,
                 PurposeLineage.BASIC,
+                PurposeLineage.UNION,
             ):
                 source_map[x.address] = set()
 

trilogy/core/processing/nodes/union_node.py

@@ -0,0 +1,50 @@
+from typing import List
+
+from trilogy.core.models import (
+    Concept,
+    QueryDatasource,
+    SourceType,
+)
+from trilogy.core.processing.nodes.base_node import StrategyNode
+
+
+class UnionNode(StrategyNode):
+    """Union nodes represent combining two keyspaces"""
+
+    source_type = SourceType.UNION
+
+    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,
+    ):
+        super().__init__(
+            input_concepts=input_concepts,
+            output_concepts=output_concepts,
+            environment=environment,
+            g=g,
+            whole_grain=whole_grain,
+            parents=parents,
+            depth=depth,
+        )
+
+    def _resolve(self) -> QueryDatasource:
+        """We need to ensure that any filtered values are removed from the output to avoid inappropriate references"""
+        base = super()._resolve()
+        return base
+
+    def copy(self) -> "UnionNode":
+        return UnionNode(
+            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

@@ -1,12 +1,11 @@
 from typing import List
 
-
 from trilogy.core.models import (
+    Concept,
+    Function,
     QueryDatasource,
     SourceType,
-    Concept,
     UnnestJoin,
-    Function,
 )
 from trilogy.core.processing.nodes.base_node import StrategyNode
 

trilogy/core/processing/nodes/window_node.py

@@ -1,8 +1,7 @@
 from typing import List
 
-
-from trilogy.core.models import SourceType, Concept
-from trilogy.core.processing.nodes.base_node import StrategyNode, QueryDatasource
+from trilogy.core.models import Concept, SourceType
+from trilogy.core.processing.nodes.base_node import QueryDatasource, StrategyNode
 
 
 class WindowNode(StrategyNode):

trilogy/core/processing/utility.py

@@ -1,52 +1,48 @@
-from typing import List, Tuple, Dict, Set, Any
+from dataclasses import dataclass
+from enum import Enum
+from logging import Logger
+from typing import Any, Dict, List, Set, Tuple
+
 import networkx as nx
+
+from trilogy.core.enums import BooleanOperator, FunctionClass, Granularity, Purpose
 from trilogy.core.models import (
-    Datasource,
-    JoinType,
+    CTE,
+    AggregateWrapper,
     BaseJoin,
+    CaseElse,
+    CaseWhen,
+    Comparison,
     Concept,
-    QueryDatasource,
-    LooseConceptList,
-    Environment,
+    ConceptPair,
     Conditional,
-    SubselectComparison,
-    Comparison,
-    Parenthetical,
-    Function,
-    FilterItem,
-    MagicConstants,
-    WindowItem,
-    AggregateWrapper,
+    Datasource,
     DataType,
-    ConceptPair,
-    UnnestJoin,
-    CaseWhen,
-    CaseElse,
-    MapWrapper,
-    ListWrapper,
-    MapType,
     DatePart,
-    NumericType,
+    Environment,
+    FilterItem,
+    Function,
+    JoinType,
     ListType,
-    TupleWrapper,
-    CTE,
+    ListWrapper,
+    LooseConceptList,
+    MagicConstants,
+    MapType,
+    MapWrapper,
     MultiSelectStatement,
-    SelectStatement,
+    NumericType,
+    Parenthetical,
     ProcessedQuery,
+    QueryDatasource,
+    SelectStatement,
+    SubselectComparison,
+    TupleWrapper,
+    UnionCTE,
+    UnnestJoin,
+    WindowItem,
 )
-
-from trilogy.core.enums import Purpose, Granularity, BooleanOperator
-from enum import Enum
 from trilogy.utility import unique
 
-from logging import Logger
-
-
-from trilogy.core.enums import FunctionClass
-
-
-from dataclasses import dataclass
-
 
 class NodeType(Enum):
     CONCEPT = 1
@@ -314,7 +310,6 @@ def get_node_joins(
     environment: Environment,
     # concepts:List[Concept],
 ):
-
     graph = nx.Graph()
     partials: dict[str, list[str]] = {}
     ds_node_map: dict[str, QueryDatasource | Datasource] = {}
@@ -536,7 +531,9 @@ def find_nullable_concepts(
     return list(sorted(final_nullable))
 
 
-def sort_select_output_processed(cte: CTE, query: ProcessedQuery) -> CTE:
+def sort_select_output_processed(
+    cte: CTE | UnionCTE, query: ProcessedQuery
+) -> CTE | UnionCTE:
     hidden_addresses = [c.address for c in query.hidden_columns]
     output_addresses = [
         c.address for c in query.output_columns if c.address not in hidden_addresses
@@ -552,8 +549,8 @@ def sort_select_output_processed(cte: CTE, query: ProcessedQuery) -> CTE:
 
 
 def sort_select_output(
-    cte: CTE, query: SelectStatement | MultiSelectStatement | ProcessedQuery
-) -> CTE:
+    cte: CTE | UnionCTE, query: SelectStatement | MultiSelectStatement | ProcessedQuery
+) -> CTE | UnionCTE:
     if isinstance(query, ProcessedQuery):
         return sort_select_output_processed(cte, query)
     hidden_addresses = [c.address for c in query.hidden_components]