pytrilogy 0.0.3.86__tar.gz → 0.0.3.89__tar.gz

{pytrilogy-0.0.3.86/pytrilogy.egg-info → pytrilogy-0.0.3.89}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytrilogy
-Version: 0.0.3.86
+Version: 0.0.3.89
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89/pytrilogy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pytrilogy
-Version: 0.0.3.86
+Version: 0.0.3.89
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/__init__.py

@@ -4,6 +4,6 @@ from trilogy.dialect.enums import Dialects
 from trilogy.executor import Executor
 from trilogy.parser import parse
 
-__version__ = "0.0.3.86"
+__version__ = "0.0.3.89"
 
 __all__ = ["parse", "Executor", "Dialects", "Environment", "CONFIG"]

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/core/enums.py

@@ -13,6 +13,11 @@ class UnnestMode(Enum):
     SNOWFLAKE = "snowflake"
 
 
+class GroupMode(Enum):
+    AUTO = "auto"
+    BY_INDEX = "by_index"
+
+
 class ConceptSource(Enum):
     MANUAL = "manual"
     CTE = "cte"
@@ -156,6 +161,10 @@ class FunctionType(Enum):
     ARRAY_TRANSFORM = "array_transform"
     ARRAY_TO_STRING = "array_to_string"
 
+    # MAP
+    MAP_KEYS = "map_keys"
+    MAP_VALUES = "map_values"
+
     # TEXT AND MAYBE MORE
     SPLIT = "split"
     LENGTH = "len"

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/core/functions.py

@@ -70,7 +70,7 @@ def get_coalesce_output_type(args: list[Any]) -> CONCRETE_TYPES:
 
 
 def get_transform_output_type(args: list[Any]) -> CONCRETE_TYPES:
-    return arg_to_datatype(args[2])
+    return ArrayType(type=arg_to_datatype(args[2]))
 
 
 def get_index_output_type(
@@ -92,7 +92,7 @@ def get_attr_datatype(
     lookup = args[1]
     datatype = arg_to_datatype(arg)
     if isinstance(datatype, StructType):
-        return arg_to_datatype(datatype.fields_map[lookup])
+        return datatype.field_types[lookup]
     return datatype
 
 
@@ -178,6 +178,20 @@ def get_date_trunc_output(
         raise InvalidSyntaxException(f"Date truncation not supported for {target}")
 
 
+def get_map_key_type(arg):
+    arg_datatype = arg_to_datatype(arg)
+    if isinstance(arg_datatype, MapType):
+        return ArrayType(type=arg_datatype.key_data_type)
+    return ArrayType(type=DataType.STRING)
+
+
+def get_map_value_type(arg):
+    arg_datatype = arg_to_datatype(arg)
+    if isinstance(arg_datatype, MapType):
+        return ArrayType(type=arg_datatype.value_data_type)
+    return ArrayType(type=DataType.STRING)
+
+
 FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
     FunctionType.ALIAS: FunctionConfig(
         arg_count=1,
@@ -254,6 +268,22 @@ FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
         output_type_function=get_index_output_type,
         arg_count=2,
     ),
+    FunctionType.MAP_KEYS: FunctionConfig(
+        valid_inputs={
+            DataType.MAP,
+        },
+        output_purpose=Purpose.PROPERTY,
+        output_type_function=lambda args: get_map_key_type(args[0]),
+        arg_count=1,
+    ),
+    FunctionType.MAP_VALUES: FunctionConfig(
+        valid_inputs={
+            DataType.MAP,
+        },
+        output_purpose=Purpose.PROPERTY,
+        output_type_function=lambda args: get_map_value_type(args[0]),
+        arg_count=1,
+    ),
     FunctionType.ARRAY_DISTINCT: FunctionConfig(
         valid_inputs={
             DataType.ARRAY,
@@ -905,6 +935,7 @@ class FunctionFactory:
                 output_purpose = Purpose.METRIC
             else:
                 output_purpose = Purpose.PROPERTY
+
         return Function(
             operator=operator,
             arguments=full_args,  # type: ignore

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/core/models/author.py

@@ -164,6 +164,26 @@ class ConceptRef(Addressable, Namespaced, DataTyped, Mergeable, BaseModel):
     def with_reference_replacement(self, source: str, target: Expr | ArgBinding):
         if self.address == source:
             return target
+
+        # a reference might be to an attribute of a struct that is bound late
+        # if the replacement is a parent in the access path; replace reference
+        # with an attribute access call
+        candidates = [f"local.{self.address}", self.address]
+        for candidate in candidates:
+            if not candidate.startswith(f"{source}."):
+                continue
+            attribute = self.address.rsplit(".", 1)[1]
+            dtype = arg_to_datatype(target)
+            if not isinstance(dtype, StructType):
+                continue
+            output_type = dtype.field_types.get(attribute, DataType.UNKNOWN)
+            return Function(
+                arguments=[target, self.address.rsplit(".", 1)[1]],
+                operator=FunctionType.ATTR_ACCESS,
+                arg_count=2,
+                output_datatype=output_type,
+                output_purpose=Purpose.PROPERTY,
+            )
         return self
 
 
@@ -1622,6 +1642,9 @@ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
     ] = None
     arguments: Sequence[FuncArgs]
 
+    class Config:
+        frozen = True
+
     def __repr__(self):
         return f'{self.operator.value}({",".join([str(a) for a in self.arguments])})'
 
@@ -1632,6 +1655,16 @@ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
     def datatype(self):
         return self.output_datatype
 
+    @field_validator("output_datatype")
+    @classmethod
+    def parse_output_datatype(cls, v, info: ValidationInfo):
+        values = info.data
+        if values.get("operator") == FunctionType.ATTR_ACCESS:
+            print(v)
+            if isinstance(v, StructType):
+                raise SyntaxError
+        return v
+
     @field_validator("arguments", mode="before")
     @classmethod
     def parse_arguments(cls, v, info: ValidationInfo):
@@ -1730,8 +1763,16 @@ class Function(DataTyped, ConceptArgs, Mergeable, Namespaced, BaseModel):
         ]
         if self.output_datatype == DataType.UNKNOWN:
             new_output = merge_datatypes([arg_to_datatype(x) for x in nargs])
+
+            if self.operator == FunctionType.ATTR_ACCESS:
+                if isinstance(new_output, StructType):
+                    new_output = new_output.field_types[str(nargs[1])]
         else:
             new_output = self.output_datatype
+        # this is not ideal - see hacky logic for datatypes above
+        # we need to figure out how to patch properly
+        # should use function factory, but does not have environment access
+        # probably move all datatype resolution to build?
         return Function.model_construct(
             operator=self.operator,
             arguments=nargs,

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/core/models/build.py

@@ -1649,8 +1649,11 @@ class Factory:
             full = self.local_concepts[base.address]
             if isinstance(full, BuildConcept):
                 return full
-        raw = self.environment.concepts[base.address]
-        return self.build(raw)
+        if base.address in self.environment.concepts:
+            raw = self.environment.concepts[base.address]
+            return self.build(raw)
+        # this will error by design - TODO - more helpful message?
+        return self.build(self.environment.concepts[base.address])
 
     @build.register
     def _(self, base: CaseWhen) -> BuildCaseWhen:

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/core/models/core.py

@@ -27,7 +27,7 @@ from pydantic_core import core_schema
 from trilogy.constants import (
     MagicConstants,
 )
-from trilogy.core.enums import Ordering
+from trilogy.core.enums import Modifier, Ordering
 
 
 class DataTyped(ABC):
@@ -211,9 +211,19 @@ class MapType(BaseModel):
         return self.key_type
 
 
+class StructComponent(BaseModel):
+    name: str
+    type: TYPEDEF_TYPES
+    modifiers: list[Modifier] = []
+
+    @field_validator("type", mode="plain")
+    def validate_Type(cls, v):
+        return v
+
+
 class StructType(BaseModel):
-    fields: Sequence[TYPEDEF_TYPES]
-    fields_map: Dict[str, DataTyped | int | float | str]
+    fields: Sequence[StructComponent | TYPEDEF_TYPES]
+    fields_map: Dict[str, DataTyped | int | float | str | StructComponent]
 
     @field_validator("fields", mode="plain")
     def validate_type(cls, v):
@@ -234,6 +244,22 @@ class StructType(BaseModel):
     def value(self):
         return self.data_type.value
 
+    @property
+    def field_types(self) -> Dict[str, CONCRETE_TYPES]:
+        out: Dict[str, CONCRETE_TYPES] = {}
+        keys = list(self.fields_map.keys())
+        for idx, field in enumerate(self.fields):
+            if isinstance(field, StructComponent):
+                out[field.name] = arg_to_datatype(field.type)
+            elif isinstance(field, DataTyped):
+                out[keys[idx]] = field.output_datatype
+            else:
+                out[keys[idx]] = field
+        return out
+
+    def __hash__(self):
+        return hash(str(self))
+
 
 class ListWrapper(Generic[VT], UserList):
     """Used to distinguish parsed list objects from other lists"""
@@ -401,6 +427,8 @@ def arg_to_datatype(arg) -> CONCRETE_TYPES:
         return DataType.STRING
     elif isinstance(arg, float):
         return DataType.FLOAT
+    elif isinstance(arg, DataType):
+        return arg
     elif isinstance(arg, NumericType):
         return arg
     elif isinstance(arg, TraitDataType):
@@ -420,6 +448,8 @@ def arg_to_datatype(arg) -> CONCRETE_TYPES:
         return DataType.DATETIME
     elif isinstance(arg, date):
         return DataType.DATE
+    elif isinstance(arg, StructComponent):
+        return arg_to_datatype(arg.type)
     else:
         raise ValueError(
             f"Cannot parse arg datatype for arg of raw type {type(arg)} value {arg}"

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/core/models/execute.py

@@ -1076,6 +1076,11 @@ class UnionCTE(BaseModel):
     def group_to_grain(self) -> bool:
         return False
 
+    @property
+    def group_concepts(self) -> List[BuildConcept]:
+        # unions should always be on unique sets
+        return []
+
     def __add__(self, other):
         if not isinstance(other, UnionCTE) or not other.name == self.name:
             raise SyntaxError("Cannot merge union CTEs")

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/dialect/base.py

@@ -15,6 +15,7 @@ from trilogy.core.enums import (
     ComparisonOperator,
     DatePart,
     FunctionType,
+    GroupMode,
     Ordering,
     ShowCategory,
     UnnestMode,
@@ -186,6 +187,9 @@ FUNCTION_MAP = {
     FunctionType.ARRAY: lambda x: f"[{', '.join(x)}]",
     FunctionType.DATE_LITERAL: lambda x: f"date '{x}'",
     FunctionType.DATETIME_LITERAL: lambda x: f"datetime '{x}'",
+    # MAP
+    FunctionType.MAP_KEYS: lambda x: f"map_keys({x[0]})",
+    FunctionType.MAP_VALUES: lambda x: f"map_values({x[0]})",
     # ARRAY
     FunctionType.ARRAY_SUM: lambda x: f"array_sum({x[0]})",
     FunctionType.ARRAY_DISTINCT: lambda x: f"array_distinct({x[0]})",
@@ -338,6 +342,7 @@ class BaseDialect:
     DATATYPE_MAP = DATATYPE_MAP
     COMPLEX_DATATYPE_MAP = COMPLEX_DATATYPE_MAP
     UNNEST_MODE = UnnestMode.CROSS_APPLY
+    GROUP_MODE = GroupMode.AUTO
 
     def __init__(self, rendering: Rendering | None = None):
         self.rendering = rendering or CONFIG.rendering
@@ -775,6 +780,33 @@ class BaseDialect:
         else:
             raise ValueError(f"Unable to render type {type(e)} {e}")
 
+    def render_cte_group_by(
+        self, cte: CTE | UnionCTE, select_columns
+    ) -> Optional[list[str]]:
+
+        if not cte.group_to_grain:
+            return None
+        base = set(
+            [self.render_concept_sql(c, cte, alias=False) for c in cte.group_concepts]
+        )
+        if self.GROUP_MODE == GroupMode.AUTO:
+            return sorted(list(base))
+
+        else:
+            # find the index of each column in the select columns
+            final = []
+            found = []
+            for idx, c in enumerate(select_columns):
+                pre_alias = c.split(" as ")[0]
+                if pre_alias in base:
+                    final.append(str(idx + 1))
+                    found.append(pre_alias)
+            if not all(c in found for c in base):
+                raise ValueError(
+                    f"Group by columns {base} not found in select columns {select_columns}"
+                )
+            return final
+
     def render_cte(self, cte: CTE | UnionCTE, auto_sort: bool = True) -> CompiledCTE:
         if isinstance(cte, UnionCTE):
             base_statement = f"\n{cte.operator}\n".join(
@@ -926,20 +958,7 @@ class BaseDialect:
                     if cte.order_by
                     else None
                 ),
-                group_by=(
-                    sorted(
-                        list(
-                            set(
-                                [
-                                    self.render_concept_sql(c, cte, alias=False)
-                                    for c in cte.group_concepts
-                                ]
-                            )
-                        )
-                    )
-                    if cte.group_to_grain
-                    else None
-                ),
+                group_by=self.render_cte_group_by(cte, select_columns),
             ),
         )
 

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/dialect/presto.py

@@ -2,7 +2,7 @@ from typing import Any, Callable, Mapping
 
 from jinja2 import Template
 
-from trilogy.core.enums import FunctionType, UnnestMode, WindowType
+from trilogy.core.enums import FunctionType, GroupMode, UnnestMode, WindowType
 from trilogy.core.models.core import DataType
 from trilogy.dialect.base import BaseDialect
 
@@ -94,6 +94,7 @@ class PrestoDialect(BaseDialect):
         DataType.STRING: "VARCHAR",
     }
     UNNEST_MODE = UnnestMode.PRESTO
+    GROUP_MODE = GroupMode.BY_INDEX
 
 
 class TrinoDialect(PrestoDialect):

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/parsing/parse_engine.py

@@ -89,6 +89,7 @@ from trilogy.core.models.core import (
     MapType,
     MapWrapper,
     NumericType,
+    StructComponent,
     StructType,
     TraitDataType,
     TupleWrapper,
@@ -443,20 +444,39 @@ class ParseToObjects(Transformer):
     def MINUS(self, args) -> str:
         return "-"
 
+    @v_args(meta=True)
+    def struct_component(self, meta: Meta, args) -> StructComponent:
+        modifiers = []
+        for arg in args:
+            if isinstance(arg, Modifier):
+                modifiers.append(arg)
+        return StructComponent(name=args[0], type=args[1], modifiers=modifiers)
+
     @v_args(meta=True)
     def struct_type(self, meta: Meta, args) -> StructType:
         final: list[
-            DataType | MapType | ArrayType | NumericType | StructType | Concept
+            DataType
+            | MapType
+            | ArrayType
+            | NumericType
+            | StructType
+            | StructComponent
+            | Concept
         ] = []
         for arg in args:
-            new = self.environment.concepts.__getitem__(  # type: ignore
-                key=arg, line_no=meta.line
-            )
-            final.append(new)
+            if isinstance(arg, StructComponent):
+                final.append(arg)
+            else:
+                new = self.environment.concepts.__getitem__(  # type: ignore
+                    key=arg, line_no=meta.line
+                )
+                final.append(new)
 
         return StructType(
             fields=final,
-            fields_map={x.name: x for x in final if isinstance(x, Concept)},
+            fields_map={
+                x.name: x for x in final if isinstance(x, (Concept, StructComponent))
+            },
         )
 
     def list_type(self, args) -> ArrayType:
@@ -681,6 +701,7 @@ class ParseToObjects(Transformer):
                 environment=self.environment,
                 metadata=metadata,
             )
+
             # let constant purposes exist to support round-tripping
             # as a build concept may end up with a constant based on constant inlining happening recursively
             if purpose == Purpose.KEY and concept.purpose != Purpose.KEY:
@@ -902,7 +923,6 @@ class ParseToObjects(Transformer):
             name=output,
             metadata=metadata,
         )
-
         return ConceptTransform(function=transformation, output=concept)
 
     @v_args(meta=True)
@@ -2026,6 +2046,16 @@ class ParseToObjects(Transformer):
     def fbool(self, meta, args):
         return self.function_factory.create_function(args, FunctionType.BOOL, meta)
 
+    @v_args(meta=True)
+    def fmap_keys(self, meta, args):
+        return self.function_factory.create_function(args, FunctionType.MAP_KEYS, meta)
+
+    @v_args(meta=True)
+    def fmap_values(self, meta, args):
+        return self.function_factory.create_function(
+            args, FunctionType.MAP_VALUES, meta
+        )
+
     @v_args(meta=True)
     def farray_sum(self, meta, args):
         return self.function_factory.create_function(args, FunctionType.ARRAY_SUM, meta)

{pytrilogy-0.0.3.86 → pytrilogy-0.0.3.89}/trilogy/parsing/trilogy.lark

@@ -181,7 +181,8 @@
     //indexing into an expression is a function
     index_access: atom "[" int_lit "]"
     map_key_access: atom "[" string_lit "]"
-    attr_access: atom  "." string_lit 
+    _ATTR_ACCESS.1 : "getattr("i
+    attr_access: ( _ATTR_ACCESS atom "," string_lit ")") | (atom  "." string_lit)
 
     ?expr: comparison_root | between_root
 
@@ -309,6 +310,14 @@
     farray_transform: _ARRAY_TRANSFORM expr ","  transform_lambda   ")"
 
     _array_functions: farray_sum | farray_distinct | farray_sort | farray_transform | farray_to_string
+
+    //map_functions
+    _MAP_KEYS.1: "map_keys("i
+    fmap_keys: _MAP_KEYS expr ")"
+    _MAP_VALUES.1: "map_values("i
+    fmap_values: _MAP_VALUES expr ")"
+
+    _map_functions: fmap_keys | fmap_values
     
     // special aggregate
     _GROUP.1: "group("i
@@ -380,7 +389,7 @@
     
     _date_functions: fdate | fdate_add | fdate_sub | fdate_diff | fdatetime | ftimestamp | fsecond | fminute | fhour | fday | fday_of_week | fweek | fmonth | fquarter | fyear | fdate_part | fdate_trunc
     
-    _static_functions:  _string_functions | _math_functions | _array_functions
+    _static_functions:  _string_functions | _math_functions | _array_functions | _map_functions
 
     custom_function:  "@" IDENTIFIER  "(" (expr ",")* expr ")"
 
@@ -424,8 +433,8 @@
     MODIFIER: /OPTIONAL|PARTIAL|NULLABLE/i
 
     SHORTHAND_MODIFIER: /~|\?/
-
-    struct_type: "struct"i "<" ((data_type | IDENTIFIER) ",")* (data_type | IDENTIFIER) ","? ">" 
+    struct_component: IDENTIFIER ":" data_type concept_nullable_modifier? metadata?
+    struct_type: "struct"i "<" ((struct_component | IDENTIFIER) ",")* (struct_component | IDENTIFIER) ","? ">" 
 
     list_type: ("list"i | "array"i) "<" (data_type | IDENTIFIER) ">"