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

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
@@ -218,7 +214,7 @@ def concept_to_relevant_joins(concepts: list[Concept]) -> List[Concept]:
             x for x in concepts if x.keys and all([key in addresses for key in x.keys])
         ]
     )
-    final = [c for c in concepts if c not in sub_props]
+    final = [c for c in concepts if c.address not in sub_props]
     return unique(final, "address")
 
 
@@ -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]

trilogy/core/query_processor.py

@@ -1,43 +1,40 @@
-from typing import List, Optional, Set, Union, Dict, Tuple
+from collections import defaultdict
+from math import ceil
+from typing import Dict, List, Optional, Set, Tuple, Union
 
-from trilogy.core.env_processor import generate_graph
-from trilogy.core.graph_models import ReferenceGraph
+from trilogy.constants import CONFIG, logger
 from trilogy.core.constants import CONSTANT_DATASET
-from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
-from trilogy.core.enums import BooleanOperator
-from trilogy.constants import CONFIG
-from trilogy.core.processing.nodes import SelectNode, StrategyNode, History
+from trilogy.core.enums import BooleanOperator, SourceType
+from trilogy.core.env_processor import generate_graph
+from trilogy.core.ergonomics import generate_cte_names
 from trilogy.core.models import (
+    CTE,
+    BaseJoin,
     Concept,
-    Environment,
-    PersistStatement,
     ConceptDeclarationStatement,
-    SelectStatement,
-    MultiSelectStatement,
-    CTE,
+    Conditional,
+    CopyStatement,
+    CTEConceptPair,
+    Datasource,
+    Environment,
+    InstantiatedUnnestJoin,
     Join,
-    UnnestJoin,
     MaterializedDataset,
+    MultiSelectStatement,
+    PersistStatement,
+    ProcessedCopyStatement,
     ProcessedQuery,
     ProcessedQueryPersist,
     QueryDatasource,
-    Datasource,
-    BaseJoin,
-    InstantiatedUnnestJoin,
-    Conditional,
-    ProcessedCopyStatement,
-    CopyStatement,
-    CTEConceptPair,
+    SelectStatement,
+    UnionCTE,
+    UnnestJoin,
 )
-
-from trilogy.utility import unique
-
-from trilogy.hooks.base_hook import BaseHook
-from trilogy.constants import logger
-from trilogy.core.ergonomics import generate_cte_names
 from trilogy.core.optimization import optimize_ctes
-from math import ceil
-from collections import defaultdict
+from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
+from trilogy.core.processing.nodes import History, SelectNode, StrategyNode
+from trilogy.hooks.base_hook import BaseHook
+from trilogy.utility import unique
 
 LOGGER_PREFIX = "[QUERY BUILD]"
 
@@ -103,7 +100,7 @@ def base_join_to_join(
 
 
 def generate_source_map(
-    query_datasource: QueryDatasource, all_new_ctes: List[CTE]
+    query_datasource: QueryDatasource, all_new_ctes: List[CTE | UnionCTE]
 ) -> Tuple[Dict[str, list[str]], Dict[str, list[str]]]:
     source_map: Dict[str, list[str]] = defaultdict(list)
     # now populate anything derived in this level
@@ -246,24 +243,44 @@ def resolve_cte_base_name_and_alias_v2(
     return None, None
 
 
-def datasource_to_ctes(
+def datasource_to_cte(
     query_datasource: QueryDatasource, name_map: dict[str, str]
-) -> List[CTE]:
-    output: List[CTE] = []
-    parents: list[CTE] = []
+) -> CTE | UnionCTE:
+    parents: list[CTE | UnionCTE] = []
+    if query_datasource.source_type == SourceType.UNION:
+        direct_parents: list[CTE | UnionCTE] = []
+        for child in query_datasource.datasources:
+            assert isinstance(child, QueryDatasource)
+            child_cte = datasource_to_cte(child, name_map=name_map)
+            direct_parents.append(child_cte)
+            parents += child_cte.parent_ctes
+        human_id = generate_cte_name(query_datasource.identifier, name_map)
+        final = UnionCTE(
+            name=human_id,
+            source=query_datasource,
+            parent_ctes=parents,
+            internal_ctes=direct_parents,
+            output_columns=[
+                c.with_grain(query_datasource.grain)
+                for c in query_datasource.output_concepts
+            ],
+            grain=direct_parents[0].grain,
+        )
+        return final
+
     if len(query_datasource.datasources) > 1 or any(
         [isinstance(x, QueryDatasource) for x in query_datasource.datasources]
     ):
-        all_new_ctes: List[CTE] = []
+        all_new_ctes: List[CTE | UnionCTE] = []
         for datasource in query_datasource.datasources:
             if isinstance(datasource, QueryDatasource):
                 sub_datasource = datasource
             else:
                 sub_datasource = datasource_to_query_datasource(datasource)
 
-            sub_cte = datasource_to_ctes(sub_datasource, name_map)
-            parents += sub_cte
-            all_new_ctes += sub_cte
+            sub_cte = datasource_to_cte(sub_datasource, name_map)
+            parents.append(sub_cte)
+            all_new_ctes.append(sub_cte)
         source_map, existence_map = generate_source_map(query_datasource, all_new_ctes)
 
     else:
@@ -284,7 +301,10 @@ def datasource_to_ctes(
 
     human_id = generate_cte_name(query_datasource.identifier, name_map)
 
-    final_joins = [base_join_to_join(join, parents) for join in query_datasource.joins]
+    final_joins = [
+        base_join_to_join(join, [x for x in parents if isinstance(x, CTE)])
+        for join in query_datasource.joins
+    ]
 
     base_name, base_alias = resolve_cte_base_name_and_alias_v2(
         human_id, query_datasource, source_map, final_joins
@@ -326,17 +346,18 @@ def datasource_to_ctes(
                 f"Missing {x.address} in {cte.source_map}, source map {cte.source.source_map.keys()} "
             )
 
-    output.append(cte)
-    return output
+    return cte
 
 
 def get_query_node(
     environment: Environment,
     statement: SelectStatement | MultiSelectStatement,
-    graph: Optional[ReferenceGraph] = None,
     history: History | None = None,
 ) -> StrategyNode:
-    graph = graph or generate_graph(environment)
+    environment = environment.duplicate()
+    for k, v in statement.local_concepts.items():
+        environment.concepts[k] = v
+    graph = generate_graph(environment)
     logger.info(
         f"{LOGGER_PREFIX} getting source datasource for query with filtering {statement.where_clause_category} and output {[str(c) for c in statement.output_components]}"
     )
@@ -380,11 +401,10 @@ def get_query_node(
 def get_query_datasources(
     environment: Environment,
     statement: SelectStatement | MultiSelectStatement,
-    graph: Optional[ReferenceGraph] = None,
     hooks: Optional[List[BaseHook]] = None,
 ) -> QueryDatasource:
 
-    ds = get_query_node(environment, statement, graph)
+    ds = get_query_node(environment, statement)
     final_qds = ds.resolve()
     if hooks:
         for hook in hooks:
@@ -393,7 +413,7 @@ def get_query_datasources(
     return final_qds
 
 
-def flatten_ctes(input: CTE) -> list[CTE]:
+def flatten_ctes(input: CTE | UnionCTE) -> list[CTE | UnionCTE]:
     output = [input]
     for cte in input.parent_ctes:
         output += flatten_ctes(cte)
@@ -456,18 +476,17 @@ def process_query(
     hooks: List[BaseHook] | None = None,
 ) -> ProcessedQuery:
     hooks = hooks or []
-    statement.refresh_bindings(environment)
-    graph = generate_graph(environment)
+
     root_datasource = get_query_datasources(
-        environment=environment, graph=graph, statement=statement, hooks=hooks
+        environment=environment, statement=statement, hooks=hooks
     )
     for hook in hooks:
         hook.process_root_datasource(root_datasource)
     # this should always return 1 - TODO, refactor
-    root_cte = datasource_to_ctes(root_datasource, environment.cte_name_map)[0]
+    root_cte = datasource_to_cte(root_datasource, environment.cte_name_map)
     for hook in hooks:
         hook.process_root_cte(root_cte)
-    raw_ctes: List[CTE] = list(reversed(flatten_ctes(root_cte)))
+    raw_ctes: List[CTE | UnionCTE] = list(reversed(flatten_ctes(root_cte)))
     seen = dict()
     # we can have duplicate CTEs at this point
     # so merge them together
@@ -479,7 +498,7 @@ def process_query(
             seen[cte.name] = seen[cte.name] + cte
     for cte in raw_ctes:
         cte.parent_ctes = [seen[x.name] for x in cte.parent_ctes]
-    deduped_ctes: List[CTE] = list(seen.values())
+    deduped_ctes: List[CTE | UnionCTE] = list(seen.values())
     root_cte.order_by = statement.order_by
     root_cte.limit = statement.limit
     root_cte.hidden_concepts = [x for x in statement.hidden_components]
@@ -497,4 +516,5 @@ def process_query(
         # we no longer do any joins at final level, this should always happen in parent CTEs
         joins=[],
         hidden_columns=[x for x in statement.hidden_components],
+        local_concepts=statement.local_concepts,
     )

trilogy/dialect/base.py

@@ -1,66 +1,67 @@
-from typing import List, Union, Optional, Dict, Any, Sequence, Callable
+from typing import Any, Callable, Dict, List, Optional, Sequence, Union
 
 from jinja2 import Template
 
-from trilogy.core.processing.utility import (
-    is_scalar_condition,
-    decompose_condition,
-    sort_select_output,
-)
-from trilogy.constants import CONFIG, logger, MagicConstants
-from trilogy.core.internal import DEFAULT_CONCEPTS
+from trilogy.constants import CONFIG, MagicConstants, logger
 from trilogy.core.enums import (
+    DatePart,
     FunctionType,
+    UnnestMode,
     WindowType,
-    DatePart,
 )
+from trilogy.core.internal import DEFAULT_CONCEPTS
 from trilogy.core.models import (
-    ListType,
-    DataType,
-    Concept,
     CTE,
-    ProcessedQuery,
-    ProcessedQueryPersist,
-    ProcessedShowStatement,
+    AggregateWrapper,
+    CaseElse,
+    CaseWhen,
+    Comparison,
     CompiledCTE,
+    Concept,
+    ConceptDeclarationStatement,
     Conditional,
-    Comparison,
-    SubselectComparison,
-    OrderItem,
-    WindowItem,
+    CopyStatement,
+    Datasource,
+    DataType,
+    Environment,
     FilterItem,
     Function,
-    AggregateWrapper,
-    Parenthetical,
-    CaseWhen,
-    CaseElse,
-    SelectStatement,
-    PersistStatement,
-    Environment,
-    RawColumnExpr,
+    ImportStatement,
+    ListType,
     ListWrapper,
-    TupleWrapper,
+    MapType,
     MapWrapper,
-    ShowStatement,
-    RowsetItem,
+    MergeStatementV2,
     MultiSelectStatement,
-    RowsetDerivationStatement,
-    ConceptDeclarationStatement,
-    ImportStatement,
-    RawSQLStatement,
-    ProcessedRawSQLStatement,
     NumericType,
-    MapType,
-    StructType,
-    MergeStatementV2,
-    Datasource,
-    CopyStatement,
+    OrderItem,
+    Parenthetical,
+    PersistStatement,
     ProcessedCopyStatement,
+    ProcessedQuery,
+    ProcessedQueryPersist,
+    ProcessedRawSQLStatement,
+    ProcessedShowStatement,
+    RawColumnExpr,
+    RawSQLStatement,
+    RowsetDerivationStatement,
+    RowsetItem,
+    SelectStatement,
+    ShowStatement,
+    StructType,
+    SubselectComparison,
+    TupleWrapper,
+    UnionCTE,
+    WindowItem,
 )
-from trilogy.core.query_processor import process_query, process_persist, process_copy
+from trilogy.core.processing.utility import (
+    decompose_condition,
+    is_scalar_condition,
+    sort_select_output,
+)
+from trilogy.core.query_processor import process_copy, process_persist, process_query
 from trilogy.dialect.common import render_join, render_unnest
 from trilogy.hooks.base_hook import BaseHook
-from trilogy.core.enums import UnnestMode
 
 LOGGER_PREFIX = "[RENDERING]"
 
@@ -130,7 +131,7 @@ FUNCTION_MAP = {
     FunctionType.SPLIT: lambda x: f"split({x[0]}, {x[1]})",
     FunctionType.IS_NULL: lambda x: f"isnull({x[0]})",
     FunctionType.BOOL: lambda x: f"CASE WHEN {x[0]} THEN TRUE ELSE FALSE END",
-    # complex
+    # Complex
     FunctionType.INDEX_ACCESS: lambda x: f"{x[0]}[{x[1]}]",
     FunctionType.MAP_ACCESS: lambda x: f"{x[0]}[{x[1]}][1]",
     FunctionType.UNNEST: lambda x: f"unnest({x[0]})",
@@ -230,7 +231,7 @@ def safe_quote(string: str, quote_char: str):
     return ".".join([f"{quote_char}{string}{quote_char}" for string in components])
 
 
-def safe_get_cte_value(coalesce, cte: CTE, c: Concept, quote_char: str):
+def safe_get_cte_value(coalesce, cte: CTE | UnionCTE, c: Concept, quote_char: str):
     address = c.address
     raw = cte.source_map.get(address, None)
 
@@ -255,15 +256,26 @@ class BaseDialect:
     UNNEST_MODE = UnnestMode.CROSS_APPLY
 
     def render_order_item(
-        self, order_item: OrderItem, cte: CTE, final: bool = False
+        self,
+        order_item: OrderItem,
+        cte: CTE | UnionCTE,
+        final: bool = False,
+        alias: bool = True,
     ) -> str:
         if final:
+            if not alias:
+                return f"{self.QUOTE_CHARACTER}{order_item.expr.safe_address}{self.QUOTE_CHARACTER} {order_item.order.value}"
+
             return f"{cte.name}.{self.QUOTE_CHARACTER}{order_item.expr.safe_address}{self.QUOTE_CHARACTER} {order_item.order.value}"
 
         return f"{self.render_concept_sql(order_item.expr, cte=cte, alias=False)} {order_item.order.value}"
 
     def render_concept_sql(
-        self, c: Concept, cte: CTE, alias: bool = True, raise_invalid: bool = False
+        self,
+        c: Concept,
+        cte: CTE | UnionCTE,
+        alias: bool = True,
+        raise_invalid: bool = False,
     ) -> str:
         result = None
         if c.pseudonyms:
@@ -290,7 +302,7 @@ class BaseDialect:
         return result
 
     def _render_concept_sql(
-        self, c: Concept, cte: CTE, raise_invalid: bool = False
+        self, c: Concept, cte: CTE | UnionCTE, raise_invalid: bool = False
     ) -> str:
         # only recurse while it's in sources of the current cte
         logger.debug(
@@ -348,6 +360,20 @@ class BaseDialect:
                         " target grain"
                     )
                     rval = f"{self.FUNCTION_GRAIN_MATCH_MAP[c.lineage.function.operator](args)}"
+            elif (
+                isinstance(c.lineage, Function)
+                and c.lineage.operator == FunctionType.UNION
+            ):
+                local_matched = [
+                    x
+                    for x in c.lineage.arguments
+                    if isinstance(x, Concept) and x.address in cte.output_columns
+                ]
+                if not local_matched:
+                    raise SyntaxError(
+                        "Could not find appropriate source element for union"
+                    )
+                rval = self.render_expr(local_matched[0], cte)
             elif (
                 isinstance(c.lineage, Function)
                 and c.lineage.operator == FunctionType.CONSTANT
@@ -447,13 +473,11 @@ class BaseDialect:
             FilterItem,
             # FilterItem
         ],
-        cte: Optional[CTE] = None,
-        cte_map: Optional[Dict[str, CTE]] = None,
+        cte: Optional[CTE | UnionCTE] = None,
+        cte_map: Optional[Dict[str, CTE | UnionCTE]] = None,
         raise_invalid: bool = False,
     ) -> str:
-
         if isinstance(e, SubselectComparison):
-
             if isinstance(e.right, Concept):
                 # we won't always have an existnce map
                 # so fall back to the normal map
@@ -559,6 +583,13 @@ class BaseDialect:
         elif isinstance(e, FilterItem):
             return f"CASE WHEN {self.render_expr(e.where.conditional,cte=cte, cte_map=cte_map, raise_invalid=raise_invalid)} THEN {self.render_expr(e.content, cte, cte_map=cte_map, raise_invalid=raise_invalid)} ELSE NULL END"
         elif isinstance(e, Concept):
+            if (
+                isinstance(e.lineage, Function)
+                and e.lineage.operator == FunctionType.CONSTANT
+                and CONFIG.rendering.parameters is True
+                and e.datatype.data_type != DataType.MAP
+            ):
+                return f":{e.safe_address}"
             if cte:
                 return self.render_concept_sql(
                     e, cte, alias=False, raise_invalid=raise_invalid
@@ -592,7 +623,18 @@ class BaseDialect:
         else:
             raise ValueError(f"Unable to render type {type(e)} {e}")
 
-    def render_cte(self, cte: CTE, auto_sort: bool = True) -> CompiledCTE:
+    def render_cte(self, cte: CTE | UnionCTE, auto_sort: bool = True) -> CompiledCTE:
+        if isinstance(cte, UnionCTE):
+            base_statement = f"\n{cte.operator}\n".join(
+                [self.render_cte(child).statement for child in cte.internal_ctes]
+            )
+            if cte.order_by:
+                ordering = [
+                    self.render_order_item(i, cte, final=True, alias=False)
+                    for i in cte.order_by.items
+                ]
+                base_statement += "\nORDER BY " + ",".join(ordering)
+            return CompiledCTE(name=cte.name, statement=base_statement)
         if self.UNNEST_MODE in (
             UnnestMode.CROSS_APPLY,
             UnnestMode.CROSS_JOIN,
@@ -863,7 +905,7 @@ class BaseDialect:
         if CONFIG.strict_mode and INVALID_REFERENCE_STRING(1) in final:
             raise ValueError(
                 f"Invalid reference string found in query: {final}, this should never"
-                " occur. Please report this issue."
+                " occur. Please create a GitHub issue to report this."
             )
         logger.info(f"{LOGGER_PREFIX} Compiled query: {final}")
         return final

trilogy/dialect/bigquery.py

@@ -1,11 +1,10 @@
-from typing import Mapping, Callable, Any
+from typing import Any, Callable, Mapping
 
 from jinja2 import Template
 
-from trilogy.core.enums import FunctionType, WindowType, UnnestMode
+from trilogy.core.enums import FunctionType, UnnestMode, WindowType
 from trilogy.dialect.base import BaseDialect
 
-
 WINDOW_FUNCTION_MAP: Mapping[WindowType, Callable[[Any, Any, Any], str]] = {}
 
 FUNCTION_MAP = {