pylegend 0.11.0__py3-none-any.whl → 0.12.0__py3-none-any.whl

pylegend/core/tds/legendql_api/frames/legendql_api_tds_frame.py

@@ -22,6 +22,7 @@ from pylegend.core.language.legendql_api.legendql_api_custom_expressions import
     LegendQLApiWindow,
     LegendQLApiPartialFrame,
     LegendQLApiWindowReference,
+    LegendQLApiWindowFrame,
 )
 from pylegend.core.language.legendql_api.legendql_api_tds_row import LegendQLApiTdsRow
 from pylegend.core.tds.tds_frame import (
@@ -60,7 +61,17 @@ class LegendQLApiTdsFrame(PyLegendTdsFrame, metaclass=ABCMeta):
         pass  # pragma: no cover
 
     @abstractmethod
-    def distinct(self) -> "LegendQLApiTdsFrame":
+    def distinct(
+            self,
+            columns: PyLegendOptional[PyLegendUnion[
+                str,
+                PyLegendList[str],
+                PyLegendCallable[
+                    [LegendQLApiTdsRow],
+                    PyLegendUnion[LegendQLApiPrimitive, PyLegendList[LegendQLApiPrimitive]]
+                ]
+            ]] = None
+    ) -> "LegendQLApiTdsFrame":
         pass  # pragma: no cover
 
     @abstractmethod
@@ -261,7 +272,8 @@ class LegendQLApiTdsFrame(PyLegendTdsFrame, metaclass=ABCMeta):
                         ]
                     ]
                 ]
-            ] = None
+            ] = None,
+            frame: PyLegendOptional[LegendQLApiWindowFrame] = None
     ) -> LegendQLApiWindow:
         pass
 
@@ -325,3 +337,22 @@ class LegendQLApiTdsFrame(PyLegendTdsFrame, metaclass=ABCMeta):
             ]
     ) -> "LegendQLApiTdsFrame":
         pass  # pragma: no cover
+
+    @abstractmethod
+    def rows(
+            self,
+            start: PyLegendUnion[str, int],
+            end: PyLegendUnion[str, int]) -> LegendQLApiWindowFrame:
+        pass  # pragma: no cover
+
+    @abstractmethod
+    def range(
+            self,
+            *,
+            number_start: PyLegendOptional[PyLegendUnion[str, int, float]] = None,
+            number_end: PyLegendOptional[PyLegendUnion[str, int, float]] = None,
+            duration_start: PyLegendOptional[PyLegendUnion[str, int, float]] = None,
+            duration_start_unit: PyLegendOptional[str] = None,
+            duration_end: PyLegendOptional[PyLegendUnion[str, int, float]] = None,
+            duration_end_unit: PyLegendOptional[str] = None) -> LegendQLApiWindowFrame:
+        pass  # pragma: no cover

pylegend/core/tds/pandas_api/frames/functions/assign_function.py

@@ -28,8 +28,11 @@ from pylegend.core.language import (
     PyLegendNumber,
     PyLegendBoolean,
     PyLegendString,
+    PyLegendDate,
+    PyLegendDateTime
 )
 from pylegend.core.language.pandas_api.pandas_api_tds_row import PandasApiTdsRow
