pytrilogy 0.0.1.118__py3-none-any.whl → 0.0.2.2__py3-none-any.whl

trilogy/core/processing/utility.py

@@ -7,6 +7,7 @@ from trilogy.core.models import (
     Concept,
     QueryDatasource,
     LooseConceptList,
+    Environment,
 )
 
 from trilogy.core.enums import Purpose, Granularity
@@ -123,9 +124,23 @@ def resolve_join_order(joins: List[BaseJoin]) -> List[BaseJoin]:
     return final_joins
 
 
+def add_node_join_concept(graph, concept, datasource, concepts):
+    # we don't need to join on a concept if all of the keys exist in the grain
+    # if concept.keys and all([x in grain for x in concept.keys]):
+    #     continue
+    concepts.append(concept)
+
+    graph.add_node(concept.address, type=NodeType.CONCEPT)
+    graph.add_edge(datasource.identifier, concept.address)
+    for k, v in concept.pseudonyms.items():
+        if v.address != concept.address:
+            add_node_join_concept(graph, v, datasource, concepts)
+
+
 def get_node_joins(
     datasources: List[QueryDatasource],
     grain: List[Concept],
+    environment: Environment,
     # concepts:List[Concept],
 ) -> List[BaseJoin]:
     graph = nx.Graph()
