Add group_by function in Legend API

pylegend/core/language/__init__.py

@@ -36,6 +36,7 @@
     PyLegendStringLiteralExpression,
     PyLegendIntegerLiteralExpression,
     PyLegendFloatLiteralExpression,
+    convert_literal_to_literal_expression,
 )
 from pylegend.core.language.column_expressions import (
     PyLegendColumnExpression,
@@ -46,6 +47,16 @@
     PyLegendFloatColumnExpression,
 )
 from pylegend.core.language.tds_row import TdsRow
+from pylegend.core.language.aggregate_specification import AggregateSpecification
+from pylegend.core.language.primitive_collection import (
+    PyLegendPrimitiveCollection,
+    PyLegendIntegerCollection,
+    PyLegendFloatCollection,
+    PyLegendNumberCollection,
+    PyLegendStringCollection,
+    PyLegendBooleanCollection,
+    create_primitive_collection,
+)
 
 __all__: PyLegendSequence[str] = [
     "PyLegendPrimitive",
@@ -73,6 +84,16 @@
     "PyLegendStringLiteralExpression",
     "PyLegendIntegerLiteralExpression",
     "PyLegendFloatLiteralExpression",
+    "convert_literal_to_literal_expression",
 
     "TdsRow",
+    "AggregateSpecification",
+
+    "PyLegendPrimitiveCollection",
+    "PyLegendIntegerCollection",
+    "PyLegendFloatCollection",
+    "PyLegendNumberCollection",
+    "PyLegendStringCollection",
+    "PyLegendBooleanCollection",
+    "create_primitive_collection",
 ]

pylegend/core/language/aggregate_specification.py

@@ -0,0 +1,53 @@
+# Copyright 2023 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pylegend._typing import (
+    PyLegendSequence,
+    PyLegendCallable
+)
+from pylegend.core.language.primitive_collection import PyLegendPrimitiveCollection
+from pylegend.core.language import (
+    TdsRow,
+    PyLegendPrimitive,
+    PyLegendPrimitiveOrPythonPrimitive,
+)
+
+__all__: PyLegendSequence[str] = [
+    "AggregateSpecification",
+]
+
+
+class AggregateSpecification:
+    __map_fn: PyLegendCallable[[TdsRow], PyLegendPrimitiveOrPythonPrimitive]
+    __aggregate_fn: PyLegendCallable[[PyLegendPrimitiveCollection], PyLegendPrimitive]
+    __name: str
+
+    def __init__(
+            self,
+            map_fn: PyLegendCallable[[TdsRow], PyLegendPrimitiveOrPythonPrimitive],
+            aggregate_fn: PyLegendCallable[[PyLegendPrimitiveCollection], PyLegendPrimitive],
+            name: str
+    ) -> None:
+        self.__map_fn = map_fn
+        self.__aggregate_fn = aggregate_fn
+        self.__name = name
+
+    def get_name(self) -> str:
+        return self.__name
+
+    def get_map_fn(self) -> PyLegendCallable[[TdsRow], PyLegendPrimitiveOrPythonPrimitive]:
+        return self.__map_fn
+
+    def get_aggregate_fn(self) -> PyLegendCallable[[PyLegendPrimitiveCollection], PyLegendPrimitive]:
+        return self.__aggregate_fn

pylegend/core/language/operations/collection_operation_expressions.py

@@ -0,0 +1,60 @@
+# Copyright 2023 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pylegend._typing import (
+    PyLegendSequence,
+    PyLegendDict,
+)
+from pylegend.core.language.expression import (
+    PyLegendExpression,
+    PyLegendExpressionIntegerReturn,
+)
+from pylegend.core.language.operations.unary_expression import PyLegendUnaryExpression
+from pylegend.core.tds.tds_frame import FrameToSqlConfig
+from pylegend.core.sql.metamodel import (
+    Expression,
+    QuerySpecification,
+    FunctionCall,
+    QualifiedName,
+)
+
+
+__all__: PyLegendSequence[str] = [
+    "PyLegendCountExpression",
+]
+
+
+class PyLegendCountExpression(PyLegendUnaryExpression, PyLegendExpressionIntegerReturn):
+
+    @staticmethod
+    def __to_sql_func(
+            expression: Expression,
+            frame_name_to_base_query_map: PyLegendDict[str, QuerySpecification],
+            config: FrameToSqlConfig
+    ) -> Expression:
+        return FunctionCall(
+            name=QualifiedName(parts=["COUNT"]),
+            arguments=[expression],
+            distinct=False,
+            filter_=None,
+            window=None
+        )
+
+    def __init__(self, operand: PyLegendExpression) -> None:
+        PyLegendExpressionIntegerReturn.__init__(self)
+        PyLegendUnaryExpression.__init__(
+            self,
+            operand,
+            PyLegendCountExpression.__to_sql_func
+        )

