pytrilogy 0.3.142__cp312-cp312-win_amd64.whl

trilogy/core/query_processor.py

@@ -0,0 +1,596 @@
+from collections import defaultdict
+from math import ceil
+from typing import Dict, List, Optional, Set, Tuple, Union
+
+from trilogy.constants import CONFIG, logger
+from trilogy.core.constants import CONSTANT_DATASET
+from trilogy.core.enums import BooleanOperator, DatasourceState, SourceType
+from trilogy.core.env_processor import generate_graph
+from trilogy.core.ergonomics import generate_cte_names
+from trilogy.core.models.author import MultiSelectLineage, SelectLineage
+from trilogy.core.models.build import (
+    BuildConcept,
+    BuildConditional,
+    BuildDatasource,
+    BuildFunction,
+    BuildMultiSelectLineage,
+    BuildParamaterizedConceptReference,
+    BuildSelectLineage,
+    Factory,
+)
+from trilogy.core.models.core import DataType
+from trilogy.core.models.datasource import Datasource
+from trilogy.core.models.environment import Environment
+from trilogy.core.models.execute import (
+    CTE,
+    BaseJoin,
+    CTEConceptPair,
+    InstantiatedUnnestJoin,
+    Join,
+    QueryDatasource,
+    RecursiveCTE,
+    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.execute import (
+    MaterializedDataset,
+    ProcessedCopyStatement,
+    ProcessedQuery,
+    ProcessedQueryPersist,
+)
+from trilogy.hooks.base_hook import BaseHook
+from trilogy.utility import unique
+
+LOGGER_PREFIX = "[QUERY BUILD]"
+
+
+def base_join_to_join(
+    base_join: BaseJoin | UnnestJoin, ctes: List[CTE | UnionCTE]
+) -> Join | InstantiatedUnnestJoin:
+    """This function converts joins at the datasource level
+    to joins at the CTE level"""
+    if isinstance(base_join, UnnestJoin):
+        object_to_unnest = base_join.parent.arguments[0]
+        if not isinstance(
+            object_to_unnest,
+            (BuildConcept | BuildParamaterizedConceptReference | BuildFunction),
+        ):
+            raise ValueError(f"Unnest join must be a concept; got {object_to_unnest}")
+        return InstantiatedUnnestJoin(
+            object_to_unnest=object_to_unnest,
+            alias=base_join.alias,
+        )
+
+    def get_datasource_cte(
+        datasource: BuildDatasource | QueryDatasource,
+    ) -> CTE | UnionCTE:
+        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
+            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)
+    else:
+        # multiple left ctes
+        left_cte = None
+    right_cte = get_datasource_cte(base_join.right_datasource)
+    if base_join.concept_pairs:
+        final_pairs = [
+            CTEConceptPair(
+                left=pair.left,
+                right=pair.right,
+                existing_datasource=pair.existing_datasource,
+                modifiers=pair.modifiers,
+                cte=get_datasource_cte(pair.existing_datasource),
+            )
+            for pair in base_join.concept_pairs
+        ]
+    elif base_join.concepts and base_join.left_datasource:
+        final_pairs = [
+            CTEConceptPair(
+                left=concept,
+                right=concept,
+                existing_datasource=base_join.left_datasource,
+                modifiers=[],
+                cte=get_datasource_cte(
+                    base_join.left_datasource,
+                ),
+            )
+            for concept in base_join.concepts
+        ]
+    else:
+        final_pairs = []
+    return Join(
+        left_cte=left_cte,
+        right_cte=right_cte,
+        jointype=base_join.join_type,
+        joinkey_pairs=final_pairs,
+        modifiers=base_join.modifiers,
+    )
+
+
+def generate_source_map(
+    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
+    for qdk, qdv in query_datasource.source_map.items():
+        unnest = [x for x in qdv if isinstance(x, UnnestJoin)]
+        for _ in unnest:
+            source_map[qdk] = []
+        if (
+            qdk not in source_map
+            and len(qdv) == 1
+            and isinstance(list(qdv)[0], UnnestJoin)
+        ):
+            source_map[qdk] = []
+        basic = [x for x in qdv if isinstance(x, BuildDatasource)]
+        for base in basic:
+            source_map[qdk].append(base.safe_identifier)
+
+        ctes = [x for x in qdv if isinstance(x, QueryDatasource)]
+        if ctes:
+            names = set([x.safe_identifier for x in ctes])
+            matches = [
+                cte for cte in all_new_ctes if cte.source.safe_identifier in names
+            ]
+
+            if not matches and names:
+                raise SyntaxError(
+                    f"Missing parent CTEs for source map; expecting {names}, have {[cte.source.safe_identifier for cte in all_new_ctes]}"
+                )
+            for cte in matches:
+                output_address = [
+                    x.address
+                    for x in cte.output_columns
+                    if x.address not in [z.address for z in cte.partial_concepts]
+                ]
+                if qdk in output_address:
+                    source_map[qdk].append(cte.safe_identifier)
+            # now do a pass that accepts partials
+            for cte in matches:
+                if qdk not in source_map:
+                    source_map[qdk] = [cte.safe_identifier]
+        if qdk not in source_map:
+            if not qdv:
+                source_map[qdk] = []
+            elif CONFIG.validate_missing:
+                raise ValueError(
+                    f"Missing {qdk} in {source_map}, source map {query_datasource.source_map} "
+                )
+
+    # existence lookups use a separate map
+    # as they cannot be referenced in row resolution
+    existence_source_map: Dict[str, list[str]] = defaultdict(list)
+    for ek, ev in query_datasource.existence_source_map.items():
+        ids = set([x.safe_identifier for x in ev])
+        ematches = [
+            cte.name for cte in all_new_ctes if cte.source.safe_identifier in ids
+        ]
+        existence_source_map[ek] = ematches
+    return {
+        k: [] if not v else list(set(v)) for k, v in source_map.items()
+    }, existence_source_map
+
+
+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]
+    concepts = unique(concepts, "address")
+    return QueryDatasource(
+        output_concepts=concepts,
+        input_concepts=concepts,
+        source_map=sub_select,
+        grain=datasource.grain,
+        datasources=[datasource],
+        joins=[],
+        partial_concepts=[x.concept for x in datasource.columns if not x.is_complete],
+    )
+
+
+def generate_cte_name(full_name: str, name_map: dict[str, str]) -> str:
+    cte_names = generate_cte_names()
+    if CONFIG.human_identifiers:
+        if full_name in name_map:
+            return name_map[full_name]
+        suffix = ""
+        idx = len(name_map)
+        if idx >= len(cte_names):
+            int = ceil(idx / len(cte_names))
+            suffix = f"_{int}"
+        valid = [x for x in cte_names if x + suffix not in name_map.values()]
+        lookup = valid[0]
+        new_name = f"{lookup}{suffix}"
+        name_map[full_name] = new_name
+        return new_name
+    else:
+        return full_name.replace("<", "").replace(">", "").replace(",", "_")
+
+
+def resolve_cte_base_name_and_alias_v2(
+    name: str,
+    source: QueryDatasource,
+    source_map: Dict[str, list[str]],
+    raw_joins: List[Join | InstantiatedUnnestJoin],
+) -> Tuple[str | None, str | None]:
+    if not source.datasources:
+        return None, None
+    if (
+        isinstance(source.datasources[0], BuildDatasource)
+        and not source.datasources[0].name == CONSTANT_DATASET
+    ):
+        ds = source.datasources[0]
+        return ds.safe_location, ds.safe_identifier
+
+    joins: List[Join] = [join for join in raw_joins if isinstance(join, Join)]
+    if joins and len(joins) > 0:
+        candidates = [x.left_cte.name for x in joins if x.left_cte]
+        for join in joins:
+            if join.joinkey_pairs:
+                candidates += [x.cte.name for x in join.joinkey_pairs if x.cte]
+        disallowed = [x.right_cte.name for x in joins]
+        try:
+            cte = [y for y in candidates if y not in disallowed][0]
+            return cte, cte
+        except IndexError:
+            raise SyntaxError(
+                f"Invalid join configuration {candidates} {disallowed} for {name}",
+            )
+    counts: dict[str, int] = defaultdict(lambda: 0)
+    output_addresses = [x.address for x in source.output_concepts]
+    input_address = [x.address for x in source.input_concepts]
+    for k, v in source_map.items():
+        for vx in v:
+            if k in output_addresses:
+                counts[vx] = counts[vx] + 1
+
+            if k in input_address:
+                counts[vx] = counts[vx] + 1
+
+            counts[vx] = counts[vx]
+    if counts:
+        return max(counts, key=counts.get), max(counts, key=counts.get)  # type: ignore
+    return None, None
+
+
+def datasource_to_cte(
+    query_datasource: QueryDatasource, name_map: dict[str, str]
+) -> 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,
+            order_by=query_datasource.ordering,
+        )
+        return final
+
+    if len(query_datasource.datasources) > 1 or any(
+        [isinstance(x, QueryDatasource) for x in query_datasource.datasources]
+    ):
+        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_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:
+        # source is the first datasource of the query datasource
+        if query_datasource.datasources:
+
+            source = query_datasource.datasources[0]
+            # this is required to ensure that constant datasets
+            # render properly on initial access; since they have
+            # no actual source
+            if source.name == CONSTANT_DATASET:
+                source_map = {k: [] for k in query_datasource.source_map}
+                existence_map = source_map
+            else:
+                source_map = {
+                    k: [] if not v else [source.safe_identifier]
+                    for k, v in query_datasource.source_map.items()
+                }
+                existence_map = source_map
+        else:
+            source_map = {k: [] for k in query_datasource.source_map}
+            existence_map = source_map
+
+    human_id = generate_cte_name(query_datasource.identifier, name_map)
+
+    final_joins = [
+        base_join_to_join(join, [x for x in parents if isinstance(x, (CTE, UnionCTE))])
+        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
+    )
+    cte_class = CTE
+
+    if query_datasource.source_type == SourceType.RECURSIVE:
+        cte_class = RecursiveCTE
+        # extra_kwargs['left_recursive_concept'] = query_datasource.left
+    cte = cte_class(
+        name=human_id,
+        source=query_datasource,
+        # output columns are what are selected/grouped by
+        output_columns=[
+            c.with_grain(query_datasource.grain)
+            for c in query_datasource.output_concepts
+        ],
+        source_map=source_map,
+        existence_source_map=existence_map,
+        # related columns include all referenced columns, such as filtering
+        joins=final_joins,
+        grain=query_datasource.grain,
+        group_to_grain=query_datasource.group_required,
+        # we restrict parent_ctes to one level
+        # as this set is used as the base for rendering the query
+        parent_ctes=parents,
+        condition=query_datasource.condition,
+        partial_concepts=query_datasource.partial_concepts,
+        nullable_concepts=query_datasource.nullable_concepts,
+        join_derived_concepts=query_datasource.join_derived_concepts,
+        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")
+    for x in cte.output_columns:
+        if (
+            x.address not in cte.source_map
+            and not any(y in cte.source_map for y in x.pseudonyms)
+            and CONFIG.validate_missing
+        ):
+            raise ValueError(
+                f"Missing {x.address} in {cte.source_map}, source map {cte.source.source_map.keys()} "
+            )
+
+    return cte
+
+
+def get_query_node(
+    environment: Environment,
+    statement: SelectLineage | MultiSelectLineage,
+    history: History | None = None,
+) -> StrategyNode:
+    if not statement.output_components:
+        raise ValueError(f"Statement has no output components {statement}")
+    history = history or History(base_environment=environment)
+    logger.info(
+        f"{LOGGER_PREFIX} building query node for {statement.output_components} grain {statement.grain}"
+    )
+    build_statement: BuildSelectLineage | BuildMultiSelectLineage = Factory(
+        environment=environment,
+    ).build(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 {build_statement.output_components} grain {build_statement.grain}"
+    )
+
+    search_concepts: list[BuildConcept] = build_statement.output_components
+
+    ods: StrategyNode = source_query_concepts(
+        output_concepts=search_concepts,
+        environment=build_environment,
+        g=graph,
+        conditions=build_statement.where_clause,
+        history=history,
+    )
+    if not ods:
+        raise ValueError(
+            f"Could not find source query concepts for {[x.address for x in search_concepts]}"
+        )
+    ds: StrategyNode = ods
+    if build_statement.having_clause:
+        final = build_statement.having_clause.conditional
+        if ds.conditions:
+            final = BuildConditional(
+                left=ds.conditions,
+                right=build_statement.having_clause.conditional,
+                operator=BooleanOperator.AND,
+            )
+        ds = SelectNode(
+            output_concepts=build_statement.output_components,
+            input_concepts=ds.usable_outputs,
+            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
+
+
+def get_query_datasources(
+    environment: Environment,
+    statement: SelectStatement | MultiSelectStatement,
+    hooks: Optional[List[BaseHook]] = None,
+) -> QueryDatasource:
+    ds = get_query_node(environment, statement.as_lineage(environment))
+
+    final_qds = ds.resolve()
+
+    if hooks:
+        for hook in hooks:
+            hook.process_root_strategy_node(ds)
+
+    return final_qds
+
+
+def flatten_ctes(input: CTE | UnionCTE) -> list[CTE | UnionCTE]:
+    output = [input]
+    for cte in input.parent_ctes:
+        output += flatten_ctes(cte)
+    return output
+
+
+def process_auto(
+    environment: Environment,
+    statement: PersistStatement | SelectStatement,
+    hooks: List[BaseHook] | None = None,
+):
+    if isinstance(statement, PersistStatement):
+        return process_persist(environment, statement, hooks)
+    elif isinstance(statement, SelectStatement):
+        return process_query(environment, statement, hooks)
+    elif isinstance(statement, ConceptDeclarationStatement):
+        return None
+    raise ValueError(f"Do not know how to process {type(statement)}")
+
+
+def process_persist(
+    environment: Environment,
+    statement: PersistStatement,
+    hooks: List[BaseHook] | None = None,
+) -> ProcessedQueryPersist:
+    ds: Datasource = environment.datasources.get(
+        statement.datasource.identifier, statement.datasource
+    )
+    original_status = ds.status
+    # set to unpublished to avoid circular refs
+    try:
+        ds.status = DatasourceState.UNPUBLISHED
+        select = process_query(
+            environment=environment, statement=statement.select, hooks=hooks
+        )
+    except:
+        raise
+    finally:
+        ds.status = original_status
+
+    # build our object to return
+    arg_dict = {k: v for k, v in select.__dict__.items()}
+    partition_by: list[str] = []
+    partition_types: list[DataType] = []
+    for addr in statement.partition_by:
+        for c in statement.datasource.columns:
+            if c.concept.address == addr and c.is_concrete:
+                partition_by.append(c.alias)  # type: ignore
+                partition_types.append(c.concept.output_datatype)
+                break
+    return ProcessedQueryPersist(
+        **arg_dict,
+        output_to=MaterializedDataset(address=statement.address),
+        persist_mode=statement.persist_mode,
+        partition_by=partition_by,
+        datasource=statement.datasource,
+        partition_types=partition_types,
+    )
+
+
+def process_copy(
+    environment: Environment,
+    statement: CopyStatement,
+    hooks: List[BaseHook] | None = None,
+) -> ProcessedCopyStatement:
+    select = process_query(
+        environment=environment, statement=statement.select, hooks=hooks
+    )
+
+    # build our object to return
+    arg_dict = {k: v for k, v in select.__dict__.items()}
+    return ProcessedCopyStatement(
+        **arg_dict,
+        target=statement.target,
+        target_type=statement.target_type,
+    )
+
+
+def process_query(
+    environment: Environment,
+    statement: SelectStatement | MultiSelectStatement,
+    hooks: List[BaseHook] | None = None,
+) -> ProcessedQuery:
+    hooks = hooks or []
+
+    root_datasource = get_query_datasources(
+        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_cte(root_datasource, environment.cte_name_map)
+
+    for hook in hooks:
+        hook.process_root_cte(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
+    for cte in raw_ctes:
+        if cte.name not in seen:
+            seen[cte.name] = cte
+        else:
+            # merge them up
+            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 | UnionCTE] = list(seen.values())
+
+    root_cte.limit = statement.limit
+    root_cte.hidden_concepts = statement.hidden_components
+
+    final_ctes = optimize_ctes(deduped_ctes, root_cte, statement)
+
+    return ProcessedQuery(
+        order_by=root_cte.order_by,
+        limit=statement.limit,
+        output_columns=statement.output_components,
+        ctes=final_ctes,
+        base=root_cte,
+        hidden_columns=set([x for x in statement.hidden_components]),
+        local_concepts=statement.local_concepts,
+        locally_derived=statement.locally_derived,
+    )

trilogy/core/statements/README.md

@@ -0,0 +1,35 @@
+# Statement Design
+
+
+## Assert statements
+Used for DQ checks.
+
+Unique in that we need to constrain to specific datasources/validation. 
+
+
+
+## Concept check
+
+# value comparison
+
+two scalar values
+
+assert expr from_datasource? = expr from_datasource?
+
+
+assert not max(len(name.split(' '))) =1; 
+
+
+assert sum(revenue) from datasource1 = sum(revenue) from datasource2;
+
+
+## parallel block
+
+begin parallel;
+
+assert a == 1;
+assert b == 2;
+
+end parallel;
+
+