Number exp/log operations

finos · Jul 26, 2023 · d89086e · d89086e
1 parent b01cee7
commit d89086e
Show file tree

Hide file tree

Showing 6 changed files with 118 additions and 0 deletions.
diff --git a/pylegend/core/databse/sql_to_string/db_extension.py b/pylegend/core/databse/sql_to_string/db_extension.py
@@ -88,6 +88,8 @@
     FloorExpression,
     SqrtExpression,
     CbrtExpression,
+    ExpExpression,
+    LogExpression,
 )
 
 
@@ -354,6 +356,10 @@ def expression_processor(
         return extension.process_sqrt_expression(expression, config)
     elif isinstance(expression, CbrtExpression):
         return extension.process_cbrt_expression(expression, config)
+    elif isinstance(expression, ExpExpression):
+        return extension.process_exp_expression(expression, config)
+    elif isinstance(expression, LogExpression):
+        return extension.process_log_expression(expression, config)
     else:
         raise ValueError("Unsupported expression type: " + str(type(expression)))  # pragma: no cover
 
@@ -986,6 +992,16 @@ def process_cbrt_expression(self, expr: CbrtExpression, config: SqlToStringConfi
             value=self.process_expression(expr.value, config)
         )
 
+    def process_exp_expression(self, expr: ExpExpression, config: SqlToStringConfig) -> str:
+        return "EXP({value})".format(
+            value=self.process_expression(expr.value, config)
+        )
+
+    def process_log_expression(self, expr: LogExpression, config: SqlToStringConfig) -> str:
+        return "LN({value})".format(
+            value=self.process_expression(expr.value, config)
+        )
+
     def process_qualified_name(self, qualified_name: QualifiedName, config: SqlToStringConfig) -> str:
         return qualified_name_processor(qualified_name, self, config)
 

diff --git a/pylegend/core/language/operations/number_operation_expressions.py b/pylegend/core/language/operations/number_operation_expressions.py
@@ -40,6 +40,8 @@
     FloorExpression,
     SqrtExpression,
     CbrtExpression,
+    ExpExpression,
+    LogExpression,
 )
 
 
@@ -59,6 +61,8 @@
     "PyLegendNumberFloorExpression",
     "PyLegendNumberSqrtExpression",
     "PyLegendNumberCbrtExpression",
+    "PyLegendNumberExpExpression",
+    "PyLegendNumberLogExpression",
 ]
 
 
@@ -363,3 +367,41 @@ def __init__(self, operand: PyLegendExpressionNumberReturn) -> None:
             operand,
             PyLegendNumberCbrtExpression.__to_sql_func
         )
+
+
+class PyLegendNumberExpExpression(PyLegendUnaryExpression, PyLegendExpressionNumberReturn):
+
+    @staticmethod
+    def __to_sql_func(
+            expression: Expression,
+            frame_name_to_base_query_map: PyLegendDict[str, QuerySpecification],
+            config: FrameToSqlConfig
+    ) -> Expression:
+        return ExpExpression(expression)
+
+    def __init__(self, operand: PyLegendExpressionNumberReturn) -> None:
+        PyLegendExpressionNumberReturn.__init__(self)
+        PyLegendUnaryExpression.__init__(
+            self,
+            operand,
+            PyLegendNumberExpExpression.__to_sql_func
+        )
+
+
+class PyLegendNumberLogExpression(PyLegendUnaryExpression, PyLegendExpressionNumberReturn):
+
+    @staticmethod
+    def __to_sql_func(
+            expression: Expression,
+            frame_name_to_base_query_map: PyLegendDict[str, QuerySpecification],
+            config: FrameToSqlConfig
+    ) -> Expression:
+        return LogExpression(expression)
+
+    def __init__(self, operand: PyLegendExpressionNumberReturn) -> None:
+        PyLegendExpressionNumberReturn.__init__(self)
+        PyLegendUnaryExpression.__init__(
+            self,
+            operand,
+            PyLegendNumberLogExpression.__to_sql_func
+        )
diff --git a/pylegend/core/language/primitives/number.py b/pylegend/core/language/primitives/number.py
@@ -43,6 +43,8 @@
     PyLegendNumberFloorExpression,
     PyLegendNumberSqrtExpression,
     PyLegendNumberCbrtExpression,
