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

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: