pytrilogy 0.0.1.104__py3-none-any.whl → 0.0.1.106__py3-none-any.whl

trilogy/core/optimization.py

@@ -0,0 +1,141 @@
+from trilogy.core.models import (
+    CTE,
+    SelectStatement,
+    PersistStatement,
+    Datasource,
+    MultiSelectStatement,
+)
+from trilogy.core.enums import PurposeLineage
+from trilogy.constants import logger
+from abc import ABC
+
+
+class OptimizationRule(ABC):
+
+    def optimize(self, cte: CTE) -> bool:
+        raise NotImplementedError
+
+    def log(self, message: str):
+        logger.info(f"[Optimization][{self.__class__.__name__}] {message}")
+
+
+class InlineDatasource(OptimizationRule):
+
+    def optimize(self, cte: CTE) -> bool:
+        if not cte.parent_ctes:
+            return False
+
+        optimized = False
+        self.log(
+            f"Checking {cte.name} for consolidating inline tables with {len(cte.parent_ctes)} parents"
+        )
+        to_inline: list[CTE] = []
+        for parent_cte in cte.parent_ctes:
+            if not parent_cte.is_root_datasource:
+                self.log(f"parent {parent_cte.name} is not root")
+                continue
+            if parent_cte.parent_ctes:
+                self.log(f"parent {parent_cte.name} has parents")
+                continue
+            raw_root = parent_cte.source.datasources[0]
+            if not isinstance(raw_root, Datasource):
+                self.log(f"parent {parent_cte.name} is not datasource")
+                continue
+            root: Datasource = raw_root
+            if not root.can_be_inlined:
+                self.log(f"parent {parent_cte.name} datasource is not inlineable")
+                continue
+            root_outputs = {x.address for x in root.output_concepts}
+            cte_outputs = {x.address for x in parent_cte.output_columns}
+            if not cte_outputs.issubset(root_outputs):
+                self.log(f"Not all {parent_cte.name} outputs are found on datasource")
+                continue
+
+            to_inline.append(parent_cte)
+
+        for replaceable in to_inline:
+            self.log(f"Inlining parent {replaceable.name}")
+            cte.inline_parent_datasource(replaceable)
+
+        return optimized
+
+
+REGISTERED_RULES: list[OptimizationRule] = [InlineDatasource()]
+
+
+def filter_irrelevant_ctes(input: list[CTE], root_cte: CTE):
+    relevant_ctes = set()
+
+    def recurse(cte: CTE):
+        relevant_ctes.add(cte.name)
+        for cte in cte.parent_ctes:
+            recurse(cte)
+
+    recurse(root_cte)
+    return [cte for cte in input if cte.name in relevant_ctes]
+
+
+def is_direct_return_eligible(
+    cte: CTE, select: SelectStatement | PersistStatement | MultiSelectStatement
+) -> bool:
+    if isinstance(select, (PersistStatement, MultiSelectStatement)):
+        return False
+    derived_concepts = [
+        c for c in cte.source.output_concepts if c not in cte.source.input_concepts
+    ]
+    eligible = True
+    conditions = (
+        set(x.address for x in select.where_clause.concept_arguments)
+        if select.where_clause
+        else set()
+    )
+    if conditions and select.limit:
+        return False
+    for x in derived_concepts:
+        if x.derivation == PurposeLineage.WINDOW:
+            return False
+        if x.derivation == PurposeLineage.AGGREGATE:
+            if x.address in conditions:
+                return False
+    logger.info(
+        f"Upleveling output select to final CTE with derived_concepts {[x.address for x in derived_concepts]}"
+    )
+    return eligible
+
+
+def sort_select_output(cte: CTE, query: SelectStatement | MultiSelectStatement):
+    hidden_addresses = [c.address for c in query.hidden_components]
+    output_addresses = [
+        c.address for c in query.output_components if c.address not in hidden_addresses
+    ]
+
+    mapping = {x.address: x for x in cte.output_columns}
+
+    new_output = []
+    for x in output_addresses:
+        new_output.append(mapping[x])
+    cte.output_columns = new_output
+
+
+def optimize_ctes(
+    input: list[CTE], root_cte: CTE, select: SelectStatement | MultiSelectStatement
+):
+    complete = False
+
+    while not complete:
+        actions_taken = False
+        for rule in REGISTERED_RULES:
+            for cte in input:
+                actions_taken = rule.optimize(cte)
+        complete = not actions_taken
+
+    if is_direct_return_eligible(root_cte, select):
+        root_cte.order_by = select.order_by
+        root_cte.limit = select.limit
+        root_cte.condition = (
+            select.where_clause.conditional if select.where_clause else None
+        )
+        root_cte.requires_nesting = False
+        sort_select_output(cte, select)
+
+    return filter_irrelevant_ctes(input, root_cte)