+    PyLegendNumberExpExpression,
+    PyLegendNumberLogExpression,
 )
 from pylegend.core.sql.metamodel import (
     Expression,
@@ -210,6 +212,12 @@ def sqrt(self) -> "PyLegendNumber":
     def cbrt(self) -> "PyLegendNumber":
         return PyLegendNumber(PyLegendNumberCbrtExpression(self.__value))
 
+    def exp(self) -> "PyLegendNumber":
+        return PyLegendNumber(PyLegendNumberExpExpression(self.__value))
+
+    def log(self) -> "PyLegendNumber":
+        return PyLegendNumber(PyLegendNumberLogExpression(self.__value))
+
     @staticmethod
     def __convert_to_number_expr(
             val: PyLegendUnion[int, float, "PyLegendInteger", "PyLegendFloat", "PyLegendNumber"]

diff --git a/pylegend/core/sql/metamodel_extension.py b/pylegend/core/sql/metamodel_extension.py
@@ -35,6 +35,8 @@
     "FloorExpression",
     "SqrtExpression",
     "CbrtExpression",
+    "ExpExpression",
+    "LogExpression",
 ]
 
 
@@ -200,3 +202,25 @@ def __init__(
     ) -> None:
         super().__init__(_type="cbrtExpression")
         self.value = value
+
+
+class ExpExpression(Expression):
+    value: "Expression"
+
+    def __init__(
+        self,
+        value: "Expression",
+    ) -> None:
+        super().__init__(_type="expExpression")
+        self.value = value
+
+
+class LogExpression(Expression):
+    value: "Expression"
+
+    def __init__(
+        self,
+        value: "Expression",
+    ) -> None:
+        super().__init__(_type="logExpression")
+        self.value = value
diff --git a/pylegend/tests/core/database/test_sql_to_string.py b/pylegend/tests/core/database/test_sql_to_string.py
@@ -85,6 +85,8 @@
     FloorExpression,
     SqrtExpression,
     CbrtExpression,
+    ExpExpression,
+    LogExpression,
 )
 
 
@@ -1256,3 +1258,17 @@ def test_process_cbrt_expression(self) -> None:
 
         expr = CbrtExpression(IntegerLiteral(10))
         assert extension.process_expression(expr, config) == "CBRT(10)"
+
+    def test_process_exp_expression(self) -> None:
+        extension = SqlToStringDbExtension()
+        config = SqlToStringConfig(SqlToStringFormat(pretty=False))
+
+        expr = ExpExpression(IntegerLiteral(10))
+        assert extension.process_expression(expr, config) == "EXP(10)"
+
+    def test_process_log_expression(self) -> None:
+        extension = SqlToStringDbExtension()
+        config = SqlToStringConfig(SqlToStringFormat(pretty=False))
+
+        expr = LogExpression(IntegerLiteral(10))
+        assert extension.process_expression(expr, config) == "LN(10)"
diff --git a/pylegend/tests/core/language/primitives/test_number.py b/pylegend/tests/core/language/primitives/test_number.py
@@ -151,6 +151,18 @@ def test_number_cbrt_expr(self) -> None:
         assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).cbrt()) == \
                'CBRT(("root".col2 + "root".col1))'
 
+    def test_number_exp_expr(self) -> None:
+        assert self.__generate_sql_string(lambda x: x.get_number("col2").exp()) == \
+               'EXP("root".col2)'
+        assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).exp()) == \
+               'EXP(("root".col2 + "root".col1))'
+
+    def test_number_log_expr(self) -> None:
+        assert self.__generate_sql_string(lambda x: x.get_number("col2").log()) == \
+               'LN("root".col2)'
+        assert self.__generate_sql_string(lambda x: (x.get_number("col2") + x.get_number("col1")).log()) == \
+               'LN(("root".col2 + "root".col1))'
+
     def __generate_sql_string(self, f) -> str:  # type: ignore
         return self.db_extension.process_expression(
             f(self.tds_row).to_sql_expression({"t": self.base_query}, self.frame_to_sql_config),