pytrilogy 0.0.2.58__py3-none-any.whl → 0.0.3.0__py3-none-any.whl

trilogy/core/processing/utility.py

@@ -6,44 +6,61 @@ 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 (
-    CTE,
-    AggregateWrapper,
-    BaseJoin,
-    CaseElse,
-    CaseWhen,
-    Comparison,
-    Concept,
-    ConceptPair,
-    Conditional,
-    Datasource,
-    DataType,
+from trilogy.constants import MagicConstants
+from trilogy.core.enums import (
+    BooleanOperator,
     DatePart,
-    Environment,
-    FilterItem,
-    Function,
+    FunctionClass,
+    Granularity,
     JoinType,
+    Purpose,
+)
+from trilogy.core.models.build import (
+    BuildAggregateWrapper,
+    BuildCaseElse,
+    BuildCaseWhen,
+    BuildComparison,
+    BuildConcept,
+    BuildConditional,
+    BuildDatasource,
+    BuildFilterItem,
+    BuildFunction,
+    BuildParenthetical,
+    BuildSubselectComparison,
+    BuildWindowItem,
+    LooseBuildConceptList,
+)
+from trilogy.core.models.build_environment import BuildEnvironment
+from trilogy.core.models.core import (
+    DataType,
     ListType,
     ListWrapper,
-    LooseConceptList,
-    MagicConstants,
     MapType,
     MapWrapper,
-    MultiSelectStatement,
     NumericType,
-    Parenthetical,
-    ProcessedQuery,
-    QueryDatasource,
-    SelectStatement,
-    SubselectComparison,
     TupleWrapper,
+)
+from trilogy.core.models.execute import (
+    CTE,
+    BaseJoin,
+    ConceptPair,
+    QueryDatasource,
     UnionCTE,
     UnnestJoin,
-    WindowItem,
 )
+from trilogy.core.statements.author import MultiSelectStatement, SelectStatement
+from trilogy.core.statements.execute import ProcessedQuery
 from trilogy.utility import unique
 
+AGGREGATE_TYPES = (BuildAggregateWrapper,)
+SUBSELECT_TYPES = (BuildSubselectComparison,)
+COMPARISON_TYPES = (BuildComparison,)
+FUNCTION_TYPES = (BuildFunction,)
+PARENTHETICAL_TYPES = (BuildParenthetical,)
+CONDITIONAL_TYPES = (BuildConditional,)
+CONCEPT_TYPES = (BuildConcept,)
+WINDOW_TYPES = (BuildWindowItem,)
+
 
 class NodeType(Enum):
     CONCEPT = 1
@@ -208,11 +225,10 @@ def resolve_join_order_v2(
     return output
 
 
-def concept_to_relevant_joins(concepts: list[Concept]) -> List[Concept]:
-    addresses = LooseConceptList(concepts=concepts)
-    sub_props = LooseConceptList(
+def concept_to_relevant_joins(concepts: list[BuildConcept]) -> List[BuildConcept]:
+    sub_props = LooseBuildConceptList(
         concepts=[
-            x for x in concepts if x.keys and all([key in addresses for key in x.keys])
+            x for x in concepts if x.keys and all([key in concepts for key in x.keys])
         ]
     )
     final = [c for c in concepts if c.address not in sub_props]
@@ -233,7 +249,7 @@ def create_log_lambda(prefix: str, depth: int, logger: Logger):
 
 
 def calculate_graph_relevance(
-    g: nx.DiGraph, subset_nodes: set[str], concepts: set[Concept]
+    g: nx.DiGraph, subset_nodes: set[str], concepts: set[BuildConcept]
 ) -> int:
     """Calculate the relevance of each node in a graph
     Relevance is used to prune irrelevant nodes from the graph
@@ -268,10 +284,10 @@ def calculate_graph_relevance(
 
 def add_node_join_concept(
     graph: nx.DiGraph,
-    concept: Concept,
-    concept_map: dict[str, Concept],
+    concept: BuildConcept,
+    concept_map: dict[str, BuildConcept],
     ds_node: str,
-    environment: Environment,
+    environment: BuildEnvironment,
 ):
     name = f"c~{concept.address}"
     graph.add_node(name, type=NodeType.CONCEPT)
@@ -294,8 +310,8 @@ def add_node_join_concept(
 
 
 def resolve_instantiated_concept(
-    concept: Concept, datasource: QueryDatasource | Datasource
-) -> Concept:
+    concept: BuildConcept, datasource: QueryDatasource | BuildDatasource
+) -> BuildConcept:
     if concept.address in datasource.output_concepts:
         return concept
     for k in concept.pseudonyms:
@@ -333,14 +349,14 @@ def reduce_concept_pairs(input: list[ConceptPair]) -> list[ConceptPair]:
 
 
 def get_node_joins(
-    datasources: List[QueryDatasource | Datasource],
-    environment: Environment,
+    datasources: List[QueryDatasource | BuildDatasource],
+    environment: BuildEnvironment,
     # concepts:List[Concept],
 ) -> List[BaseJoin]:
     graph = nx.Graph()
     partials: dict[str, list[str]] = {}
-    ds_node_map: dict[str, QueryDatasource | Datasource] = {}
-    concept_map: dict[str, Concept] = {}
+    ds_node_map: dict[str, QueryDatasource | BuildDatasource] = {}
+    concept_map: dict[str, BuildConcept] = {}
     for datasource in datasources:
         ds_node = f"ds~{datasource.identifier}"
         ds_node_map[ds_node] = datasource
@@ -367,7 +383,7 @@ def get_node_joins(
             concepts=[] if not j.keys else None,
             concept_pairs=reduce_concept_pairs(
                 [
-                    ConceptPair(
+                    ConceptPair.model_construct(
                         left=resolve_instantiated_concept(
                             concept_map[concept], ds_node_map[k]
                         ),
@@ -386,7 +402,7 @@ def get_node_joins(
 
 
 def get_disconnected_components(
-    concept_map: Dict[str, Set[Concept]]
+    concept_map: Dict[str, Set[BuildConcept]]
 ) -> Tuple[int, List]:
     """Find if any of the datasources are not linked"""
     import networkx as nx
@@ -414,18 +430,18 @@ def is_scalar_condition(
         | date
         | datetime
         | list[Any]
-        | WindowItem
-        | FilterItem
-        | Concept
-        | Comparison
-        | Conditional
-        | Parenthetical
-        | Function
-        | AggregateWrapper
+        | BuildConcept
+        | BuildWindowItem
+        | BuildFilterItem
+        | BuildConditional
+        | BuildComparison
+        | BuildParenthetical
+        | BuildFunction
+        | BuildAggregateWrapper
+        | BuildCaseWhen
+        | BuildCaseElse
         | MagicConstants
         | DataType
-        | CaseWhen
-        | CaseElse
         | MapWrapper[Any, Any]
         | ListType
         | MapType
@@ -436,64 +452,76 @@ def is_scalar_condition(
     ),
     materialized: set[str] | None = None,
 ) -> bool:
-    if isinstance(element, Parenthetical):
+    if isinstance(element, PARENTHETICAL_TYPES):
         return is_scalar_condition(element.content, materialized)
-    elif isinstance(element, SubselectComparison):
+    elif isinstance(element, SUBSELECT_TYPES):
         return True
-    elif isinstance(element, Comparison):
+    elif isinstance(element, COMPARISON_TYPES):
         return is_scalar_condition(element.left, materialized) and is_scalar_condition(
             element.right, materialized
         )
-    elif isinstance(element, Function):
+    elif isinstance(element, FUNCTION_TYPES):
         if element.operator in FunctionClass.AGGREGATE_FUNCTIONS.value:
             return False
         return all([is_scalar_condition(x, materialized) for x in element.arguments])
-    elif isinstance(element, Concept):
+    elif isinstance(element, CONCEPT_TYPES):
         if materialized and element.address in materialized:
             return True
-        if element.lineage and isinstance(element.lineage, AggregateWrapper):
+        if element.lineage and isinstance(element.lineage, AGGREGATE_TYPES):
             return is_scalar_condition(element.lineage, materialized)
-        if element.lineage and isinstance(element.lineage, Function):
+        if element.lineage and isinstance(element.lineage, FUNCTION_TYPES):
             return is_scalar_condition(element.lineage, materialized)
         return True
-    elif isinstance(element, AggregateWrapper):
+    elif isinstance(element, AGGREGATE_TYPES):
         return is_scalar_condition(element.function, materialized)
-    elif isinstance(element, Conditional):
+    elif isinstance(element, CONDITIONAL_TYPES):
         return is_scalar_condition(element.left, materialized) and is_scalar_condition(
             element.right, materialized
         )
-    elif isinstance(element, CaseWhen):
+    elif isinstance(element, (BuildCaseWhen,)):
         return is_scalar_condition(
             element.comparison, materialized
         ) and is_scalar_condition(element.expr, materialized)
-    elif isinstance(element, CaseElse):
+    elif isinstance(element, (BuildCaseElse,)):
         return is_scalar_condition(element.expr, materialized)
     elif isinstance(element, MagicConstants):
         return True
     return True
 
 
+CONDITION_TYPES = (
+    BuildSubselectComparison,
+    BuildComparison,
+    BuildConditional,
+    BuildParenthetical,
+)
+
+
 def decompose_condition(
-    conditional: Conditional | Comparison | Parenthetical,
-) -> list[SubselectComparison | Comparison | Conditional | Parenthetical]:
-    chunks: list[SubselectComparison | Comparison | Conditional | Parenthetical] = []
-    if not isinstance(conditional, Conditional):
+    conditional: BuildConditional | BuildComparison | BuildParenthetical,
+) -> list[
+    BuildSubselectComparison | BuildComparison | BuildConditional | BuildParenthetical
+]:
+    chunks: list[
+        BuildSubselectComparison
+        | BuildComparison
+        | BuildConditional
+        | BuildParenthetical
+    ] = []
+    if not isinstance(conditional, BuildConditional):
         return [conditional]
     if conditional.operator == BooleanOperator.AND:
         if not (
-            isinstance(
-                conditional.left,
-                (SubselectComparison, Comparison, Conditional, Parenthetical),
-            )
+            isinstance(conditional.left, CONDITION_TYPES)
             and isinstance(
                 conditional.right,
-                (SubselectComparison, Comparison, Conditional, Parenthetical),
+                CONDITION_TYPES,
             )
         ):
             chunks.append(conditional)
         else:
             for val in [conditional.left, conditional.right]:
-                if isinstance(val, Conditional):
+                if isinstance(val, BuildConditional):
                     chunks.extend(decompose_condition(val))
                 else:
                     chunks.append(val)
@@ -503,8 +531,8 @@ def decompose_condition(
 
 
 def find_nullable_concepts(
-    source_map: Dict[str, set[Datasource | QueryDatasource | UnnestJoin]],
-    datasources: List[Datasource | QueryDatasource],
+    source_map: Dict[str, set[BuildDatasource | QueryDatasource | UnnestJoin]],
+    datasources: List[BuildDatasource | QueryDatasource],
     joins: List[BaseJoin | UnnestJoin],
 ) -> List[str]:
     """give a set of datasources and joins, find the concepts
@@ -514,7 +542,7 @@ def find_nullable_concepts(
     datasource_map = {
         x.identifier: x
         for x in datasources
-        if isinstance(x, (Datasource, QueryDatasource))
+        if isinstance(x, (BuildDatasource, QueryDatasource))
     }
     for join in joins:
         is_on_nullable_condition = False

trilogy/core/query_processor.py

@@ -7,32 +7,40 @@ from trilogy.core.constants import CONSTANT_DATASET
 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 (
+from trilogy.core.models.author import MultiSelectLineage, SelectLineage
+from trilogy.core.models.build import (
+    BuildConcept,
+    BuildConditional,
+    BuildDatasource,
+    Factory,
+)
+from trilogy.core.models.environment import Environment
+from trilogy.core.models.execute import (
     CTE,
     BaseJoin,
-    Concept,
-    ConceptDeclarationStatement,
-    Conditional,
-    CopyStatement,
     CTEConceptPair,
-    Datasource,
-    Environment,
     InstantiatedUnnestJoin,
     Join,
-    MaterializedDataset,
-    MultiSelectStatement,
-    PersistStatement,
-    ProcessedCopyStatement,
-    ProcessedQuery,
-    ProcessedQueryPersist,
     QueryDatasource,
-    SelectStatement,
     UnionCTE,
     UnnestJoin,
 )
 from trilogy.core.optimization import optimize_ctes
 from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
 from trilogy.core.processing.nodes import History, SelectNode, StrategyNode
+from trilogy.core.statements.author import (
+    ConceptDeclarationStatement,
+    CopyStatement,
+    MultiSelectStatement,
+    PersistStatement,
+    SelectStatement,
+)
+from trilogy.core.statements.common import MaterializedDataset
+from trilogy.core.statements.execute import (
+    ProcessedCopyStatement,
+    ProcessedQuery,
+    ProcessedQueryPersist,
+)
 from trilogy.hooks.base_hook import BaseHook
 from trilogy.utility import unique
 
@@ -50,14 +58,19 @@ def base_join_to_join(
             alias=base_join.alias,
         )
 
-    def get_datasource_cte(datasource: Datasource | QueryDatasource) -> CTE:
+    def get_datasource_cte(datasource: BuildDatasource | QueryDatasource) -> CTE:
+        eligible = set()
         for cte in ctes:
             if cte.source.identifier == datasource.identifier:
                 return cte
+            eligible.add(cte.source.identifier)
         for cte in ctes:
             if cte.source.datasources[0].identifier == datasource.identifier:
                 return cte
-        raise ValueError(f"Could not find CTE for datasource {datasource.identifier}")
+            eligible.add(cte.source.datasources[0].identifier)
+        raise ValueError(
+            f"Could not find CTE for datasource {datasource.identifier}; have {eligible}"
+        )
 
     if base_join.left_datasource is not None:
         left_cte = get_datasource_cte(base_join.left_datasource)
@@ -114,7 +127,7 @@ def generate_source_map(
             and isinstance(list(qdv)[0], UnnestJoin)
         ):
             source_map[qdk] = []
-        basic = [x for x in qdv if isinstance(x, Datasource)]
+        basic = [x for x in qdv if isinstance(x, BuildDatasource)]
         for base in basic:
             source_map[qdk].append(base.safe_identifier)
 
@@ -163,8 +176,8 @@ def generate_source_map(
     }, existence_source_map
 
 
-def datasource_to_query_datasource(datasource: Datasource) -> QueryDatasource:
-    sub_select: Dict[str, Set[Union[Datasource, QueryDatasource, UnnestJoin]]] = {
+def datasource_to_query_datasource(datasource: BuildDatasource) -> QueryDatasource:
+    sub_select: Dict[str, Set[Union[BuildDatasource, QueryDatasource, UnnestJoin]]] = {
         **{c.address: {datasource} for c in datasource.concepts},
     }
     concepts = [c for c in datasource.concepts]
@@ -206,7 +219,7 @@ def resolve_cte_base_name_and_alias_v2(
     raw_joins: List[Join | InstantiatedUnnestJoin],
 ) -> Tuple[str | None, str | None]:
     if (
-        isinstance(source.datasources[0], Datasource)
+        isinstance(source.datasources[0], BuildDatasource)
         and not source.datasources[0].name == CONSTANT_DATASET
     ):
         ds = source.datasources[0]
@@ -265,6 +278,7 @@ def datasource_to_cte(
                 for c in query_datasource.output_concepts
             ],
             grain=direct_parents[0].grain,
+            order_by=query_datasource.ordering,
         )
         return final
 
@@ -333,6 +347,7 @@ def datasource_to_cte(
         hidden_concepts=query_datasource.hidden_concepts,
         base_name_override=base_name,
         base_alias_override=base_alias,
+        order_by=query_datasource.ordering,
     )
     if cte.grain != query_datasource.grain:
         raise ValueError("Grain was corrupted in CTE generation")
@@ -351,26 +366,30 @@ def datasource_to_cte(
 
 def get_query_node(
     environment: Environment,
-    statement: SelectStatement | MultiSelectStatement,
+    statement: SelectLineage | MultiSelectLineage,
     history: History | None = None,
 ) -> StrategyNode:
-    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 grain {statement.grain}"
-    )
     if not statement.output_components:
         raise ValueError(f"Statement has no output components {statement}")
 
-    search_concepts: list[Concept] = statement.output_components
+    history = history or History(base_environment=environment)
+    build_statement = Factory(environment=environment).build(statement)
+    # build_statement = statement
+    build_environment = environment.materialize_for_select(
+        build_statement.local_concepts
+    )
+    graph = generate_graph(build_environment)
+    logger.info(
+        f"{LOGGER_PREFIX} getting source datasource for outputs {statement.output_components} grain {build_statement.grain}"
+    )
+
+    search_concepts: list[BuildConcept] = build_statement.output_components
 
     ods: StrategyNode = source_query_concepts(
-        search_concepts,
-        environment=environment,
+        output_concepts=search_concepts,
+        environment=build_environment,
         g=graph,
-        conditions=(statement.where_clause if statement.where_clause else None),
+        conditions=build_statement.where_clause,
         history=history,
     )
     if not ods:
@@ -378,22 +397,26 @@ def get_query_node(
             f"Could not find source query concepts for {[x.address for x in search_concepts]}"
         )
     ds: StrategyNode = ods
-    if statement.having_clause:
-        final = statement.having_clause.conditional
+    if build_statement.having_clause:
+        final = build_statement.having_clause.conditional
         if ds.conditions:
-            final = Conditional(
+            final = BuildConditional(
                 left=ds.conditions,
-                right=statement.having_clause.conditional,
+                right=build_statement.having_clause.conditional,
                 operator=BooleanOperator.AND,
             )
         ds = SelectNode(
-            output_concepts=statement.output_components,
+            output_concepts=build_statement.output_components,
             input_concepts=ds.output_concepts,
             parents=[ds],
             environment=ds.environment,
             partial_concepts=ds.partial_concepts,
             conditions=final,
         )
+    ds.hidden_concepts = build_statement.hidden_components
+    ds.ordering = build_statement.order_by
+    # TODO: avoid this
+    ds.rebuild_cache()
     return ds
 
 
@@ -402,8 +425,7 @@ def get_query_datasources(
     statement: SelectStatement | MultiSelectStatement,
     hooks: Optional[List[BaseHook]] = None,
 ) -> QueryDatasource:
-
-    ds = get_query_node(environment, statement)
+    ds = get_query_node(environment, statement.as_lineage(environment))
     final_qds = ds.resolve()
     if hooks:
         for hook in hooks:
@@ -498,22 +520,18 @@ def process_query(
     for cte in raw_ctes:
         cte.parent_ctes = [seen[x.name] for x in cte.parent_ctes]
     deduped_ctes: List[CTE | UnionCTE] = list(seen.values())
-    root_cte.order_by = statement.order_by
+
     root_cte.limit = statement.limit
     root_cte.hidden_concepts = statement.hidden_components
 
     final_ctes = optimize_ctes(deduped_ctes, root_cte, statement)
+    mapping = {x.address: x for x in cte.output_columns}
     return ProcessedQuery(
-        order_by=statement.order_by,
-        grain=statement.grain,
+        order_by=root_cte.order_by,
         limit=statement.limit,
-        where_clause=statement.where_clause,
-        having_clause=statement.having_clause,
-        output_columns=statement.output_components,
+        output_columns=[mapping[x.address] for x in statement.output_components],
         ctes=final_ctes,
         base=root_cte,
-        # we no longer do any joins at final level, this should always happen in parent CTEs
-        joins=[],
         hidden_columns=set([x for x in statement.hidden_components]),
         local_concepts=statement.local_concepts,
     )