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

trilogy/core/optimization.py

@@ -1,17 +1,18 @@
+from trilogy.constants import CONFIG, logger
+from trilogy.core.enums import BooleanOperator, PurposeLineage
 from trilogy.core.models import (
     CTE,
-    SelectStatement,
-    MultiSelectStatement,
     Conditional,
+    MultiSelectStatement,
+    SelectStatement,
+    UnionCTE,
 )
-from trilogy.core.enums import PurposeLineage, BooleanOperator
-from trilogy.constants import logger, CONFIG
 from trilogy.core.optimizations import (
-    OptimizationRule,
     InlineConstant,
+    InlineDatasource,
+    OptimizationRule,
     PredicatePushdown,
     PredicatePushdownRemove,
-    InlineDatasource,
 )
 from trilogy.core.processing.utility import sort_select_output
 
@@ -61,12 +62,12 @@ def reorder_ctes(
 
 
 def filter_irrelevant_ctes(
-    input: list[CTE],
-    root_cte: CTE,
+    input: list[CTE | UnionCTE],
+    root_cte: CTE | UnionCTE,
 ):
     relevant_ctes = set()
 
-    def recurse(cte: CTE, inverse_map: dict[str, list[CTE]]):
+    def recurse(cte: CTE | UnionCTE, inverse_map: dict[str, list[CTE | UnionCTE]]):
         # TODO: revisit this
         # if parent := is_locally_irrelevant(cte):
         #     logger.info(
@@ -88,6 +89,9 @@ def filter_irrelevant_ctes(
         relevant_ctes.add(cte.name)
         for cte in cte.parent_ctes:
             recurse(cte, inverse_map)
+        if isinstance(cte, UnionCTE):
+            for cte in cte.internal_ctes:
+                recurse(cte, inverse_map)
 
     inverse_map = gen_inverse_map(input)
     recurse(root_cte, inverse_map)
@@ -97,22 +101,31 @@ def filter_irrelevant_ctes(
     return filter_irrelevant_ctes(final, root_cte)
 
 
-def gen_inverse_map(input: list[CTE]) -> dict[str, list[CTE]]:
-    inverse_map: dict[str, list[CTE]] = {}
+def gen_inverse_map(input: list[CTE | UnionCTE]) -> dict[str, list[CTE | UnionCTE]]:
+    inverse_map: dict[str, list[CTE | UnionCTE]] = {}
     for cte in input:
-        for parent in cte.parent_ctes:
-            if parent.name not in inverse_map:
-                inverse_map[parent.name] = []
-            inverse_map[parent.name].append(cte)
+        if isinstance(cte, UnionCTE):
+            for internal in cte.internal_ctes:
+                if internal.name not in inverse_map:
+                    inverse_map[internal.name] = []
+                inverse_map[internal.name].append(cte)
+        else:
+            for parent in cte.parent_ctes:
+                if parent.name not in inverse_map:
+                    inverse_map[parent.name] = []
+                inverse_map[parent.name].append(cte)
+
     return inverse_map
 
 
-def is_direct_return_eligible(cte: CTE) -> CTE | None:
+def is_direct_return_eligible(cte: CTE | UnionCTE) -> CTE | UnionCTE | None:
     # if isinstance(select, (PersistStatement, MultiSelectStatement)):
     #     return False
     if len(cte.parent_ctes) != 1:
         return None
     direct_parent = cte.parent_ctes[0]
+    if isinstance(direct_parent, UnionCTE):
+        return None
 
     output_addresses = set([x.address for x in cte.output_columns])
     parent_output_addresses = set([x.address for x in direct_parent.output_columns])
@@ -120,6 +133,8 @@ def is_direct_return_eligible(cte: CTE) -> CTE | None:
         return None
     if not direct_parent.grain == cte.grain:
         return None
+
+    assert isinstance(cte, CTE)
     derived_concepts = [
         c
         for c in cte.source.output_concepts + cte.source.hidden_concepts
@@ -155,10 +170,11 @@ def is_direct_return_eligible(cte: CTE) -> CTE | None:
 
 
 def optimize_ctes(
-    input: list[CTE], root_cte: CTE, select: SelectStatement | MultiSelectStatement
-) -> list[CTE]:
-
-    direct_parent: CTE | None = root_cte
+    input: list[CTE | UnionCTE],
+    root_cte: CTE | UnionCTE,
+    select: SelectStatement | MultiSelectStatement,
+) -> list[CTE | UnionCTE]:
+    direct_parent: CTE | UnionCTE | None = root_cte
     while CONFIG.optimizations.direct_return and (
         direct_parent := is_direct_return_eligible(root_cte)
     ):
@@ -178,7 +194,7 @@ def optimize_ctes(
                 direct_parent.condition = root_cte.condition
         root_cte = direct_parent
 
-    sort_select_output(root_cte, select)
+        sort_select_output(root_cte, select)
 
     REGISTERED_RULES: list["OptimizationRule"] = []
     if CONFIG.optimizations.constant_inlining:

trilogy/core/optimizations/__init__.py

@@ -1,7 +1,7 @@
+from .base_optimization import OptimizationRule
 from .inline_constant import InlineConstant
 from .inline_datasource import InlineDatasource
 from .predicate_pushdown import PredicatePushdown, PredicatePushdownRemove
-from .base_optimization import OptimizationRule
 
 __all__ = [
     "OptimizationRule",

trilogy/core/optimizations/base_optimization.py

@@ -1,13 +1,13 @@
-from trilogy.core.models import (
-    CTE,
-)
-from trilogy.constants import logger
 from abc import ABC
 
+from trilogy.constants import logger
+from trilogy.core.models import CTE, UnionCTE
+
 
 class OptimizationRule(ABC):
-
-    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
+    def optimize(
+        self, cte: CTE | UnionCTE, inverse_map: dict[str, list[CTE | UnionCTE]]
+    ) -> bool:
         raise NotImplementedError
 
     def log(self, message: str):

trilogy/core/optimizations/inline_constant.py

@@ -1,15 +1,18 @@
+from trilogy.core.enums import PurposeLineage
 from trilogy.core.models import (
     CTE,
     Concept,
+    UnionCTE,
 )
-from trilogy.core.enums import PurposeLineage
-
 from trilogy.core.optimizations.base_optimization import OptimizationRule
 
 
 class InlineConstant(OptimizationRule):
-
-    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
+    def optimize(
+        self, cte: CTE | UnionCTE, inverse_map: dict[str, list[CTE | UnionCTE]]
+    ) -> bool:
+        if isinstance(cte, UnionCTE):
+            return any(self.optimize(x, inverse_map) for x in cte.internal_ctes)
 
         to_inline: list[Concept] = []
         for x in cte.source.input_concepts:

trilogy/core/optimizations/inline_datasource.py

@@ -1,21 +1,28 @@
+from collections import defaultdict
+
+from trilogy.constants import CONFIG
 from trilogy.core.models import (
     CTE,
     Datasource,
+    UnionCTE,
 )
-
 from trilogy.core.optimizations.base_optimization import OptimizationRule
-from collections import defaultdict
-from trilogy.constants import CONFIG
 
 
 class InlineDatasource(OptimizationRule):
-
     def __init__(self):
         super().__init__()
         self.candidates = defaultdict(lambda: set())
         self.count = defaultdict(lambda: 0)
 
-    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
+    def optimize(
+        self, cte: CTE | UnionCTE, inverse_map: dict[str, list[CTE | UnionCTE]]
+    ) -> bool:
+        if isinstance(cte, UnionCTE):
+            return any(
+                self.optimize(x, inverse_map=inverse_map) for x in cte.internal_ctes
+            )
+
         if not cte.parent_ctes:
             return False
 
@@ -25,6 +32,8 @@ class InlineDatasource(OptimizationRule):
         to_inline: list[CTE] = []
         force_group = False
         for parent_cte in cte.parent_ctes:
+            if isinstance(parent_cte, UnionCTE):
+                continue
             if not parent_cte.is_root_datasource:
                 self.debug(f"parent {parent_cte.name} is not root")
                 continue

trilogy/core/optimizations/predicate_pushdown.py

@@ -1,11 +1,12 @@
 from trilogy.core.models import (
     CTE,
-    Conditional,
     BooleanOperator,
-    Datasource,
-    ConceptArgs,
     Comparison,
+    ConceptArgs,
+    Conditional,
+    Datasource,
     Parenthetical,
+    UnionCTE,
     WindowItem,
 )
 from trilogy.core.optimizations.base_optimization import OptimizationRule
@@ -25,20 +26,21 @@ def is_child_of(a, comparison):
 
 
 class PredicatePushdown(OptimizationRule):
-
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.complete: dict[str, bool] = {}
 
     def _check_parent(
         self,
-        cte: CTE,
-        parent_cte: CTE,
+        cte: CTE | UnionCTE,
+        parent_cte: CTE | UnionCTE,
         candidate: Conditional | Comparison | Parenthetical | None,
-        inverse_map: dict[str, list[CTE]],
+        inverse_map: dict[str, list[CTE | UnionCTE]],
     ):
         if not isinstance(candidate, ConceptArgs):
             return False
+        if not isinstance(parent_cte, CTE):
+            return False
         row_conditions = {x.address for x in candidate.row_arguments}
         existence_conditions = {
             y.address for x in candidate.existence_arguments for y in x
@@ -112,7 +114,12 @@ class PredicatePushdown(OptimizationRule):
         )
         return False
 
-    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
+    def optimize(
+        self, cte: CTE | UnionCTE, inverse_map: dict[str, list[CTE | UnionCTE]]
+    ) -> bool:
+        # TODO - pushdown through unions
+        if isinstance(cte, UnionCTE):
+            return False
         optimized = False
 
         if not cte.parent_ctes:
@@ -167,12 +174,15 @@ class PredicatePushdown(OptimizationRule):
 
 
 class PredicatePushdownRemove(OptimizationRule):
-
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.complete: dict[str, bool] = {}
 
-    def optimize(self, cte: CTE, inverse_map: dict[str, list[CTE]]) -> bool:
+    def optimize(
+        self, cte: CTE | UnionCTE, inverse_map: dict[str, list[CTE | UnionCTE]]
+    ) -> bool:
+        if isinstance(cte, UnionCTE):
+            return False
         optimized = False
 
         if not cte.parent_ctes:

trilogy/core/processing/concept_strategies_v3.py

@@ -1,41 +1,42 @@
 from collections import defaultdict
+from enum import Enum
 from typing import List, Optional, Protocol, Union
 
 from trilogy.constants import logger
-from trilogy.core.enums import PurposeLineage, Granularity, FunctionType
+from trilogy.core.enums import FunctionType, Granularity, PurposeLineage
 from trilogy.core.env_processor import generate_graph
 from trilogy.core.graph_models import ReferenceGraph
 from trilogy.core.models import (
     Concept,
     Environment,
     Function,
-    WhereClause,
     RowsetItem,
+    UndefinedConcept,
+    WhereClause,
 )
-from trilogy.core.processing.utility import (
-    get_disconnected_components,
+from trilogy.core.processing.node_generators import (
+    gen_basic_node,
+    gen_filter_node,
+    gen_group_node,
+    gen_group_to_node,
+    gen_merge_node,
+    gen_multiselect_node,
+    gen_rowset_node,
+    gen_union_node,
+    gen_unnest_node,
+    gen_window_node,
 )
-from trilogy.utility import unique
 from trilogy.core.processing.nodes import (
     ConstantNode,
-    MergeNode,
     GroupNode,
-    StrategyNode,
     History,
+    MergeNode,
+    StrategyNode,
 )
-from trilogy.core.processing.node_generators import (
-    gen_filter_node,
-    gen_window_node,
-    gen_group_node,
-    gen_basic_node,
-    gen_unnest_node,
-    gen_merge_node,
-    gen_group_to_node,
-    gen_rowset_node,
-    gen_multiselect_node,
+from trilogy.core.processing.utility import (
+    get_disconnected_components,
 )
-
-from enum import Enum
+from trilogy.utility import unique
 
 
 class ValidationResult(Enum):
@@ -107,6 +108,9 @@ def get_priority_concept(
             +
             # then rowsets to remove them from scope, as they cannot get partials
             [c for c in remaining_concept if c.derivation == PurposeLineage.ROWSET]
+            +
+            # then rowsets to remove them from scope, as they cannot get partials
+            [c for c in remaining_concept if c.derivation == PurposeLineage.UNION]
             # we should be home-free here
             +
             # then aggregates to remove them from scope, as they cannot get partials
@@ -250,6 +254,20 @@ def generate_node(
             history,
             conditions=conditions,
         )
+    elif concept.derivation == PurposeLineage.UNION:
+        logger.info(
+            f"{depth_to_prefix(depth)}{LOGGER_PREFIX} for {concept.address}, generating union node with optional {[x.address for x in local_optional]} and condition {conditions}"
+        )
+        return gen_union_node(
+            concept,
+            local_optional,
+            environment,
+            g,
+            depth + 1,
+            source_concepts,
+            history,
+            conditions=conditions,
+        )
     elif concept.derivation == PurposeLineage.AGGREGATE:
         # don't push constants up before aggregation
         # if not required
@@ -495,11 +513,9 @@ def validate_concept(
     seen: set[str],
     environment: Environment,
 ):
-
     found_map[str(node)].add(concept)
     seen.add(concept.address)
     if concept not in node.partial_concepts:
-
         found_addresses.add(concept.address)
         non_partial_addresses.add(concept.address)
         # remove it from our partial tracking
@@ -658,7 +674,6 @@ def search_concepts(
     history: History | None = None,
     conditions: WhereClause | None = None,
 ) -> StrategyNode | None:
-
     history = history or History()
     hist = history.get_history(
         search=mandatory_list, accept_partial=accept_partial, conditions=conditions
@@ -698,10 +713,11 @@ def _search_concepts(
     accept_partial: bool = False,
     conditions: WhereClause | None = None,
 ) -> StrategyNode | None:
-
     # these are the concepts we need in the output projection
     mandatory_list = unique(mandatory_list, "address")
-
+    for x in mandatory_list:
+        if isinstance(x, UndefinedConcept):
+            raise SyntaxError(f"Undefined concept {x.address}")
     all_mandatory = set(c.address for c in mandatory_list)
 
     must_evaluate_condition_on_this_level_not_push_down = False

trilogy/core/processing/graph_utils.py

@@ -1,18 +1,17 @@
+from collections import defaultdict
 from typing import Dict, List
+
 from trilogy.core.models import Concept
-from collections import defaultdict
 from trilogy.utility import unique
 
 
 def extract_required_subgraphs(
     assocs: defaultdict[str, list], path: List[str]
 ) -> defaultdict[str, list]:
-
     ds = path[0]
     current: list[str] = []
     for idx, val in enumerate(path):
         if val.startswith("ds~"):
-
             if current:
                 assocs[ds] += current
                 current = [path[idx - 1]] if idx > 0 else []

trilogy/core/processing/node_generators/__init__.py

@@ -1,13 +1,14 @@
+from .basic_node import gen_basic_node
 from .filter_node import gen_filter_node
-from .window_node import gen_window_node
 from .group_node import gen_group_node
 from .group_to_node import gen_group_to_node
-from .basic_node import gen_basic_node
-from .select_node import gen_select_node
-from .unnest_node import gen_unnest_node
+from .multiselect_node import gen_multiselect_node
 from .node_merge_node import gen_merge_node
 from .rowset_node import gen_rowset_node
-from .multiselect_node import gen_multiselect_node
+from .select_node import gen_select_node
+from .union_node import gen_union_node
+from .unnest_node import gen_unnest_node
+from .window_node import gen_window_node
 
 __all__ = [
     "gen_filter_node",
@@ -16,6 +17,7 @@ __all__ = [
     "gen_select_node",
     "gen_basic_node",
     "gen_unnest_node",
+    "gen_union_node",
     "gen_merge_node",
     "gen_group_to_node",
     "gen_rowset_node",

trilogy/core/processing/node_generators/basic_node.py

@@ -1,18 +1,18 @@
 # directly select out a basic derivation
 from typing import List
 
+from trilogy.constants import logger
+from trilogy.core.enums import SourceType
 from trilogy.core.models import (
     Concept,
-    WhereClause,
     Function,
     FunctionClass,
+    WhereClause,
 )
-from trilogy.core.processing.nodes import StrategyNode, History
 from trilogy.core.processing.node_generators.common import (
     resolve_function_parent_concepts,
 )
-from trilogy.constants import logger
-from trilogy.core.enums import SourceType
+from trilogy.core.processing.nodes import History, StrategyNode
 
 LOGGER_PREFIX = "[GEN_BASIC_NODE]"
 

trilogy/core/processing/node_generators/common.py

@@ -1,23 +1,23 @@
-from typing import List, Tuple, Callable
+from collections import defaultdict
+from typing import Callable, List, Tuple
 
-from trilogy.core.enums import PurposeLineage, Purpose
+from trilogy.core.enums import Purpose, PurposeLineage
 from trilogy.core.models import (
-    Concept,
-    Function,
     AggregateWrapper,
-    FilterItem,
+    Concept,
     Environment,
+    FilterItem,
+    Function,
     LooseConceptList,
     WhereClause,
 )
-from trilogy.utility import unique
-from trilogy.core.processing.nodes.base_node import StrategyNode
-from trilogy.core.processing.nodes.merge_node import MergeNode
-from trilogy.core.processing.nodes import History
 from trilogy.core.processing.nodes import (
+    History,
     NodeJoin,
 )
-from collections import defaultdict
+from trilogy.core.processing.nodes.base_node import StrategyNode
+from trilogy.core.processing.nodes.merge_node import MergeNode
+from trilogy.utility import unique
 
 
 def resolve_function_parent_concepts(concept: Concept) -> List[Concept]:
@@ -151,7 +151,6 @@ def gen_enrichment_node(
     history: History | None = None,
     conditions: WhereClause | None = None,
 ):
-
     local_opts = LooseConceptList(concepts=local_optional)
 
     extra_required = [

trilogy/core/processing/node_generators/filter_node.py

@@ -1,20 +1,18 @@
 from typing import List
 
-
+from trilogy.constants import logger
 from trilogy.core.models import Concept, Environment, FilterItem, Grain, WhereClause
+from trilogy.core.processing.node_generators.common import (
+    resolve_filter_parent_concepts,
+)
 from trilogy.core.processing.nodes import (
     FilterNode,
-    MergeNode,
     History,
-    StrategyNode,
+    MergeNode,
     SelectNode,
+    StrategyNode,
 )
-from trilogy.core.processing.node_generators.common import (
-    resolve_filter_parent_concepts,
-)
-from trilogy.constants import logger
-from trilogy.core.processing.utility import padding, unique
-from trilogy.core.processing.utility import is_scalar_condition
+from trilogy.core.processing.utility import is_scalar_condition, padding, unique
 
 LOGGER_PREFIX = "[GEN_FILTER_NODE]"
 

trilogy/core/processing/node_generators/group_node.py

@@ -1,23 +1,22 @@
+from typing import List
+
+from trilogy.constants import logger
 from trilogy.core.models import (
+    AggregateWrapper,
     Concept,
     Environment,
-    LooseConceptList,
-    WhereClause,
     Function,
-    AggregateWrapper,
     Grain,
+    LooseConceptList,
+    WhereClause,
 )
-from trilogy.utility import unique
-from trilogy.core.processing.nodes import GroupNode, StrategyNode, History
-from typing import List
-from trilogy.core.processing.node_generators.common import (
-    resolve_function_parent_concepts,
-)
-from trilogy.constants import logger
-from trilogy.core.processing.utility import padding, create_log_lambda
 from trilogy.core.processing.node_generators.common import (
     gen_enrichment_node,
+    resolve_function_parent_concepts,
 )
+from trilogy.core.processing.nodes import GroupNode, History, StrategyNode
+from trilogy.core.processing.utility import create_log_lambda, padding
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[GEN_GROUP_NODE]"
 

trilogy/core/processing/node_generators/group_to_node.py

@@ -1,13 +1,13 @@
+from typing import List
+
+from trilogy.constants import logger
 from trilogy.core.models import Concept, Environment, Function, WhereClause
 from trilogy.core.processing.nodes import (
     GroupNode,
-    StrategyNode,
-    MergeNode,
     History,
+    MergeNode,
+    StrategyNode,
 )
-from typing import List
-
-from trilogy.constants import logger
 from trilogy.core.processing.utility import padding
 
 LOGGER_PREFIX = "[GEN_GROUP_TO_NODE]"

trilogy/core/processing/node_generators/multiselect_node.py

@@ -1,22 +1,20 @@
+from collections import defaultdict
+from itertools import combinations
+from typing import List
+
+from trilogy.constants import logger
+from trilogy.core.enums import BooleanOperator, JoinType, Purpose
 from trilogy.core.models import (
     Concept,
+    Conditional,
     Environment,
     MultiSelectStatement,
     WhereClause,
-    Conditional,
 )
-from trilogy.core.processing.nodes import MergeNode, NodeJoin, History
-from trilogy.core.processing.nodes.base_node import concept_list_to_grain, StrategyNode
-from typing import List
-
-from trilogy.core.enums import JoinType
-from trilogy.constants import logger
-from trilogy.core.processing.utility import padding
-from trilogy.core.processing.utility import concept_to_relevant_joins
-from collections import defaultdict
-from itertools import combinations
-from trilogy.core.enums import Purpose, BooleanOperator
 from trilogy.core.processing.node_generators.common import resolve_join_order
+from trilogy.core.processing.nodes import History, MergeNode, NodeJoin
+from trilogy.core.processing.nodes.base_node import StrategyNode, concept_list_to_grain
+from trilogy.core.processing.utility import concept_to_relevant_joins, padding
 
 LOGGER_PREFIX = "[GEN_MULTISELECT_NODE]"