pytrilogy 0.0.1.102__py3-none-any.whl

trilogy/core/functions.py

@@ -0,0 +1,320 @@
+from trilogy.core.models import (
+    Function,
+    Concept,
+    AggregateWrapper,
+    Parenthetical,
+    arg_to_datatype,
+    WindowItem,
+    DataType,
+    ListType,
+    StructType,
+    MapType,
+)
+from trilogy.core.enums import FunctionType, Purpose, Granularity, DatePart
+from trilogy.core.exceptions import InvalidSyntaxException
+from trilogy.constants import MagicConstants
+from typing import Optional
+
+
+def create_function_derived_concept(
+    name: str,
+    namespace: str,
+    operator: FunctionType,
+    arguments: list[Concept],
+    output_type: Optional[DataType | ListType | StructType | MapType] = None,
+    output_purpose: Optional[Purpose] = None,
+) -> Concept:
+    purpose = (
+        function_args_to_output_purpose(arguments)
+        if output_purpose is None
+        else output_purpose
+    )
+    output_type = arg_to_datatype(arguments[0]) if output_type is None else output_type
+    return Concept(
+        name=name,
+        namespace=namespace,
+        datatype=output_type,
+        purpose=purpose,
+        lineage=Function(
+            operator=operator,
+            arguments=arguments,
+            output_datatype=output_type,
+            output_purpose=purpose,
+            arg_count=len(arguments),
+        ),
+    )
+
+
+def argument_to_purpose(arg) -> Purpose:
+    if isinstance(arg, Function):
+        return arg.output_purpose
+    elif isinstance(arg, AggregateWrapper):
+        return arg.function.output_purpose
+    elif isinstance(arg, Parenthetical):
+        return argument_to_purpose(arg.content)
+    elif isinstance(arg, WindowItem):
+        return Purpose.PROPERTY
+    elif isinstance(arg, Concept):
+        return arg.purpose
+    elif isinstance(arg, (int, float, str, bool, list)):
+        return Purpose.CONSTANT
+    elif isinstance(arg, DataType):
+        return Purpose.CONSTANT
+    elif isinstance(arg, DatePart):
+        return Purpose.CONSTANT
+    elif isinstance(arg, MagicConstants):
+        return Purpose.CONSTANT
+    else:
+        raise ValueError(f"Cannot parse arg purpose for {arg} of type {type(arg)}")
+
+
+def function_args_to_output_purpose(args) -> Purpose:
+    has_metric = False
+    has_non_constant = False
+    has_non_single_row_constant = False
+    if not args:
+        return Purpose.CONSTANT
+    for arg in args:
+        purpose = argument_to_purpose(arg)
+        if purpose == Purpose.METRIC:
+            has_metric = True
+        if purpose != Purpose.CONSTANT:
+            has_non_constant = True
+        if isinstance(arg, Concept) and arg.granularity != Granularity.SINGLE_ROW:
+            has_non_single_row_constant = True
+    if args and not has_non_constant and not has_non_single_row_constant:
+        return Purpose.CONSTANT
+    if has_metric:
+        return Purpose.METRIC
+    return Purpose.PROPERTY
+
+
+def Unnest(args: list[Concept]) -> Function:
+    output = arg_to_datatype(args[0])
+    if isinstance(output, (ListType)):
+        output = output.value_data_type
+    return Function(
+        operator=FunctionType.UNNEST,
+        arguments=args,
+        output_datatype=output,
+        output_purpose=Purpose.KEY,
+        arg_count=1,
+        valid_inputs={
+            DataType.ARRAY,
+            DataType.LIST,
+            ListType(type=DataType.STRING),
+            ListType(type=DataType.INTEGER),
+        },
+    )
+
+
+def Group(args: list[Concept]) -> Function:
+    output = args[0]
+    return Function(
+        operator=FunctionType.GROUP,
+        arguments=args,
+        output_datatype=output.datatype,
+        output_purpose=Purpose.PROPERTY,
+        arg_count=-1,
+    )
+
+
+def Count(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.COUNT,
+        arguments=args,
+        output_datatype=DataType.INTEGER,
+        output_purpose=Purpose.METRIC,
+        arg_count=1,
+    )
+
+
+def CountDistinct(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.COUNT_DISTINCT,
+        arguments=args,
+        output_datatype=DataType.INTEGER,
+        output_purpose=Purpose.METRIC,
+        arg_count=1,
+    )
+
+
+def Max(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.MAX,
+        arguments=args,
+        output_datatype=args[0].datatype,
+        output_purpose=Purpose.METRIC,
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.DATE,
+            DataType.DATETIME,
+            DataType.TIMESTAMP,
+        },
+        arg_count=1,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def Min(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.MIN,
+        arguments=args,
+        output_datatype=args[0].datatype,
+        output_purpose=Purpose.METRIC,
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.DATE,
+            DataType.DATETIME,
+            DataType.TIMESTAMP,
+        },
+        arg_count=1,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def Split(args: list[Concept]) -> Function:
+    # TODO: overload this for non-string types?
+    return Function(
+        operator=FunctionType.SPLIT,
+        arguments=args,
+        # first arg sets properties
+        output_datatype=ListType(type=DataType.STRING),
+        output_purpose=function_args_to_output_purpose(args),
+        valid_inputs={DataType.STRING},
+        arg_count=2,
+    )
+
+
+def IndexAccess(args: list[Concept]):
+    return Function(
+        operator=FunctionType.INDEX_ACCESS,
+        arguments=args,
+        output_datatype=(
+            args[0].datatype.value_data_type
+            if isinstance(args[0].datatype, ListType)
+            else args[0].datatype
+        ),
+        output_purpose=Purpose.PROPERTY,
+        valid_inputs=[
+            {
+                DataType.LIST,
+                ListType(type=DataType.STRING),
+                ListType(type=DataType.INTEGER),
+            },
+            {
+                DataType.INTEGER,
+            },
+        ],
+        arg_count=2,
+    )
+
+
+def AttrAccess(args: list[Concept]):
+    return Function(
+        operator=FunctionType.ATTR_ACCESS,
+        arguments=args,
+        output_datatype=args[0].field_map[args[1]].datatype,  # type: ignore
+        output_purpose=Purpose.PROPERTY,
+        valid_inputs=[
+            {DataType.STRUCT},
+            {
+                DataType.STRING,
+            },
+        ],
+        arg_count=2,
+    )
+
+
+def Abs(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.ABS,
+        arguments=args,
+        output_datatype=args[0].datatype,
+        output_purpose=function_args_to_output_purpose(args),
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+        },
+        arg_count=1,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def Coalesce(args: list[Concept]) -> Function:
+    non_null = [x for x in args if not x == MagicConstants.NULL]
+    if not len(set(arg_to_datatype(x) for x in non_null if x)) == 1:
+        raise InvalidSyntaxException(
+            f"All arguments to coalesce must be of the same type, have {set(arg_to_datatype(x) for x in args)}"
+        )
+    return Function(
+        operator=FunctionType.COALESCE,
+        arguments=args,
+        output_datatype=arg_to_datatype(non_null[0]),
+        output_purpose=function_args_to_output_purpose(non_null),
+        arg_count=-1,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def CurrentDate(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.CURRENT_DATE,
+        arguments=args,
+        output_datatype=DataType.DATE,
+        output_purpose=Purpose.CONSTANT,
+        arg_count=0,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def CurrentDatetime(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.CURRENT_DATETIME,
+        arguments=args,
+        output_datatype=DataType.DATE,
+        output_purpose=Purpose.CONSTANT,
+        arg_count=0,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def IsNull(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.IS_NULL,
+        arguments=args,
+        output_datatype=DataType.BOOL,
+        output_purpose=function_args_to_output_purpose(args),
+        arg_count=1,
+        # output_grain=Grain(components=arguments),
+    )
+
+
+def StrPos(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.STRPOS,
+        arguments=args,
+        output_datatype=DataType.INTEGER,
+        output_purpose=function_args_to_output_purpose(args),
+        arg_count=2,
+        valid_inputs=[
+            {DataType.STRING},
+            {DataType.STRING},
+        ],
+    )
+
+
+def SubString(args: list[Concept]) -> Function:
+    return Function(
+        operator=FunctionType.SUBSTRING,
+        arguments=args,
+        output_datatype=DataType.STRING,
+        output_purpose=function_args_to_output_purpose(args),
+        arg_count=3,
+        valid_inputs=[{DataType.STRING}, {DataType.INTEGER}, {DataType.INTEGER}],
+    )

trilogy/core/graph_models.py

@@ -0,0 +1,55 @@
+import networkx as nx
+
+from trilogy.core.models import Concept, Datasource
+
+
+def concept_to_node(input: Concept) -> str:
+    # if input.purpose == Purpose.METRIC:
+    #     return f"c~{input.namespace}.{input.name}@{input.grain}"
+    return f"c~{input.namespace}.{input.name}@{input.grain}"
+
+
+def datasource_to_node(input: Datasource) -> str:
+    # if isinstance(input, JoinedDataSource):
+    #     return "ds~join~" + ",".join(
+    #         [datasource_to_node(sub) for sub in input.datasources]
+    #     )
+    return f"ds~{input.namespace}.{input.identifier}"
+
+
+class ReferenceGraph(nx.DiGraph):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def add_node(self, node_for_adding, **attr):
+        if isinstance(node_for_adding, Concept):
+            node_name = concept_to_node(node_for_adding)
+            attr["type"] = "concept"
+            attr["concept"] = node_for_adding
+            attr["grain"] = node_for_adding.grain
+        elif isinstance(node_for_adding, Datasource):
+            node_name = datasource_to_node(node_for_adding)
+            attr["type"] = "datasource"
+            attr["ds"] = node_for_adding
+            attr["grain"] = node_for_adding.grain
+        else:
+            node_name = node_for_adding
+        super().add_node(node_name, **attr)
+
+    def add_edge(self, u_of_edge, v_of_edge, **attr):
+        if isinstance(u_of_edge, Concept):
+            orig = u_of_edge
+            u_of_edge = concept_to_node(u_of_edge)
+            if u_of_edge not in self.nodes:
+                self.add_node(orig)
+        elif isinstance(u_of_edge, Datasource):
+            u_of_edge = datasource_to_node(u_of_edge)
+
+        if isinstance(v_of_edge, Concept):
+            orig = v_of_edge
+            v_of_edge = concept_to_node(v_of_edge)
+            if v_of_edge not in self.nodes:
+                self.add_node(orig)
+        elif isinstance(v_of_edge, Datasource):
+            v_of_edge = datasource_to_node(v_of_edge)
+        super().add_edge(u_of_edge, v_of_edge, **attr)

trilogy/core/internal.py

@@ -0,0 +1,37 @@
+from trilogy.core.models import Concept, DataType, Function
+from trilogy.core.enums import Purpose, FunctionType
+from trilogy.core.constants import ALL_ROWS_CONCEPT, INTERNAL_NAMESPACE
+
+
+DEFAULT_CONCEPTS = {
+    ALL_ROWS_CONCEPT: Concept(
+        name=ALL_ROWS_CONCEPT,
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.INTEGER,
+        purpose=Purpose.CONSTANT,
+        lineage=Function(
+            operator=FunctionType.CONSTANT,
+            arguments=[1],
+            output_datatype=DataType.INTEGER,
+            output_purpose=Purpose.CONSTANT,
+        ),
+    ),
+    "concept_name": Concept(
+        name="concept_name",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+    ),
+    "datasource": Concept(
+        name="datasource",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+    ),
+    "query_text": Concept(
+        name="query_text",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+    ),
+}