pytrilogy 0.0.1.103__py3-none-any.whl → 0.0.1.105__py3-none-any.whl

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,35 +26,64 @@ 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,
         history=history,
     )
+    if not enrich_node:
+        return filter_node
     x = MergeNode(
         input_concepts=[concept, immediate_parent] + local_optional,
         output_concepts=[
@@ -73,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/node_generators/select_node.py

@@ -2,7 +2,13 @@ from itertools import combinations
 from typing import List, Optional
 
 from trilogy.core.enums import PurposeLineage
-from trilogy.core.models import Concept, Environment, Grain, LooseConceptList
+from trilogy.core.models import (
+    Concept,
+    Environment,
+    Grain,
+    LooseConceptList,
+    Datasource,
+)
 from trilogy.core.processing.nodes import (
     StrategyNode,
     SelectNode,
@@ -15,10 +21,211 @@ import networkx as nx
 from trilogy.core.graph_models import concept_to_node, datasource_to_node
 from trilogy.constants import logger
 from trilogy.core.processing.utility import padding
+from dataclasses import dataclass
 
 LOGGER_PREFIX = "[GEN_SELECT_NODE]"
 
 
+@dataclass
+class DatasourceMatch:
+    key: str
+    datasource: Datasource
+    matched: LooseConceptList
+    partial: LooseConceptList
+
+    def __repr__(self):
+        return f"DatasourceMatch({self.key}, {self.datasource.identifier}, {str(self.matched)}, {str(self.partial)})"
+
+
+def dm_to_strategy_node(
+    dm: DatasourceMatch,
+    target_grain: Grain,
+    environment: Environment,
+    g: nx.DiGraph,
+    depth: int,
+    accept_partial: bool = False,
+) -> StrategyNode:
+    datasource = dm.datasource
+    if target_grain and target_grain.issubset(datasource.grain):
+        if all([x in dm.matched for x in target_grain.components]):
+            force_group = False
+        # if we are not returning the grain
+        # we have to group
+        else:
+            logger.info(
+                f"{padding(depth)}{LOGGER_PREFIX} not all grain components are in output {str(dm.matched)}, group to actual grain"
+            )
+            force_group = True
+    elif all([x in dm.matched for x in datasource.grain.components]):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} query output includes all grain components, no reason to group further"
+        )
+        force_group = False
+    else:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} target grain is not subset of datasource grain {datasource.grain}, required to group"
+        )
+        force_group = True
+    bcandidate: StrategyNode = SelectNode(
+        input_concepts=[c.concept for c in datasource.columns],
+        output_concepts=dm.matched.concepts,
+        environment=environment,
+        g=g,
+        parents=[],
+        depth=depth,
+        partial_concepts=dm.partial.concepts,
+        accept_partial=accept_partial,
+        datasource=datasource,
+        grain=Grain(components=dm.matched.concepts),
+    )
+    # we need to nest the group node one further
+    if force_group is True:
+        candidate: StrategyNode = GroupNode(
+            output_concepts=dm.matched.concepts,
+            input_concepts=dm.matched.concepts,
+            environment=environment,
+            g=g,
+            parents=[bcandidate],
+            depth=depth,
+            partial_concepts=bcandidate.partial_concepts,
+        )
+    else:
+        candidate = bcandidate
+    return candidate
+
+
+def gen_select_nodes_from_tables_v2(
+    mandatory_concept: Concept,
+    all_concepts: List[Concept],
+    g: nx.DiGraph,
+    environment: Environment,
+    depth: int,
+    target_grain: Grain,
+    accept_partial: bool = False,
+) -> tuple[bool, list[Concept], list[StrategyNode]]:
+    # if we have only constants
+    # we don't need a table
+    # so verify nothing, select node will render
+    all_lcl = LooseConceptList(concepts=all_concepts)
+    if all([c.derivation == PurposeLineage.CONSTANT for c in all_lcl]):
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} All concepts {[x.address for x in all_lcl]} are constants, returning constant node"
+        )
+        return (
+            True,
+            all_lcl.concepts,
+            [
+                ConstantNode(
+                    output_concepts=all_lcl.concepts,
+                    input_concepts=[],
+                    environment=environment,
+                    g=g,
+                    parents=[],
+                    depth=depth,
+                    # no partial for constants
+                    partial_concepts=[],
+                    force_group=False,
+                )
+            ],
+        )
+    # otherwise, we need to look for a table
+    nodes_to_find = [concept_to_node(x.with_default_grain()) for x in all_lcl.concepts]
+    matches: dict[str, DatasourceMatch] = {}
+    for k, datasource in environment.datasources.items():
+        matched = []
+        matched_paths = []
+        for idx, req_concept in enumerate(nodes_to_find):
+            try:
+                path = nx.shortest_path(
+                    g,
+                    source=datasource_to_node(datasource),
+                    target=req_concept,
+                )
+                ds_valid = (
+                    sum(
+                        [
+                            1 if g.nodes[node]["type"] == "datasource" else 0
+                            for node in path
+                        ]
+                    )
+                    == 1
+                )
+                address_valid = (
+                    sum(
+                        [
+                            (
+                                1
+                                if g.nodes[node]["type"] == "concept"
+                                and g.nodes[node]["concept"].address
+                                != all_lcl.concepts[idx].address
+                                else 0
+                            )
+                            for node in path
+                        ]
+                    )
+                    == 0
+                )
+                if ds_valid and address_valid:
+                    matched_paths.append(path)
+                    matched.append(all_lcl.concepts[idx])
+            except nx.NodeNotFound:
+                continue
+            except nx.exception.NetworkXNoPath:
+                continue
+        dm = DatasourceMatch(
+            key=k,
+            datasource=datasource,
+            matched=LooseConceptList(concepts=matched),
+            partial=LooseConceptList(
+                concepts=[
+                    c.concept
+                    for c in datasource.columns
+                    if not c.is_complete and c.concept.address in all_lcl
+                ]
+            ),
+        )
+        if not matched:
+            continue
+        if mandatory_concept.address not in dm.matched:
+            continue
+        if not accept_partial and dm.partial.addresses:
+            continue
+        matches[k] = dm
+    found: set[str] = set()
+    all_found = False
+    all_checked = False
+    to_return: list[StrategyNode] = []
+    if not matches:
+        return False, [], []
+    while not all_found and not all_checked:
+        final_key: str = max(
+            matches,
+            key=lambda x: len(
+                [m for m in matches[x].matched.addresses if m not in found]
+            )
+            - 0.1 * len(matches[x].partial.addresses),
+        )
+        final: DatasourceMatch = matches[final_key]
+        candidate = dm_to_strategy_node(
+            final,
+            target_grain=Grain(
+                components=[
+                    x for x in target_grain.components if x.address in final.matched
+                ]
+            ),
+            environment=environment,
+            g=g,
+            depth=depth,
+            accept_partial=accept_partial,
+        )
+        to_return.append(candidate)
+        del matches[final_key]
+        found = found.union(final.matched.addresses)
+        all_found = all_lcl.addresses.issubset(found)
+        all_checked = len(matches) == 0
+    return all_found, [x for x in all_concepts if x.address in found], to_return
+
+
 def gen_select_node_from_table(
     target_concept: Concept,
     all_concepts: List[Concept],
@@ -166,58 +373,15 @@ def gen_select_node_from_table(
     return candidates[final]
 
 
-def gen_select_node(
-    concept: Concept,
+def gen_select_nodes_from_tables(
     local_optional: List[Concept],
-    environment: Environment,
-    g,
     depth: int,
-    accept_partial: bool = False,
-    fail_if_not_found: bool = True,
-    accept_partial_optional: bool = True,
-    target_grain: Grain | None = None,
-) -> StrategyNode | None:
-    all_concepts = [concept] + local_optional
-    all_lcl = LooseConceptList(concepts=all_concepts)
-    materialized_lcl = LooseConceptList(
-        concepts=[
-            x
-            for x in all_concepts
-            if x.address in [z.address for z in environment.materialized_concepts]
-            or x.derivation == PurposeLineage.CONSTANT
-        ]
-    )
-    if not target_grain:
-        target_grain = Grain()
-        for ac in all_concepts:
-            target_grain += ac.grain
-    if materialized_lcl != all_lcl:
-        logger.info(
-            f"{padding(depth)}{LOGGER_PREFIX} Skipping select node generation for {concept.address} "
-            f" as it + optional (looking for all {all_lcl}) includes non-materialized concepts {all_lcl.difference(materialized_lcl)} vs materialized: {materialized_lcl}"
-        )
-        if fail_if_not_found:
-            raise NoDatasourceException(f"No datasource exists for {concept}")
-        return None
-
-    ds: StrategyNode | None = None
-
-    # attempt to select all concepts from table
-    ds = gen_select_node_from_table(
-        concept,
-        [concept] + local_optional,
-        g=g,
-        environment=environment,
-        depth=depth,
-        accept_partial=accept_partial,
-        target_grain=target_grain,
-    )
-    if ds:
-        logger.info(
-            f"{padding(depth)}{LOGGER_PREFIX} Found select node with all target concepts, force group is {ds.force_group}, target grain {target_grain}"
-        )
-        return ds
-    # if we cannot find a match
+    concept: Concept,
+    environment: Environment,
+    g: nx.DiGraph,
+    accept_partial: bool,
+    all_concepts: List[Concept],
+) -> tuple[bool, list[Concept], list[StrategyNode]]:
     parents: List[StrategyNode] = []
     found: List[Concept] = []
     logger.info(
@@ -238,7 +402,7 @@ def gen_select_node(
         )
         if not ds:
             unreachable.append(opt_con.address)
-    # actual search
+    all_found = False
     for x in reversed(range(1, len(local_optional) + 1)):
         if all_found:
             break
@@ -275,6 +439,64 @@ def gen_select_node(
                         f"{padding(depth)}{LOGGER_PREFIX} found all optional {[c.address for c in local_optional]}"
                     )
                     all_found = True
+    return all_found, found, parents
+
+
+def gen_select_node(
+    concept: Concept,
+    local_optional: List[Concept],
+    environment: Environment,
+    g,
+    depth: int,
+    accept_partial: bool = False,
+    fail_if_not_found: bool = True,
+    accept_partial_optional: bool = True,
+    target_grain: Grain | None = None,
+) -> StrategyNode | None:
+    all_concepts = [concept] + local_optional
+    all_lcl = LooseConceptList(concepts=all_concepts)
+    materialized_lcl = LooseConceptList(
+        concepts=[
+            x
+            for x in all_concepts
+            if x.address in [z.address for z in environment.materialized_concepts]
+            or x.derivation == PurposeLineage.CONSTANT
+        ]
+    )
+    if not target_grain:
+        target_grain = Grain()
+        for ac in all_concepts:
+            target_grain += ac.grain
+    if materialized_lcl != all_lcl:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} Skipping select node generation for {concept.address} "
+            f" as it + optional (looking for all {all_lcl}) includes non-materialized concepts {all_lcl.difference(materialized_lcl)} vs materialized: {materialized_lcl}"
+        )
+        if fail_if_not_found:
+            raise NoDatasourceException(f"No datasource exists for {concept}")
+        return None
+
+    ds: StrategyNode | None = None
+
+    # attempt to select all concepts from table
+    ds = gen_select_node_from_table(
+        concept,
+        [concept] + local_optional,
+        g=g,
+        environment=environment,
+        depth=depth,
+        accept_partial=accept_partial,
+        target_grain=target_grain,
+    )
+    if ds:
+        logger.info(
+            f"{padding(depth)}{LOGGER_PREFIX} Found select node with all target concepts, force group is {ds.force_group}, target grain {target_grain}"
+        )
+        return ds
+    # if we cannot find a match
+    all_found, found, parents = gen_select_nodes_from_tables_v2(
+        concept, all_concepts, g, environment, depth, target_grain, accept_partial
+    )
     if parents and (all_found or accept_partial_optional):
         if all_found:
             logger.info(
@@ -282,7 +504,7 @@ def gen_select_node(
             )
         else:
             logger.info(
-                f"{padding(depth)}{LOGGER_PREFIX} found some optional {[c.address for c in found]}, returning"
+                f"{padding(depth)}{LOGGER_PREFIX} found some optional, returning"
             )
         all_partial = [
             c

trilogy/core/processing/nodes/__init__.py

@@ -6,11 +6,12 @@ from .window_node import WindowNode
 from .base_node import StrategyNode, NodeJoin
 from .unnest_node import UnnestNode
 from pydantic import BaseModel, Field, ConfigDict
-from trilogy.core.models import Concept
+from trilogy.core.models import Concept, Environment
 
 
 class History(BaseModel):
     history: dict[str, StrategyNode | None] = Field(default_factory=dict)
+    select_history: dict[str, StrategyNode | None] = Field(default_factory=dict)
     started: set[str] = Field(default_factory=set)
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -60,6 +61,58 @@ class History(BaseModel):
             in self.started
         )
 
+    def _select_concepts_to_lookup(
+        self,
+        main: Concept,
+        search: list[Concept],
+        accept_partial: bool,
+        fail_if_not_found: bool,
+        accept_partial_optional: bool,
+    ) -> str:
+        return (
+            str(main.address)
+            + "|"
+            + "-".join([c.address for c in search])
+            + str(accept_partial)
+            + str(fail_if_not_found)
+            + str(accept_partial_optional)
+        )
+
+    def gen_select_node(
+        self,
+        concept: Concept,
+        local_optional: list[Concept],
+        environment: Environment,
+        g,
+        depth: int,
+        fail_if_not_found: bool = False,
+        accept_partial: bool = False,
+        accept_partial_optional: bool = False,
+    ) -> StrategyNode | None:
+        from trilogy.core.processing.node_generators.select_node import gen_select_node
+
+        fingerprint = self._select_concepts_to_lookup(
+            concept,
+            local_optional,
+            accept_partial,
+            fail_if_not_found,
+            accept_partial_optional,
+        )
+        if fingerprint in self.select_history:
+            return self.select_history[fingerprint]
+        gen = gen_select_node(
+            concept,
+            local_optional,
+            environment,
+            g,
+            depth + 1,
+            fail_if_not_found=fail_if_not_found,
+            accept_partial=accept_partial,
+            accept_partial_optional=accept_partial_optional,
+        )
+        self.select_history[fingerprint] = gen
+        return gen
+
 
 __all__ = [
     "FilterNode",

trilogy/dialect/base.py

@@ -22,6 +22,7 @@ from trilogy.core.models import (
     CompiledCTE,
     Conditional,
     Comparison,
+    SubselectComparison,
     OrderItem,
     WindowItem,
     FilterItem,
@@ -273,14 +274,13 @@ class BaseDialect:
                 ]
                 rval = f"{self.WINDOW_FUNCTION_MAP[c.lineage.type](concept = self.render_concept_sql(c.lineage.content, cte=cte, alias=False), window=','.join(rendered_over_components), sort=','.join(rendered_order_components))}"  # noqa: E501
             elif isinstance(c.lineage, FilterItem):
-                rval = f"CASE WHEN {self.render_expr(c.lineage.where.conditional)} THEN {self.render_concept_sql(c.lineage.content, cte=cte, alias=False)} ELSE NULL END"
+                rval = f"CASE WHEN {self.render_expr(c.lineage.where.conditional, cte=cte)} THEN {self.render_concept_sql(c.lineage.content, cte=cte, alias=False)} ELSE NULL END"
             elif isinstance(c.lineage, RowsetItem):
                 rval = f"{self.render_concept_sql(c.lineage.content, cte=cte, alias=False)}"
             elif isinstance(c.lineage, MultiSelectStatement):
                 rval = f"{self.render_concept_sql(c.lineage.find_source(c, cte), cte=cte, alias=False)}"
             elif isinstance(c.lineage, MergeStatement):
                 rval = f"{self.render_concept_sql(c.lineage.find_source(c, cte), cte=cte, alias=False)}"
-                #  rval = f"{self.FUNCTION_MAP[FunctionType.COALESCE](*[self.render_concept_sql(parent, cte=cte, alias=False) for parent in c.lineage.find_sources(c, cte)])}"
             elif isinstance(c.lineage, AggregateWrapper):
                 args = [
                     self.render_expr(v, cte)  # , alias=False)
@@ -330,6 +330,7 @@ class BaseDialect:
             Function,
             Conditional,
             Comparison,
+            SubselectComparison,
             Concept,
             str,
             int,
@@ -358,7 +359,15 @@ class BaseDialect:
         # if isinstance(e, Concept):
         #     cte = cte or cte_map.get(e.address, None)
 
-        if isinstance(e, Comparison):
+        if isinstance(e, SubselectComparison):
+            assert cte, "Subselects must be rendered with a CTE in context"
+            if isinstance(e.right, Concept):
+                return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} (select {self.render_expr(e.right, cte=cte, cte_map=cte_map)} from {cte.source_map[e.right.address][0]})"
+            else:
+                raise NotImplementedError(
+                    f"Subselects must be a concept, got {e.right}"
+                )
+        elif isinstance(e, Comparison):
             return f"{self.render_expr(e.left, cte=cte, cte_map=cte_map)} {e.operator.value} {self.render_expr(e.right, cte=cte, cte_map=cte_map)}"
         elif isinstance(e, Conditional):
             # conditions need to be nested in parentheses

trilogy/parsing/common.py

@@ -174,3 +174,33 @@ def agg_wrapper_to_concept(
         keys=tuple(parent.by) if parent.by else keys,
     )
     return out
+
+
+def arbitrary_to_concept(
+    parent: (
+        AggregateWrapper
+        | WindowItem
+        | FilterItem
+        | Function
+        | ListWrapper
+        | int
+        | float
+        | str
+    ),
+    namespace: str,
+    name: str,
+    metadata: Metadata | None = None,
+    purpose: Purpose | None = None,
+) -> Concept:
+    if isinstance(parent, AggregateWrapper):
+        return agg_wrapper_to_concept(parent, namespace, name, metadata, purpose)
+    elif isinstance(parent, WindowItem):
+        return window_item_to_concept(parent, name, namespace, purpose, metadata)
+    elif isinstance(parent, FilterItem):
+        return filter_item_to_concept(parent, name, namespace, purpose, metadata)
+    elif isinstance(parent, Function):
+        return function_to_concept(parent, name, namespace)
+    elif isinstance(parent, ListWrapper):
+        return constant_to_concept(parent, name, namespace, purpose, metadata)
+    else:
+        return constant_to_concept(parent, name, namespace, purpose, metadata)