pylegend/core/language/primitive_collection.py

@@ -0,0 +1,97 @@
+# Copyright 2023 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABCMeta
+from pylegend._typing import (
+    PyLegendSequence,
+)
+from pylegend.core.language.primitives.primitive import PyLegendPrimitiveOrPythonPrimitive
+from pylegend.core.language import (
+    PyLegendInteger,
+    PyLegendFloat,
+    PyLegendNumber,
+    PyLegendString,
+    PyLegendBoolean,
+    convert_literal_to_literal_expression,
+)
+from pylegend.core.language.operations.collection_operation_expressions import PyLegendCountExpression
+
+
+__all__: PyLegendSequence[str] = [
+    "PyLegendPrimitiveCollection",
+    "PyLegendIntegerCollection",
+    "PyLegendFloatCollection",
+    "PyLegendNumberCollection",
+    "PyLegendStringCollection",
+    "PyLegendBooleanCollection",
+    "create_primitive_collection"
+]
+
+
+class PyLegendPrimitiveCollection(metaclass=ABCMeta):
+    __nested: PyLegendPrimitiveOrPythonPrimitive
+
+    def __init__(self, nested: PyLegendPrimitiveOrPythonPrimitive) -> None:
+        self.__nested = nested
+
+    def count(self) -> "PyLegendInteger":
+        nested_expr = (
+            convert_literal_to_literal_expression(self.__nested) if isinstance(self.__nested, (bool, int, float, str))
+            else self.__nested.value()
+        )
+        return PyLegendInteger(PyLegendCountExpression(nested_expr))
+
+
+class PyLegendIntegerCollection(PyLegendPrimitiveCollection):
+    def __init__(self, nested: PyLegendPrimitiveOrPythonPrimitive) -> None:
+        super().__init__(nested)
+
+
+class PyLegendFloatCollection(PyLegendPrimitiveCollection):
+    def __init__(self, nested: PyLegendPrimitiveOrPythonPrimitive) -> None:
+        super().__init__(nested)
+
+
+class PyLegendNumberCollection(PyLegendPrimitiveCollection):
+    def __init__(self, nested: PyLegendPrimitiveOrPythonPrimitive) -> None:
+        super().__init__(nested)
+
+
+class PyLegendStringCollection(PyLegendPrimitiveCollection):
+    def __init__(self, nested: PyLegendPrimitiveOrPythonPrimitive) -> None:
+        super().__init__(nested)
+
+
+class PyLegendBooleanCollection(PyLegendPrimitiveCollection):
+    def __init__(self, nested: PyLegendPrimitiveOrPythonPrimitive) -> None:
+        super().__init__(nested)
+
+
+def create_primitive_collection(nested: PyLegendPrimitiveOrPythonPrimitive) -> PyLegendPrimitiveCollection:
+    if isinstance(nested, (int, PyLegendInteger)):
+        return PyLegendIntegerCollection(nested)
+
+    if isinstance(nested, (float, PyLegendFloat)):
+        return PyLegendFloatCollection(nested)
+
+    if isinstance(nested, PyLegendNumber):
+        return PyLegendNumberCollection(nested)
+
+    if isinstance(nested, (str, PyLegendString)):
+        return PyLegendStringCollection(nested)
+
+    if isinstance(nested, (bool, PyLegendBoolean)):
+        return PyLegendBooleanCollection(nested)
+
+    raise RuntimeError(f"Not supported type - {type(nested)}")  # pragma: no cover

pylegend/core/tds/legend_api/frames/functions/extend_function.py

@@ -23,23 +23,16 @@
     QuerySpecification,
     SingleColumn,
 )
-from pylegend.core.tds.tds_column import TdsColumn, PrimitiveTdsColumn
+from pylegend.core.tds.tds_column import TdsColumn
 from pylegend.core.tds.tds_frame import FrameToSqlConfig
 from pylegend.core.tds.legend_api.frames.legend_api_base_tds_frame import LegendApiBaseTdsFrame
 from pylegend.core.language import (
     TdsRow,
     PyLegendPrimitive,
     PyLegendPrimitiveOrPythonPrimitive,
-    PyLegendBoolean,
-    PyLegendString,
-    PyLegendNumber,
-    PyLegendInteger,
-    PyLegendFloat,
-    PyLegendBooleanLiteralExpression,
-    PyLegendIntegerLiteralExpression,
-    PyLegendFloatLiteralExpression,
-    PyLegendStringLiteralExpression,
+    convert_literal_to_literal_expression,
 )
+from pylegend.core.tds.legend_api.frames.functions.function_helpers import tds_column_for_primitive
 
 __all__: PyLegendSequence[str] = [
     "ExtendFunction"
@@ -78,15 +71,15 @@ def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
         tds_row = TdsRow.from_tds_frame("frame", self.__base_frame)
         for (func, name) in zip(self.__functions_list, self.__column_names_list):
             col_expr = func(tds_row)
-            if isinstance(col_expr, bool):
-                col_expr = PyLegendBoolean(PyLegendBooleanLiteralExpression(col_expr))
-            elif isinstance(col_expr, int):
-                col_expr = PyLegendInteger(PyLegendIntegerLiteralExpression(col_expr))
-            elif isinstance(col_expr, float):
-                col_expr = PyLegendFloat(PyLegendFloatLiteralExpression(col_expr))
-            elif isinstance(col_expr, str):
-                col_expr = PyLegendString(PyLegendStringLiteralExpression(col_expr))
-            col_sql_expr = col_expr.to_sql_expression({"frame": new_query}, config)
+
+            if isinstance(col_expr, (bool, int, float, str)):
+                col_sql_expr = convert_literal_to_literal_expression(col_expr).to_sql_expression(
+                    {"frame": new_query},
+                    config
+                )
+            else:
+                col_sql_expr = col_expr.to_sql_expression({"frame": new_query}, config)
+
             new_query.select.selectItems.append(
                 SingleColumn(alias=db_extension.quote_identifier(name), expression=col_sql_expr)
             )
@@ -105,18 +98,7 @@ def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
         tds_row = TdsRow.from_tds_frame("frame", self.__base_frame)
         for (func, name) in zip(self.__functions_list, self.__column_names_list):
             result = func(tds_row)
-            if isinstance(result, (bool, PyLegendBoolean)):
-                new_columns.append(PrimitiveTdsColumn.boolean_column(name))
-            elif isinstance(result, (str, PyLegendString)):
-                new_columns.append(PrimitiveTdsColumn.string_column(name))
-            elif isinstance(result, (int, PyLegendInteger)):
-                new_columns.append(PrimitiveTdsColumn.integer_column(name))
-            elif isinstance(result, (float, PyLegendFloat)):
-                new_columns.append(PrimitiveTdsColumn.float_column(name))
-            elif isinstance(result, PyLegendNumber):
-                new_columns.append(PrimitiveTdsColumn.number_column(name))
-            else:
-                raise RuntimeError("Unhandled type: " + str(type(result)))
+            new_columns.append(tds_column_for_primitive(name, result))
 
         return new_columns
 

pylegend/core/tds/legend_api/frames/functions/function_helpers.py

@@ -0,0 +1,45 @@
+# Copyright 2023 Goldman Sachs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pylegend._typing import (
+    PyLegendSequence
+)
+from pylegend.core.tds.tds_column import TdsColumn, PrimitiveTdsColumn
+from pylegend.core.language import (
+    PyLegendPrimitiveOrPythonPrimitive,
+    PyLegendBoolean,
+    PyLegendString,
+    PyLegendNumber,
+    PyLegendInteger,
+    PyLegendFloat,
+)
+
+__all__: PyLegendSequence[str] = [
+    "tds_column_for_primitive",
+]
+
+
+def tds_column_for_primitive(name: str, result: PyLegendPrimitiveOrPythonPrimitive) -> TdsColumn:
+    if isinstance(result, (bool, PyLegendBoolean)):
+        return PrimitiveTdsColumn.boolean_column(name)
+    elif isinstance(result, (str, PyLegendString)):
+        return PrimitiveTdsColumn.string_column(name)
+    elif isinstance(result, (int, PyLegendInteger)):
+        return PrimitiveTdsColumn.integer_column(name)
+    elif isinstance(result, (float, PyLegendFloat)):
+        return PrimitiveTdsColumn.float_column(name)
+    elif isinstance(result, PyLegendNumber):
+        return PrimitiveTdsColumn.number_column(name)
+    else:
+        raise RuntimeError("Unhandled type: " + str(type(result)))  # pragma: no cover