@@ -133,12 +148,14 @@ def get_node_joins(
     for datasource in datasources:
         graph.add_node(datasource.identifier, type=NodeType.NODE)
         for concept in datasource.output_concepts:
+            add_node_join_concept(graph, concept, datasource, concepts)
             # we don't need to join on a concept if all of the keys exist in the grain
             # if concept.keys and all([x in grain for x in concept.keys]):
             #     continue
-            concepts.append(concept)
-            graph.add_node(concept.address, type=NodeType.CONCEPT)
-            graph.add_edge(datasource.identifier, concept.address)
+            # concepts.append(concept)
+
+            # graph.add_node(concept.address, type=NodeType.CONCEPT)
+            # graph.add_edge(datasource.identifier, concept.address)
 
     # add edges for every constant to every datasource
     for datasource in datasources:
@@ -149,15 +166,55 @@ def get_node_joins(
                         graph.add_edge(node, concept.address)
 
     joins: defaultdict[str, set] = defaultdict(set)
-    identifier_map = {x.identifier: x for x in datasources}
+    identifier_map: dict[str, Datasource | QueryDatasource] = {
+        x.identifier: x for x in datasources
+    }
+
+    grain_pseudonyms: set[str] = set()
+    for g in grain:
+        env_lookup = environment.concepts[g.address]
+        # if we're looking up a pseudonym, we would have gotten the remapped value
+        # so double check we got what we were looking for
+        if env_lookup.address == g.address:
+            grain_pseudonyms.update(env_lookup.pseudonyms.keys())
 
     node_list = sorted(
         [x for x in graph.nodes if graph.nodes[x]["type"] == NodeType.NODE],
         # sort so that anything with a partial match on the target is later
         key=lambda x: len(
-            [x for x in identifier_map[x].partial_concepts if x in grain]
+            [
+                partial
+                for partial in identifier_map[x].partial_concepts
+                if partial in grain
+            ]
+            + [
+                output
+                for output in identifier_map[x].output_concepts
+                if output.address in grain_pseudonyms
+            ]
         ),
     )
+
+    node_map = {
+        x[0:20]: len(
+            [
+                partial
+                for partial in identifier_map[x].partial_concepts
+                if partial in grain
+            ]
+            + [
+                output
+                for output in identifier_map[x].output_concepts
+                if output.address in grain_pseudonyms
+            ]
+        )
+        for x in node_list
+    }
+    print("NODE MAP")
+    print(node_map)
+    print([x.address for x in grain])
+    print(grain_pseudonyms)
+
     for left in node_list:
         # the constant dataset is a special case
         # and can never be on the left of a join
@@ -203,14 +260,49 @@ def get_node_joins(
                 c for c in local_concepts if c.granularity != Granularity.SINGLE_ROW
             ]
 
-            # if concept.keys and all([x in grain for x in concept.keys]):
-            #     continue
+        relevant = concept_to_relevant_joins(local_concepts)
+        left_datasource = identifier_map[left]
+        right_datasource = identifier_map[right]
+        join_tuples = []
+        for joinc in relevant:
+            left_arg = joinc
+            right_arg = joinc
+            if joinc.address not in [
+                c.address for c in left_datasource.output_concepts
+            ]:
+                try:
+                    left_arg = [
+                        x
+                        for x in left_datasource.output_concepts
+                        if x.address in joinc.pseudonyms
+                        or joinc.address in x.pseudonyms
+                    ].pop()
+                except IndexError:
+                    raise SyntaxError(
+                        f"Could not find {joinc.address} in {left_datasource.identifier} output {[c.address for c in left_datasource.output_concepts]}"
+                    )
+            if joinc.address not in [
+                c.address for c in right_datasource.output_concepts
+            ]:
+                try:
+                    right_arg = [
+                        x
+                        for x in right_datasource.output_concepts
+                        if x.address in joinc.pseudonyms
+                        or joinc.address in x.pseudonyms
+                    ].pop()
+                except IndexError:
+                    raise SyntaxError(
+                        f"Could not find {joinc.address} in {right_datasource.identifier} output {[c.address for c in right_datasource.output_concepts]}"
+                    )
+            join_tuples.append((left_arg, right_arg))
         final_joins_pre.append(
             BaseJoin(
                 left_datasource=identifier_map[left],
                 right_datasource=identifier_map[right],
                 join_type=join_type,
-                concepts=concept_to_relevant_joins(local_concepts),
+                concepts=[],
+                concept_pairs=join_tuples,
             )
         )
     final_joins = resolve_join_order(final_joins_pre)

trilogy/core/query_processor.py

@@ -4,8 +4,11 @@ from trilogy.core.env_processor import generate_graph
 from trilogy.core.graph_models import ReferenceGraph
 from trilogy.core.constants import CONSTANT_DATASET
 from trilogy.core.processing.concept_strategies_v3 import source_query_concepts
+from trilogy.core.enums import SelectFiltering
 from trilogy.constants import CONFIG, DEFAULT_NAMESPACE
+from trilogy.core.processing.nodes import GroupNode, SelectNode, StrategyNode
 from trilogy.core.models import (
+    Concept,
     Environment,
     PersistStatement,
     SelectStatement,
@@ -24,13 +27,14 @@ from trilogy.core.models import (
 )
 
 from trilogy.utility import unique
-from collections import defaultdict
+
 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
+from collections import defaultdict
+from random import shuffle
 
 LOGGER_PREFIX = "[QUERY BUILD]"
 
@@ -79,6 +83,7 @@ def base_join_to_join(
         right_cte=right_cte,
         joinkeys=[JoinKey(concept=concept) for concept in base_join.concepts],
         jointype=base_join.join_type,
+        joinkey_pairs=base_join.concept_pairs if base_join.concept_pairs else None,
     )
 
 
@@ -107,7 +112,9 @@ def generate_source_map(
             matches = [cte for cte in all_new_ctes if cte.source.name in names]
 
             if not matches and names:
-                raise SyntaxError(query_datasource.source_map)
+                raise SyntaxError(
+                    f"Missing parent CTEs for source map; expecting {names}, have {[cte.source.name for cte in all_new_ctes]}"
+                )
             for cte in matches:
                 output_address = [
                     x.address
@@ -260,7 +267,7 @@ def datasource_to_ctes(
 
     human_id = generate_cte_name(query_datasource.full_name, name_map)
     logger.info(
-        f"Finished building source map for {human_id} with {len(parents)} parents, have {source_map}, parent had non-empty keys {[k for k, v in query_datasource.source_map.items() if v]} "
+        f"Finished building source map for {human_id} with {len(parents)} parents, have {source_map}, query_datasource had non-empty keys {[k for k, v in query_datasource.source_map.items() if v]} "
     )
     final_joins = [
         x
@@ -306,6 +313,28 @@ def datasource_to_ctes(
     return output
 
 
+def append_existence_check(
+    node: StrategyNode, environment: Environment, graph: ReferenceGraph
+):
+    # we if we have a where clause doing an existence check
+    # treat that as separate subquery
+    if (where := node.conditions) and where.existence_arguments:
+        for subselect in where.existence_arguments:
+            if not subselect:
+                continue
+            logger.info(
+                f"{LOGGER_PREFIX} fetching existance clause inputs {[str(c) for c in subselect]}"
+            )
+            eds = source_query_concepts([*subselect], environment=environment, g=graph)
+
+            final_eds = eds.resolve()
+            first_parent = node.resolve()
+            first_parent.datasources.append(final_eds)
+            for x in final_eds.output_concepts:
+                if x.address not in first_parent.existence_source_map:
+                    first_parent.existence_source_map[x.address] = {final_eds}
+
+
 def get_query_datasources(
     environment: Environment,
     statement: SelectStatement | MultiSelectStatement,
@@ -318,33 +347,57 @@ def get_query_datasources(
     )
     if not statement.output_components:
         raise ValueError(f"Statement has no output components {statement}")
-    ds = source_query_concepts(
-        statement.output_components,
+
+    search_concepts: list[Concept] = statement.output_components
+    nest_where = statement.where_clause_category == SelectFiltering.IMPLICIT
+    if nest_where and statement.where_clause:
+        search_concepts = unique(
+            statement.where_clause.row_arguments + search_concepts, "address"
+        )
+        nest_where = True
+
+    ods = source_query_concepts(
+        search_concepts,
         environment=environment,
         g=graph,
     )
+    ds: GroupNode | SelectNode
+    if nest_where and statement.where_clause:
+        ods.conditions = statement.where_clause.conditional
+        ods.output_concepts = search_concepts
+        # ods.hidden_concepts = where_delta
+        ods.rebuild_cache()
+        append_existence_check(ods, environment, graph)
+        ds = GroupNode(
+            output_concepts=statement.output_components,
+            input_concepts=search_concepts,
+            parents=[ods],
+            environment=ods.environment,
+            g=ods.g,
+            partial_concepts=ods.partial_concepts,
+        )
+        # we can still check existence here.
+
+    elif statement.where_clause:
+        ds = SelectNode(
+            output_concepts=statement.output_components,
+            input_concepts=ods.input_concepts,
+            parents=[ods],
+            environment=ods.environment,
+            g=ods.g,
+            partial_concepts=ods.partial_concepts,
+            conditions=statement.where_clause.conditional,
+        )
+        append_existence_check(ds, environment, graph)
+
+    else:
+        ds = ods
+
+    final_qds = ds.resolve()
     if hooks:
         for hook in hooks:
             hook.process_root_strategy_node(ds)
-    final_qds = ds.resolve()
 
-    # we if we have a where clause doing an existence check
-    # treat that as separate subquery
-    if (where := statement.where_clause) and where.existence_arguments:
-        for subselect in where.existence_arguments:
-            if not subselect:
-                continue
-            logger.info(
-                f"{LOGGER_PREFIX} fetching existance clause inputs {[str(c) for c in subselect]}"
-            )
-            eds = source_query_concepts([*subselect], environment=environment, g=graph)
-
-            final_eds = eds.resolve()
-            first_parent = final_qds
-            first_parent.datasources.append(final_eds)
-            for x in final_eds.output_concepts:
-                if x.address not in first_parent.existence_source_map:
-                    first_parent.existence_source_map[x.address] = {final_eds}
     return final_qds
 
 

trilogy/dialect/base.py

@@ -5,7 +5,6 @@ from jinja2 import Template
 from trilogy.constants import CONFIG, logger, MagicConstants
 from trilogy.core.internal import DEFAULT_CONCEPTS
 from trilogy.core.enums import (
-    Purpose,
     FunctionType,
     WindowType,
     DatePart,
@@ -40,18 +39,17 @@ from trilogy.core.models import (
     ShowStatement,
     RowsetItem,
     MultiSelectStatement,
-    MergeStatement,
     RowsetDerivationStatement,
     ConceptDeclarationStatement,
     ImportStatement,
     RawSQLStatement,
     ProcessedRawSQLStatement,
     NumericType,
+    MergeStatementV2,
 )
 from trilogy.core.query_processor import process_query, process_persist
 from trilogy.dialect.common import render_join
 from trilogy.hooks.base_hook import BaseHook
-from trilogy.utility import unique
 from trilogy.core.enums import UnnestMode
 
 LOGGER_PREFIX = "[RENDERING]"
@@ -264,7 +262,10 @@ 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):
                 # for cases when we've optimized this
-                if len(cte.output_columns) == 1:
+                if (
+                    len(cte.output_columns) == 1
+                    and cte.condition == c.lineage.where.conditional
+                ):
                     rval = self.render_expr(c.lineage.content, cte=cte)
                 else:
                     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"
@@ -272,8 +273,6 @@ class BaseDialect:
                 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)}"
             elif isinstance(c.lineage, AggregateWrapper):
                 args = [
                     self.render_expr(v, cte)  # , alias=False)
@@ -509,35 +508,7 @@ class BaseDialect:
                         set(
                             [
                                 self.render_concept_sql(c, cte, alias=False)
-                                for c in unique(
-                                    cte.grain.components
-                                    + [
-                                        c
-                                        for c in cte.output_columns
-                                        if c.purpose in (Purpose.PROPERTY, Purpose.KEY)
-                                        and c.address
-                                        not in [x.address for x in cte.grain.components]
-                                    ]
-                                    + [
-                                        c
-                                        for c in cte.output_columns
-                                        if c.purpose == Purpose.METRIC
-                                        and any(
-                                            [
-                                                c.with_grain(cte.grain)
-                                                in cte.output_columns
-                                                for cte in cte.parent_ctes
-                                            ]
-                                        )
-                                    ]
-                                    + [
-                                        c
-                                        for c in cte.output_columns
-                                        if c.purpose == Purpose.CONSTANT
-                                        and cte.source_map[c.address] != []
-                                    ],
-                                    "address",
-                                )
+                                for c in cte.group_concepts
                             ]
                         )
                     )
@@ -563,9 +534,9 @@ class BaseDialect:
             | ShowStatement
             | ConceptDeclarationStatement
             | RowsetDerivationStatement
-            | MergeStatement
             | ImportStatement
             | RawSQLStatement
+            | MergeStatementV2
         ],
         hooks: Optional[List[BaseHook]] = None,
     ) -> List[
@@ -626,7 +597,7 @@ class BaseDialect:
                 statement,
                 (
                     ConceptDeclarationStatement,
-                    MergeStatement,
+                    MergeStatementV2,
                     ImportStatement,
                     RowsetDerivationStatement,
                 ),
@@ -675,7 +646,7 @@ class BaseDialect:
         # where assignment
         output_where = False
         if query.where_clause:
-            found = False
+            # found = False
             filter = set(
                 [
                     str(x.address)
@@ -684,16 +655,10 @@ class BaseDialect:
                 ]
             )
             query_output = set([str(z.address) for z in query.output_columns])
+            # if it wasn't an output
+            # we would have forced it up earlier and we don't need to render at this point
             if filter.issubset(query_output):
                 output_where = True
-                found = True
-
-            if not found:
-                raise NotImplementedError(
-                    f"Cannot generate query with filtering on row arguments {filter} that is"
-                    f" not a subset of the query output grain {query_output}. Try a"
-                    " filtered concept instead, or include it in the select clause"
-                )
             for ex_set in query.where_clause.existence_arguments:
                 for c in ex_set:
                     if c.address not in cte_output_map:

trilogy/dialect/bigquery.py

@@ -43,7 +43,7 @@ CREATE OR REPLACE TABLE {{ output.address.location }} AS
 {% endif %}{%- if ctes %}
 WITH {% for cte in ctes %}
 {{cte.name}} as ({{cte.statement}}){% if not loop.last %},{% endif %}{% endfor %}{% endif %}
-{%- if full_select -%}
+{%- if full_select %}
 {{full_select}}
 {% else -%}
 

trilogy/dialect/common.py

@@ -39,6 +39,17 @@ def render_join(
         )
         for key in join.joinkeys
     ]
+    if join.joinkey_pairs:
+        base_joinkeys.extend(
+            [
+                null_wrapper(
+                    f"{left_name}.{quote_character}{join.left_cte.get_alias(left_concept) if isinstance(join.left_cte, Datasource) else left_concept.safe_address}{quote_character}",
+                    f"{right_name}.{quote_character}{join.right_cte.get_alias(right_concept) if isinstance(join.right_cte, Datasource) else right_concept.safe_address}{quote_character}",
+                    left_concept,
+                )
+                for left_concept, right_concept in join.joinkey_pairs
+            ]
+        )
     if not base_joinkeys:
         base_joinkeys = ["1=1"]
     joinkeys = " AND ".join(base_joinkeys)

trilogy/dialect/duckdb.py

@@ -59,7 +59,7 @@ SELECT
     {{ join }}{% endfor %}{% endif %}
 {% if where %}WHERE
     {{ where }}
-{% endif %}
+{% endif -%}
 {%- if group_by %}
 GROUP BY {% for group in group_by %}
     {{group}}{% if not loop.last %},{% endif %}{% endfor %}{% endif %}

trilogy/dialect/presto.py

@@ -14,6 +14,7 @@ FUNCTION_MAP = {
     FunctionType.SUM: lambda x: f"sum({x[0]})",
     FunctionType.LENGTH: lambda x: f"length({x[0]})",
     FunctionType.AVG: lambda x: f"avg({x[0]})",
+    FunctionType.INDEX_ACCESS: lambda x: f"element_at({x[0]},{x[1]})",
     FunctionType.LIKE: lambda x: (
         f" CASE WHEN {x[0]} like {x[1]} THEN True ELSE False END"
     ),

trilogy/hooks/graph_hook.py

@@ -1,5 +1,5 @@
 from trilogy.hooks.base_hook import BaseHook
-from networkx import DiGraph
+import networkx as nx
 
 
 class GraphHook(BaseHook):
@@ -10,8 +10,12 @@ class GraphHook(BaseHook):
         except ImportError:
             raise ImportError("GraphHook requires matplotlib and scipy to be installed")
 
-    def query_graph_built(self, graph: DiGraph):
-        from networkx import draw_kamada_kawai
+    def query_graph_built(
+        self,
+        graph: nx.DiGraph,
+        target: str | None = None,
+        highlight_nodes: list[str] | None = None,
+    ):
         from matplotlib import pyplot as plt
 
         graph = graph.copy()
@@ -19,6 +23,47 @@ class GraphHook(BaseHook):
         for node in nodes:
             if "__preql_internal" in node:
                 graph.remove_node(node)
-        draw_kamada_kawai(graph, with_labels=True, connectionstyle="arc3, rad = 0.1")
-        # draw_spring(graph, with_labels=True, connectionstyle='arc3, rad = 0.1')
+        graph.remove_nodes_from(list(nx.isolates(graph)))
+        color_map = []
+        highlight_nodes = highlight_nodes or []
+        for node in graph:
+            if node in highlight_nodes:
+                color_map.append("orange")
+            elif str(node).startswith("ds"):
+                color_map.append("blue")
+            else:
+                color_map.append("green")
+        # pos = nx.kamada_kawai_layout(graph, scale=2)
+        pos = nx.spring_layout(graph)
+        kwargs = {}
+        if target:
+            edge_colors = []
+            descendents = nx.descendants(graph, target)
+            for edge in graph.edges():
+                if edge[0] == target:
+                    edge_colors.append("blue")
+                elif edge[1] == target:
+                    edge_colors.append("blue")
+                elif edge[1] in descendents:
+                    edge_colors.append("green")
+                else:
+                    edge_colors.append("black")
+            kwargs["edge_color"] = edge_colors
+        nx.draw(
+            graph,
+            pos=pos,
+            node_color=color_map,
+            connectionstyle="arc3, rad = 0.1",
+            **kwargs
+        )  # Draw the original graph
+        # Please note, the code below uses the original idea of re-calculating a dictionary of adjusted label positions per node.
+        pos_labels = {}
+        # For each node in the Graph
+        for aNode in graph.nodes():
+            # Get the node's position from the layout
+            x, y = pos[aNode]
+            # pos_labels[aNode] = (x+slopeX*label_ratio, y+slopeY*label_ratio)
+            pos_labels[aNode] = (x, y)
+        # Finally, redraw the labels at their new position.
+        nx.draw_networkx_labels(graph, pos=pos_labels, font_size=10)
         plt.show()

trilogy/hooks/query_debugger.py

@@ -130,4 +130,5 @@ class DebuggingHook(BaseHook):
         if self.process_nodes != PrintMode.OFF:
             printed = print_recursive_nodes(node, mode=self.process_nodes)
             for row in printed:
+                logger.info("".join([str(v) for v in row]))
                 print("".join([str(v) for v in row]))

trilogy/parsing/common.py

@@ -53,6 +53,7 @@ def constant_to_concept(
         output_purpose=Purpose.CONSTANT,
         arguments=[parent],
     )
+    fmetadata = metadata or Metadata()
     return Concept(
         name=name,
         datatype=const_function.output_datatype,
@@ -60,7 +61,7 @@ def constant_to_concept(
         lineage=const_function,
         grain=const_function.output_grain,
         namespace=namespace,
-        metadata=metadata,
+        metadata=fmetadata,
     )
 
 
@@ -105,13 +106,13 @@ def filter_item_to_concept(
     purpose: Purpose | None = None,
     metadata: Metadata | None = None,
 ) -> Concept:
-
+    fmetadata = metadata or Metadata()
     return Concept(
         name=name,
         datatype=parent.content.datatype,
         purpose=parent.content.purpose,
         lineage=parent,
-        metadata=metadata,
+        metadata=fmetadata,
         namespace=namespace,
         # filtered copies cannot inherit keys
         keys=None,
@@ -130,6 +131,7 @@ def window_item_to_concept(
     purpose: Purpose | None = None,
     metadata: Metadata | None = None,
 ) -> Concept:
+    fmetadata = metadata or Metadata()
     local_purpose, keys = get_purpose_and_keys(purpose, (parent.content,))
     if parent.order_by:
         grain = parent.over + [parent.content.output]
@@ -142,7 +144,7 @@ def window_item_to_concept(
         datatype=parent.content.datatype,
         purpose=local_purpose,
         lineage=parent,
-        metadata=metadata,
+        metadata=fmetadata,
         # filters are implicitly at the grain of the base item
         grain=Grain(components=grain),
         namespace=namespace,
@@ -162,12 +164,13 @@ def agg_wrapper_to_concept(
     )
     # anything grouped to a grain should be a property
     # at that grain
+    fmetadata = metadata or Metadata()
     aggfunction = parent.function
     out = Concept(
         name=name,
         datatype=aggfunction.output_datatype,
         purpose=Purpose.METRIC,
-        metadata=metadata,
+        metadata=fmetadata,
         lineage=parent,
         grain=Grain(components=parent.by) if parent.by else Grain(),
         namespace=namespace,