trilogy/core/processing/concept_strategies_v3.py

@@ -317,6 +317,7 @@ def generate_node(
         return gen_basic_node(
             concept, local_optional, environment, g, depth + 1, source_concepts, history
         )
+
     elif concept.derivation == PurposeLineage.ROOT:
         logger.info(
             f"{depth_to_prefix(depth)}{LOGGER_PREFIX} for {concept.address}, generating select node with optional {[x.address for x in local_optional]}"

trilogy/core/processing/node_generators/common.py

@@ -45,21 +45,33 @@ def resolve_function_parent_concepts(concept: Concept) -> List[Concept]:
     return unique(concept.lineage.concept_arguments, "address")
 
 
-def resolve_filter_parent_concepts(concept: Concept) -> Tuple[Concept, List[Concept]]:
+def resolve_filter_parent_concepts(
+    concept: Concept,
+) -> Tuple[Concept, List[Concept], List[Concept]]:
     if not isinstance(concept.lineage, FilterItem):
-        raise ValueError
+        raise ValueError(
+            f"Concept {concept} lineage is not filter item, is {type(concept.lineage)}"
+        )
     direct_parent = concept.lineage.content
-    base = [direct_parent]
-    base += concept.lineage.where.concept_arguments
+    base_existence = []
+    base_rows = [direct_parent]
+    base_rows += concept.lineage.where.row_arguments
+    base_existence += concept.lineage.where.existence_arguments
     if direct_parent.grain:
-        base += direct_parent.grain.components_copy
+        base_rows += direct_parent.grain.components_copy
     if (
         isinstance(direct_parent, Concept)
         and direct_parent.purpose == Purpose.PROPERTY
         and direct_parent.keys
     ):
-        base += direct_parent.keys
-    return concept.lineage.content, unique(base, "address")
+        base_rows += direct_parent.keys
+    if concept.lineage.where.existence_arguments:
+        return (
+            concept.lineage.content,
+            unique(base_rows, "address"),
+            unique(base_existence, "address"),
+        )
+    return concept.lineage.content, unique(base_rows, "address"), []
 
 
 def gen_property_enrichment_node(

trilogy/core/processing/node_generators/filter_node.py

@@ -11,7 +11,7 @@ from trilogy.core.processing.node_generators.common import (
     resolve_filter_parent_concepts,
 )
 from trilogy.constants import logger
-from trilogy.core.processing.utility import padding
+from trilogy.core.processing.utility import padding, unique
 from trilogy.core.processing.node_generators.common import concept_to_relevant_joins
 
 LOGGER_PREFIX = "[GEN_FILTER_NODE]"
@@ -26,30 +26,57 @@ def gen_filter_node(
     source_concepts,
     history: History | None = None,
 ) -> MergeNode | FilterNode | None:
-    immediate_parent, parent_concepts = resolve_filter_parent_concepts(concept)
+    immediate_parent, parent_row_concepts, parent_existence_concepts = (
+        resolve_filter_parent_concepts(concept)
+    )
 
-    logger.info(f"{padding(depth)}{LOGGER_PREFIX} fetching filter node parents")
+    logger.info(
+        f"{padding(depth)}{LOGGER_PREFIX} fetching filter node row parents {[x.address for x in parent_row_concepts]}"
+    )
+    core_parents = []
     parent = source_concepts(
-        mandatory_list=parent_concepts,
+        mandatory_list=parent_row_concepts,
         environment=environment,
         g=g,
         depth=depth + 1,
         history=history,
     )
+
     if not parent:
         return None
+    core_parents.append(parent)
+    if parent_existence_concepts:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} fetching filter node existence parents {[x.address for x in parent_existence_concepts]}"
+        )
+        parent_existence = source_concepts(
+            mandatory_list=parent_existence_concepts,
+            environment=environment,
+            g=g,
+            depth=depth + 1,
+            history=history,
+        )
+        if not parent_existence:
+            return None
+        core_parents.append(parent_existence)
+
     filter_node = FilterNode(
-        input_concepts=[immediate_parent] + parent_concepts,
-        output_concepts=[concept, immediate_parent] + parent_concepts,
+        input_concepts=unique(
+            [immediate_parent] + parent_row_concepts + parent_existence_concepts,
+            "address",
+        ),
+        output_concepts=[concept, immediate_parent] + parent_row_concepts,
         environment=environment,
         g=g,
-        parents=[parent],
+        parents=core_parents,
     )
-    if not local_optional:
+    if not local_optional or all(
+        [x.address in [y.address for y in parent_row_concepts] for x in local_optional]
+    ):
         return filter_node
     enrich_node = source_concepts(  # this fetches the parent + join keys
         # to then connect to the rest of the query
-        mandatory_list=[immediate_parent] + parent_concepts + local_optional,
+        mandatory_list=[immediate_parent] + parent_row_concepts + local_optional,
         environment=environment,
         g=g,
         depth=depth + 1,
@@ -75,7 +102,7 @@ def gen_filter_node(
                 left_node=enrich_node,
                 right_node=filter_node,
                 concepts=concept_to_relevant_joins(
-                    [immediate_parent] + parent_concepts
+                    [immediate_parent] + parent_row_concepts
                 ),
                 join_type=JoinType.LEFT_OUTER,
                 filter_to_mutual=False,

trilogy/core/processing/node_generators/merge_node.py

@@ -87,8 +87,18 @@ def gen_merge_node(
 ) -> Optional[MergeNode]:
     join_candidates: List[PathInfo] = []
     # anchor on datasources
+    final_all_concepts = []
+    # implicit_upstream = {}
+    for x in all_concepts:
+        # if x.derivation in (PurposeLineage.AGGREGATE, PurposeLineage.BASIC):
+        #     final_all_concepts +=resolve_function_parent_concepts(x)
+        # elif x.derivation == PurposeLineage.FILTER:
+        #     final_all_concepts +=resolve_filter_parent_concepts(x)
+        # else:
+        #     final_all_concepts.append(x)
+        final_all_concepts.append(x)
     for datasource in environment.datasources.values():
-        path = identify_ds_join_paths(all_concepts, g, datasource, accept_partial)
+        path = identify_ds_join_paths(final_all_concepts, g, datasource, accept_partial)
         if path and path.reduced_concepts:
             join_candidates.append(path)
     join_candidates.sort(key=lambda x: sum([len(v) for v in x.paths.values()]))

trilogy/core/processing/nodes/base_node.py

@@ -45,7 +45,7 @@ def concept_list_to_grain(
 
 
 def resolve_concept_map(
-    inputs: List[QueryDatasource],
+    inputs: List[QueryDatasource | Datasource],
     targets: List[Concept],
     inherited_inputs: List[Concept],
     full_joins: List[Concept] | None = None,
@@ -156,7 +156,9 @@ class StrategyNode:
         return f"{self.__class__.__name__}<{contents}>"
 
     def _resolve(self) -> QueryDatasource:
-        parent_sources = [p.resolve() for p in self.parents]
+        parent_sources: List[QueryDatasource | Datasource] = [
+            p.resolve() for p in self.parents
+        ]
 
         # if conditional:
         #     for condition in conditions[1:]:

trilogy/core/processing/nodes/group_node.py

@@ -4,6 +4,7 @@ from trilogy.constants import logger
 from trilogy.core.models import (
     Grain,
     QueryDatasource,
+    Datasource,
     SourceType,
     Concept,
     Environment,
@@ -45,7 +46,9 @@ class GroupNode(StrategyNode):
         )
 
     def _resolve(self) -> QueryDatasource:
-        parent_sources: list[QueryDatasource] = [p.resolve() for p in self.parents]
+        parent_sources: List[QueryDatasource | Datasource] = [
+            p.resolve() for p in self.parents
+        ]
 
         grain = concept_list_to_grain(self.output_concepts, [])
         comp_grain = Grain()
@@ -66,7 +69,7 @@ class GroupNode(StrategyNode):
                 len(parent_sources) == 1
                 and LooseConceptList(concepts=parent_sources[0].output_concepts)
                 == self.output_lcl
-            ):
+            ) and isinstance(parent_sources[0], QueryDatasource):
                 logger.info(
                     f"{self.logging_prefix}{LOGGER_PREFIX} No group by required, returning parent node"
                 )

trilogy/core/processing/nodes/merge_node.py

@@ -7,6 +7,7 @@ from trilogy.core.models import (
     Grain,
     JoinType,
     QueryDatasource,
+    Datasource,
     SourceType,
     Concept,
     UnnestJoin,
@@ -24,8 +25,8 @@ LOGGER_PREFIX = "[CONCEPT DETAIL - MERGE NODE]"
 
 
 def deduplicate_nodes(
-    merged: dict[str, QueryDatasource], logging_prefix: str
-) -> tuple[bool, dict[str, QueryDatasource], set[str]]:
+    merged: dict[str, QueryDatasource | Datasource], logging_prefix: str
+) -> tuple[bool, dict[str, QueryDatasource | Datasource], set[str]]:
     duplicates = False
     removed: set[str] = set()
     set_map: dict[str, set[str]] = {}
@@ -65,9 +66,9 @@ def deduplicate_nodes(
 
 def deduplicate_nodes_and_joins(
     joins: List[NodeJoin] | None,
-    merged: dict[str, QueryDatasource],
+    merged: dict[str, QueryDatasource | Datasource],
     logging_prefix: str,
-) -> Tuple[List[NodeJoin] | None, dict[str, QueryDatasource]]:
+) -> Tuple[List[NodeJoin] | None, dict[str, QueryDatasource | Datasource]]:
     # it's possible that we have more sources than we need
     duplicates = True
     while duplicates:
@@ -211,8 +212,10 @@ class MergeNode(StrategyNode):
         return joins
 
     def _resolve(self) -> QueryDatasource:
-        parent_sources = [p.resolve() for p in self.parents]
-        merged: dict[str, QueryDatasource] = {}
+        parent_sources: List[QueryDatasource | Datasource] = [
+            p.resolve() for p in self.parents
+        ]
+        merged: dict[str, QueryDatasource | Datasource] = {}
         final_joins = self.node_joins
         for source in parent_sources:
             if source.full_name in merged:
@@ -228,14 +231,15 @@ class MergeNode(StrategyNode):
             final_joins, merged, self.logging_prefix
         )
         # early exit if we can just return the parent
-        final_datasets: List[QueryDatasource] = list(merged.values())
+        final_datasets: List[QueryDatasource | Datasource] = list(merged.values())
 
         if len(merged.keys()) == 1:
-            final: QueryDatasource = list(merged.values())[0]
+            final: QueryDatasource | Datasource = list(merged.values())[0]
             if (
                 set([c.address for c in final.output_concepts])
                 == set([c.address for c in self.output_concepts])
                 and not self.conditions
+                and isinstance(final, QueryDatasource)
             ):
                 logger.info(
                     f"{self.logging_prefix}{LOGGER_PREFIX} Merge node has only one parent with the same"
@@ -255,6 +259,7 @@ class MergeNode(StrategyNode):
             if (
                 all([c.address in output_set for c in self.all_concepts])
                 and not self.conditions
+                and isinstance(dataset, QueryDatasource)
             ):
                 logger.info(
                     f"{self.logging_prefix}{LOGGER_PREFIX} Merge node not required as parent node {dataset.source_type}"

trilogy/core/query_processor.py

@@ -29,6 +29,7 @@ from trilogy.hooks.base_hook import BaseHook
 from trilogy.constants import logger
 from random import shuffle
 from trilogy.core.ergonomics import CTE_NAMES
+from trilogy.core.optimization import optimize_ctes
 from math import ceil
 
 LOGGER_PREFIX = "[QUERY BUILD]"
@@ -186,7 +187,7 @@ def datasource_to_ctes(
             source_map = {k: "" for k in query_datasource.source_map}
         else:
             source_map = {
-                k: "" if not v else source.full_name
+                k: "" if not v else source.identifier
                 for k, v in query_datasource.source_map.items()
             }
     human_id = generate_cte_name(query_datasource.full_name, name_map)
@@ -315,7 +316,9 @@ 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]
-    final_ctes: List[CTE] = list(seen.values())
+    deduped_ctes: List[CTE] = list(seen.values())
+
+    final_ctes = optimize_ctes(deduped_ctes, root_cte, statement)
 
     return ProcessedQuery(
         order_by=statement.order_by,