+from pylegend.core.language.shared.literal_expressions import convert_literal_to_literal_expression
 from pylegend.core.sql.metamodel import (
     QuerySpecification,
     SingleColumn,
@@ -50,7 +53,7 @@ class AssignFunction(PandasApiAppliedFunction):
 
     @classmethod
     def name(cls) -> str:
-        return "assign"
+        return "assign"  # pragma: no cover
 
     def __init__(
             self,
@@ -74,22 +77,53 @@ class AssignFunction(PandasApiAppliedFunction):
             copy_query(base_query)
         )
 
-        tds_row = PandasApiTdsRow.from_tds_frame("frame", self.__base_frame)
+        base_cols = {c.get_name() for c in self.__base_frame.columns()}
+        tds_row = PandasApiTdsRow.from_tds_frame("c", self.__base_frame)
         for col, func in self.__col_definitions.items():
             res = func(tds_row)
-            if not isinstance(res, PyLegendPrimitive):
-                raise RuntimeError("Constants not supported")
-            new_col_expr = res.to_sql_expression(
-                {"frame": base_query},
+            res_expr = res if isinstance(res, PyLegendPrimitive) else convert_literal_to_literal_expression(res)
+            new_col_expr = res_expr.to_sql_expression(
+                {"c": base_query},
                 config
             )
-            new_query.select.selectItems.append(
-                SingleColumn(alias=db_extension.quote_identifier(col), expression=new_col_expr)
-            )
+
+            alias = db_extension.quote_identifier(col)
+            if col in base_cols:
+                for i, si in enumerate(new_query.select.selectItems):
+                    if isinstance(si, SingleColumn) and si.alias == alias:
+                        new_query.select.selectItems[i] = SingleColumn(alias=alias, expression=new_col_expr)
+
+            else:
+                new_query.select.selectItems.append(SingleColumn(alias=alias, expression=new_col_expr))
         return new_query
 
     def to_pure(self, config: FrameToPureConfig) -> str:
-        raise NotImplementedError("to_pure is not implemented yet")  # pragma: no cover
+        tds_row = PandasApiTdsRow.from_tds_frame("c", self.__base_frame)
+        base_cols = [c.get_name() for c in self.__base_frame.columns()]
+
+        assigned_exprs: PyLegendDict[str, str] = {}
+        for col, func in self.__col_definitions.items():
+            res = func(tds_row)
+            res_expr = res if isinstance(res, PyLegendPrimitive) else convert_literal_to_literal_expression(res)
+            assigned_exprs[col] = res_expr.to_pure_expression(config)
+
+        # build project clauses
+        clauses: PyLegendList[str] = []
+
+        for col in base_cols:
+            if col in assigned_exprs:
+                clauses.append(f"{col}:c|{assigned_exprs[col]}")
+            else:
+                clauses.append(f"{col}:c|$c.{col}")
+
+        for col, pure_expr in assigned_exprs.items():
+            if col not in base_cols:
+                clauses.append(f"{col}:c|{pure_expr}")
+
+        return (
+            f"{self.__base_frame.to_pure(config)}{config.separator(1)}"
+            f"->project(~[{', '.join(clauses)}])"
+        )
 
     def base_frame(self) -> PandasApiBaseTdsFrame:
         return self.__base_frame
@@ -99,21 +133,29 @@ class AssignFunction(PandasApiAppliedFunction):
 
     def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
         new_cols = [c.copy() for c in self.__base_frame.columns()]
+        base_cols = {c.get_name() for c in self.__base_frame.columns()}
         tds_row = PandasApiTdsRow.from_tds_frame("frame", self.__base_frame)
         for col, func in self.__col_definitions.items():
-            res = func(tds_row)
-            if isinstance(res, (int, PyLegendInteger)):
-                new_cols.append(PrimitiveTdsColumn.integer_column(col))
-            elif isinstance(res, (float, PyLegendFloat)):
-                new_cols.append(PrimitiveTdsColumn.float_column(col))
-            elif isinstance(res, PyLegendNumber):
-                new_cols.append(PrimitiveTdsColumn.number_column(col))
-            elif isinstance(res, (bool, PyLegendBoolean)):
-                new_cols.append(PrimitiveTdsColumn.boolean_column(col))
-            elif isinstance(res, (str, PyLegendString)):
-                new_cols.append(PrimitiveTdsColumn.string_column(col))
-            else:
-                raise RuntimeError("Type not supported")
+            if col not in base_cols:
+                res = func(tds_row)
+                if isinstance(res, (int, PyLegendInteger)):
+                    new_cols.append(PrimitiveTdsColumn.integer_column(col))
+                elif isinstance(res, (float, PyLegendFloat)):
+                    new_cols.append(PrimitiveTdsColumn.float_column(col))
+                elif isinstance(res, PyLegendNumber):
+                    new_cols.append(PrimitiveTdsColumn.number_column(col))  # pragma: no cover
+                elif isinstance(res, (bool, PyLegendBoolean)):
+                    new_cols.append(
+                        PrimitiveTdsColumn.boolean_column(col)
+                    )  # pragma: no cover (Boolean column not supported in PURE)
+                elif isinstance(res, (str, PyLegendString)):
+                    new_cols.append(PrimitiveTdsColumn.string_column(col))
+                elif isinstance(res, (datetime, PyLegendDateTime)):
+                    new_cols.append(PrimitiveTdsColumn.datetime_column(col))
+                elif isinstance(res, (date, PyLegendDate)):
+                    new_cols.append(PrimitiveTdsColumn.date_column(col))
+                else:
+                    raise RuntimeError("Type not supported")
         return new_cols
 
     def validate(self) -> bool:

pylegend/core/tds/pandas_api/frames/functions/drop.py

@@ -160,10 +160,10 @@ class PandasApiDropFunction(PandasApiAppliedFunction):
             self.__columns = _normalize_columns(self.__columns)  # type: ignore
 
         if isinstance(self.__inplace, (bool, PyLegendBoolean)):
-            if self.__inplace is False:
-                raise NotImplementedError(f"Only inplace=True is supported. Got inplace={self.__inplace!r}")
+            if self.__inplace is True:
+                raise NotImplementedError(f"Only inplace=False is supported. Got inplace={self.__inplace!r}")
         else:
-            raise TypeError(f"Inplace must be True. Got inplace={self.__inplace!r}")  # pragma: no cover
+            raise TypeError(f"Inplace must be False. Got inplace={self.__inplace!r}")  # pragma: no cover
 
         if valid_paramters == 0:
             raise ValueError("Need to specify at least one of 'labels' or 'columns'")

pylegend/core/tds/pandas_api/frames/functions/dropna.py

@@ -0,0 +1,167 @@
+# Copyright 2025 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 functools import reduce
+
+from pylegend._typing import (
+    PyLegendSequence,
+    PyLegendOptional,
+    PyLegendList,
+    PyLegendUnion,
+)
+from pylegend.core.language.pandas_api.pandas_api_tds_row import PandasApiTdsRow
+from pylegend.core.sql.metamodel import (
+    QuerySpecification,
+    LogicalBinaryType,
+    LogicalBinaryExpression, BooleanLiteral,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import (
+    PandasApiAppliedFunction,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import (
+    PandasApiBaseTdsFrame,
+)
+from pylegend.core.tds.sql_query_helpers import copy_query
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToSqlConfig, FrameToPureConfig
+
+__all__: PyLegendSequence[str] = ["PandasApiDropnaFunction"]
+
+
+class PandasApiDropnaFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __axis: PyLegendUnion[int, str]
+    __how: str
+    __thresh: PyLegendOptional[int]
+    __subset: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str]]]
+    __inplace: bool
+    __ignore_index: bool
+
+    @classmethod
+    def name(cls) -> str:
+        return "dropna"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            axis: PyLegendUnion[int, str],
+            how: str,
+            thresh: PyLegendOptional[int],
+            subset: PyLegendOptional[PyLegendUnion[str, PyLegendSequence[str]]],
+            inplace: bool,
+            ignore_index: bool
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__axis = axis
+        self.__how = how
+        self.__thresh = thresh
+        self.__subset = subset
+        self.__inplace = inplace
+        self.__ignore_index = ignore_index
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        base_query = self.__base_frame.to_sql_query_object(config)
+        new_query = copy_query(base_query)
+
+        if self.__subset is not None:
+            cols_to_check = self.__subset
+        else:
+            cols_to_check = [c.get_name() for c in self.__base_frame.columns()]
+
+        if not cols_to_check:
+            if self.__how == 'all':
+                new_query.where = BooleanLiteral(value=False)
+            return new_query
+
+        tds_row = PandasApiTdsRow.from_tds_frame("c", self.__base_frame)
+
+        filter_expr = None
+
+        conditions = [tds_row[col].is_not_null() for col in cols_to_check]
+        if self.__how == "any":
+            filter_expr = reduce(lambda x, y: x & y, conditions)
+        else:  # "all"
+            filter_expr = reduce(lambda x, y: x | y, conditions)
+
+        sql_expr = filter_expr.to_sql_expression({"c": new_query}, config)
+        if new_query.where is None:
+            new_query.where = sql_expr
+        else:
+            new_query.where = LogicalBinaryExpression(
+                type_=LogicalBinaryType.AND,
+                left=new_query.where,
+                right=sql_expr
+            )
+
+        return new_query
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        base_pure = self.__base_frame.to_pure(config)
+        if self.__subset is not None:
+            cols_to_check = self.__subset
+        else:
+            cols_to_check = [c.get_name() for c in self.__base_frame.columns()]
+
+        if not cols_to_check:
+            if self.__how == 'all':
+                return f"{base_pure}{config.separator(1)}->filter(c|1!=1)"
+            return base_pure
+
+        tds_row = PandasApiTdsRow.from_tds_frame("c", self.__base_frame)
+        conditions = [tds_row[col].is_not_null() for col in cols_to_check]
+
+        if self.__how == "any":
+            filter_expr = reduce(lambda x, y: x & y, conditions)
+        else:  # "all"
+            filter_expr = reduce(lambda x, y: x | y, conditions)
+
+        pure_expr = filter_expr.to_pure_expression(config)
+        return f"{base_pure}{config.separator(1)}->filter(c|{pure_expr})"
+
+    def base_frame(self) -> PandasApiBaseTdsFrame:
+        return self.__base_frame
+
+    def tds_frame_parameters(self) -> PyLegendList["PandasApiBaseTdsFrame"]:
+        return []
+
+    def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
+        return [c.copy() for c in self.__base_frame.columns()]
+
+    def validate(self) -> bool:
+        if self.__axis not in (0, 1, "index", "columns"):
+            raise ValueError(f"No axis named {self.__axis} for object type TdsFrame")
+        if self.__axis in (1, "columns"):
+            raise NotImplementedError("axis=1 is not supported yet in Pandas API dropna")
+
+        if self.__thresh is not None:
+            raise NotImplementedError("thresh parameter is not supported yet in Pandas API dropna")
+
+        if self.__how not in ("any", "all"):
+            raise ValueError(f"invalid how option: {self.__how}")
+
+        if self.__subset is not None:
+            if not isinstance(self.__subset, (list, tuple, set)):
+                raise TypeError(f"subset must be a list, tuple or set of column names. Got {type(self.__subset)}")
+            valid_cols = {c.get_name() for c in self.__base_frame.columns()}
+            invalid_cols = [s for s in self.__subset if s not in valid_cols]
+            if invalid_cols:
+                raise KeyError(f"{invalid_cols}")
+
+        if self.__inplace:
+            raise NotImplementedError("inplace=True is not supported yet in Pandas API dropna")
+
+        if self.__ignore_index:
+            raise NotImplementedError("ignore_index=True is not supported yet in Pandas API dropna")
+
+        return True

pylegend/core/tds/pandas_api/frames/functions/fillna.py

@@ -0,0 +1,162 @@
+# Copyright 2025 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 datetime import date, datetime
+
+from pylegend._typing import (
+    PyLegendSequence,
+    PyLegendOptional,
+    PyLegendList,
+    PyLegendUnion,
+    PyLegendDict
+)
+from pylegend.core.language import (
+    convert_literal_to_literal_expression
+)
+from pylegend.core.language.pandas_api.pandas_api_tds_row import PandasApiTdsRow
+from pylegend.core.sql.metamodel import (
+    QuerySpecification,
+    SingleColumn, FunctionCall, QualifiedName,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_applied_function_tds_frame import (
+    PandasApiAppliedFunction,
+)
+from pylegend.core.tds.pandas_api.frames.pandas_api_base_tds_frame import (
+    PandasApiBaseTdsFrame,
+)
+from pylegend.core.tds.sql_query_helpers import copy_query
+from pylegend.core.tds.tds_column import TdsColumn
+from pylegend.core.tds.tds_frame import FrameToSqlConfig, FrameToPureConfig
+
+__all__: PyLegendSequence[str] = ["PandasApiFillnaFunction"]
+
+
+class PandasApiFillnaFunction(PandasApiAppliedFunction):
+    __base_frame: PandasApiBaseTdsFrame
+    __value: PyLegendUnion[
+        int, float, str, bool, date, datetime,
+        PyLegendDict[str, PyLegendUnion[int, float, str, bool, date, datetime]]
+    ]
+    __axis: PyLegendOptional[PyLegendUnion[int, str]]
+    __inplace: bool
+    __limit: PyLegendOptional[int]
+
+    @classmethod
+    def name(cls) -> str:
+        return "fillna"  # pragma: no cover
+
+    def __init__(
+            self,
+            base_frame: PandasApiBaseTdsFrame,
+            value: PyLegendUnion[
+                int, float, str, bool, date, datetime,
+                PyLegendDict[str, PyLegendUnion[int, float, str, bool, date, datetime]]
+            ],
+            axis: PyLegendOptional[PyLegendUnion[int, str]],
+            inplace: bool,
+            limit: PyLegendOptional[int]
+    ) -> None:
+        self.__base_frame = base_frame
+        self.__value = value
+        self.__axis = axis
+        self.__inplace = inplace
+        self.__limit = limit
+
+    def to_sql(self, config: FrameToSqlConfig) -> QuerySpecification:
+        base_query = self.__base_frame.to_sql_query_object(config)
+        new_query = copy_query(base_query)
+
+        tds_row = PandasApiTdsRow.from_tds_frame("c", self.__base_frame)
+        db_extension = config.sql_to_string_generator().get_db_extension()
+        select_items = []
+
+        for col in self.__base_frame.columns():
+            col_name = col.get_name()
+            fill_value = self.__value if not isinstance(self.__value, dict) else self.__value.get(col_name)
+            col_expr = tds_row[col_name]
+            col_sql_expr = col_expr.to_sql_expression({"c": new_query}, config)
+
+            if fill_value is not None:
+                fill_expr = convert_literal_to_literal_expression(fill_value)
+                fill_sql_expr = fill_expr.to_sql_expression({"c": new_query}, config)
+                sql_expr = FunctionCall(
+                    name=QualifiedName(parts=['coalesce']),
+                    distinct=False,
+                    arguments=[col_sql_expr, fill_sql_expr],
+                    filter_=None,
+                    window=None
+                )
+            else:
+                sql_expr = col_sql_expr  # type: ignore
+
+            select_items.append(SingleColumn(alias=db_extension.quote_identifier(col_name), expression=sql_expr))
+
+        new_query.select.selectItems = select_items  # type: ignore
+        return new_query
+
+    def to_pure(self, config: FrameToPureConfig) -> str:
+        base_pure = self.__base_frame.to_pure(config)
+        projections = []
+
+        for col in self.__base_frame.columns():
+            col_name = col.get_name()
+            fill_value = self.__value if not isinstance(self.__value, dict) else self.__value.get(col_name)
+
+            if fill_value is not None:
+                fill_expr = convert_literal_to_literal_expression(fill_value)
+                fill_pure_expr = fill_expr.to_pure_expression(config)
+                projections.append(f"'{col_name}':c|coalesce($c.{col_name}, {fill_pure_expr})")
+            else:
+                projections.append(f"'{col_name}':c|$c.{col_name}")
+
+        projection_string = ", ".join(projections)
+        return f"{base_pure}{config.separator(1)}->project(~[{projection_string}])"
+
+    def base_frame(self) -> PandasApiBaseTdsFrame:
+        return self.__base_frame
+
+    def tds_frame_parameters(self) -> PyLegendList["PandasApiBaseTdsFrame"]:
+        return []
+
+    def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
+        return [c.copy() for c in self.__base_frame.columns()]
+
+    def validate(self) -> bool:
+        if self.__value is None:
+            raise ValueError("Must specify a fill 'value'")
+
+        if not isinstance(self.__value, (int, float, str, bool, date, datetime, dict)):
+            raise TypeError(f"'value' parameter must be a scalar or dict, but you passed a {type(self.__value)}")
+        if isinstance(self.__value, dict):
+            for k, v in self.__value.items():
+                if not isinstance(k, str):
+                    raise TypeError(
+                        "All keys in 'value' dict must be strings representing column names, "
+                        f"but found key of type {type(k)}"
+                    )
+                if not isinstance(v, (int, float, str, bool, date, datetime)):
+                    raise TypeError(f"Non-scalar value of type {type(v)} passed for column '{k}' in 'value' parameter")
+
+        if self.__axis not in (0, 1, "index", "columns"):
+            raise ValueError(f"No axis named {self.__axis} for object type TdsFrame")
+        if self.__axis in (1, "columns"):
+            raise NotImplementedError("axis=1 is not supported yet in Pandas API fillna")
+
+        if self.__inplace:
+            raise NotImplementedError("inplace=True is not supported yet in Pandas API fillna")
+
+        if self.__limit is not None:
+            raise NotImplementedError("limit parameter is not supported yet in Pandas API fillna")
+
+        return True

pylegend/core/tds/pandas_api/frames/functions/filter.py

@@ -101,9 +101,13 @@ class PandasApiFilterFunction(PandasApiAppliedFunction):
             new_cols_with_index: PyLegendList[PyLegendTuple[int, SelectItem]] = []
             for col in base_query.select.selectItems:
                 if not isinstance(col, SingleColumn):
-                    raise ValueError("Select operation not supported for queries with columns other than SingleColumn")
+                    raise ValueError(
+                        "Select operation not supported for queries with columns other than SingleColumn"
+                    )  # pragma: no cover
                 if col.alias is None:
-                    raise ValueError("Select operation not supported for queries with SingleColumns with missing alias")
+                    raise ValueError(
+                        "Select operation not supported for queries with SingleColumns with missing alias"
+                    )  # pragma: no cover
                 if col.alias in columns_to_retain:
                     new_cols_with_index.append((columns_to_retain.index(col.alias), col))
 
@@ -130,10 +134,11 @@ class PandasApiFilterFunction(PandasApiAppliedFunction):
     def calculate_columns(self) -> PyLegendSequence["TdsColumn"]:
         base_cols = [c.copy() for c in self.__base_frame.columns()]
         desired_col_names = self.__get_desired_columns([c.get_name() for c in base_cols])
+        base_col_map = {c.get_name(): c for c in base_cols}
         return [
-            base_col.copy()
-            for base_col in base_cols
-            if base_col.get_name() in desired_col_names
+            base_col_map[name].copy()
+            for name in desired_col_names
+            if name in base_col_map
         ]
 
     def validate(self) -> bool: