Merge pull request #13 from data-apis/remove-column-add-expression

Separate lazy and eager

Author: MarcoGorelli

Diff for: .github/workflows/tox.yml

@@ -9,7 +9,7 @@ jobs:
   tox:
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11"]
         os: [windows-latest, ubuntu-latest]
 
     runs-on: ${{ matrix.os }}

Diff for: dataframe_api_compat/pandas_standard/__init__.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import re
 from typing import Any
+from typing import Literal
 from typing import TYPE_CHECKING
 
 import pandas as pd
@@ -10,12 +12,24 @@
 from dataframe_api_compat.pandas_standard.pandas_standard import PandasColumn
 from dataframe_api_compat.pandas_standard.pandas_standard import PandasDataFrame
 from dataframe_api_compat.pandas_standard.pandas_standard import PandasGroupBy
+from dataframe_api_compat.pandas_standard.pandas_standard import PandasPermissiveColumn
+from dataframe_api_compat.pandas_standard.pandas_standard import PandasPermissiveFrame
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from dataframe_api._types import DType
+
+
+def col(name: str) -> PandasColumn:
+    return PandasColumn(
+        root_names=[name], output_name=name, base_call=lambda df: df.loc[:, name]
+    )
+
 
 Column = PandasColumn
+PermissiveColumn = PandasPermissiveColumn
 DataFrame = PandasDataFrame
+PermissiveFrame = PandasPermissiveFrame
 GroupBy = PandasGroupBy
 
 
@@ -67,35 +81,82 @@ class String:
     ...
 
 
-DTYPE_MAP = {
-    "int64": Int64(),
-    "Int64": Int64(),
-    "int32": Int32(),
-    "Int32": Int32(),
-    "int16": Int16(),
-    "Int16": Int16(),
-    "int8": Int8(),
-    "Int8": Int8(),
-    "uint64": UInt64(),
-    "UInt64": UInt64(),
-    "uint32": UInt32(),
-    "UInt32": UInt32(),
-    "uint16": UInt16(),
-    "UInt16": UInt16(),
-    "uint8": UInt8(),
-    "UInt8": UInt8(),
-    "float64": Float64(),
-    "Float64": Float64(),
-    "float32": Float32(),
-    "Float32": Float32(),
-    "bool": Bool(),
-    "boolean": Bool(),
-    "object": String(),
-    "string": String(),
-}
-
-
-def map_standard_dtype_to_pandas_dtype(dtype: Any) -> Any:
+class Date:
+    ...
+
+
+class Datetime:
+    def __init__(self, time_unit, time_zone=None):
+        self.time_unit = time_unit
+        # todo validate time zone
+        self.time_zone = time_zone
+
+
+class Duration:
+    def __init__(self, time_unit):
+        self.time_unit = time_unit
+
+
+def map_pandas_dtype_to_standard_dtype(dtype: Any) -> DType:
+    if dtype == "int64":
+        return Int64()
+    if dtype == "Int64":
+        return Int64()
+    if dtype == "int32":
+        return Int32()
+    if dtype == "Int32":
+        return Int32()
+    if dtype == "int16":
+        return Int16()
+    if dtype == "Int16":
+        return Int16()
+    if dtype == "int8":
+        return Int8()
+    if dtype == "Int8":
+        return Int8()
+    if dtype == "uint64":
+        return UInt64()
+    if dtype == "UInt64":
+        return UInt64()
+    if dtype == "uint32":
+        return UInt32()
+    if dtype == "UInt32":
+        return UInt32()
+    if dtype == "uint16":
+        return UInt16()
+    if dtype == "UInt16":
+        return UInt16()
+    if dtype == "uint8":
+        return UInt8()
+    if dtype == "UInt8":
+        return UInt8()
+    if dtype == "float64":
+        return Float64()
+    if dtype == "Float64":
+        return Float64()
+    if dtype == "float32":
+        return Float32()
+    if dtype == "Float32":
+        return Float32()
+    if dtype == "bool":
+        # 'boolean' not yet covered, as the default dtype in pandas is still 'bool'
+        return Bool()
+    if dtype == "object":
+        return String()
+    if dtype == "string":
+        return String()
+    if dtype == "datetime64[s]":
+        return Date()
+    if dtype.startswith("datetime64["):
+        time_unit = re.search(r"datetime64\[(\w{1,2})", dtype).group(1)
+        return Datetime(time_unit)
+    if dtype.startswith("timedelta64["):
+        time_unit = re.search(r"timedelta64\[(\w{1,2})", dtype).group(1)
+        return Duration(time_unit)
+    raise AssertionError(f"Unsupported dtype! {dtype}")
+
+
+def map_standard_dtype_to_pandas_dtype(dtype: DType) -> Any:
     if isinstance(dtype, Int64):
         return "int64"
     if isinstance(dtype, Int32):
@@ -120,9 +181,26 @@ def map_standard_dtype_to_pandas_dtype(dtype: Any) -> Any:
         return "bool"
     if isinstance(dtype, String):
         return "object"
+    if isinstance(dtype, Datetime):
+        if dtype.time_zone is not None:  # pragma: no cover (todo)
+            return f"datetime64[{dtype.time_unit}, {dtype.time_zone}]"
+        return f"datetime64[{dtype.time_unit}]"
+    if isinstance(dtype, Duration):
+        return f"timedelta64[{dtype.time_unit}]"
     raise AssertionError(f"Unknown dtype: {dtype}")
 
 
+def convert_to_standard_compliant_column(
+    ser: pd.Series, api_version: str | None = None
+) -> PandasDataFrame:
+    if api_version is None:  # pragma: no cover
+        api_version = LATEST_API_VERSION
+    if ser.name is not None and not isinstance(ser.name, str):
+        raise ValueError(f"Expected column with string name, got: {ser.name}")
+    name = ser.name or ""
+    return PandasPermissiveColumn(ser.rename(name), api_version=api_version)
+
+
 def convert_to_standard_compliant_dataframe(
     df: pd.DataFrame, api_version: str | None = None
 ) -> PandasDataFrame:
@@ -131,13 +209,6 @@ def convert_to_standard_compliant_dataframe(
     return PandasDataFrame(df, api_version=api_version)
 
 
-def convert_to_standard_compliant_column(
-    df: pd.Series[Any],
-    api_version: str | None = None,
-) -> PandasColumn[Any]:
-    return PandasColumn(df, api_version=api_version or LATEST_API_VERSION)
-
-
 def concat(dataframes: Sequence[PandasDataFrame]) -> PandasDataFrame:
     dtypes = dataframes[0].dataframe.dtypes
     dfs = []
@@ -164,16 +235,30 @@ def concat(dataframes: Sequence[PandasDataFrame]) -> PandasDataFrame:
 
 def column_from_sequence(
     sequence: Sequence[Any], *, dtype: Any, name: str, api_version: str | None = None
-) -> PandasColumn[Any]:
+) -> PandasPermissiveColumn[Any]:
     ser = pd.Series(sequence, dtype=map_standard_dtype_to_pandas_dtype(dtype), name=name)
-    return PandasColumn(ser, api_version=LATEST_API_VERSION)
+    return PandasPermissiveColumn(ser, api_version=api_version or LATEST_API_VERSION)
+
+
+def dataframe_from_dict(
+    data: dict[str, PandasPermissiveColumn[Any]], api_version: str | None = None
+) -> PandasDataFrame:
+    for _, col in data.items():
+        if not isinstance(col, PandasPermissiveColumn):  # pragma: no cover
+            raise TypeError(f"Expected PandasPermissiveColumn, got {type(col)}")
+    return PandasDataFrame(
+        pd.DataFrame(
+            {label: column.column.rename(label) for label, column in data.items()}
+        ),
+        api_version=api_version or LATEST_API_VERSION,
+    )
 
 
 def column_from_1d_array(
     data: Any, *, dtype: Any, name: str | None = None, api_version: str | None = None
-) -> PandasColumn[Any]:  # pragma: no cover
+) -> PandasPermissiveColumn[Any]:  # pragma: no cover
     ser = pd.Series(data, dtype=map_standard_dtype_to_pandas_dtype(dtype), name=name)
-    return PandasColumn(ser, api_version=api_version or LATEST_API_VERSION)
+    return PandasPermissiveColumn(ser, api_version=api_version or LATEST_API_VERSION)
 
 
 def dataframe_from_2d_array(
@@ -189,20 +274,6 @@ def dataframe_from_2d_array(
     return PandasDataFrame(df, api_version=api_version or LATEST_API_VERSION)
 
 
-def dataframe_from_dict(
-    data: dict[str, PandasColumn[Any]], api_version: str | None = None
-) -> PandasDataFrame:
-    for _, col in data.items():
-        if not isinstance(col, PandasColumn):  # pragma: no cover
-            raise TypeError(f"Expected PandasColumn, got {type(col)}")
-    return PandasDataFrame(
-        pd.DataFrame(
-            {label: column.column.rename(label) for label, column in data.items()}
-        ),
-        api_version=api_version or LATEST_API_VERSION,
-    )
-
-
 def is_null(value: Any) -> bool:
     return value is null
 
@@ -223,3 +294,47 @@ def is_dtype(dtype: Any, kind: str | tuple[str, ...]) -> bool:
         if _kind == "string":
             dtypes.add(String)
     return isinstance(dtype, tuple(dtypes))
+
+
+def any_rowwise(*columns: str, skip_nulls: bool = True) -> PandasColumn:
+    # todo: accept expressions
+    def func(df):
+        return df.loc[:, list(columns) or df.columns.tolist()].any(axis=1)
+
+    return PandasColumn(root_names=list(columns), output_name="any", base_call=func)
+
+
+def all_rowwise(*columns: str, skip_nulls: bool = True) -> PandasColumn:
+    def func(df: pd.DataFrame) -> pd.Series:
+        return df.loc[:, list(columns) or df.columns.tolist()].all(axis=1)
+
+    return PandasColumn(root_names=list(columns), output_name="all", base_call=func)
+
+
+def sorted_indices(
+    *keys: str,
+    ascending: Sequence[bool] | bool = True,
+    nulls_position: Literal["first", "last"] = "last",
+) -> Column:
+    def func(df: pd.DataFrame) -> pd.Series:
+        if ascending:
+            return (
+                df.loc[:, list(keys)]
+                .sort_values(list(keys))
+                .index.to_series()
+                .reset_index(drop=True)
+            )
+        return (
+            df.loc[:, list(keys)]
+            .sort_values(list(keys))
+            .index.to_series()[::-1]
+            .reset_index(drop=True)
+        )
+
+    return PandasColumn(root_names=list(keys), output_name="indices", base_call=func)
+
+
+def unique_indices(
+    keys: str | list[str] | None = None, *, skip_nulls: bool = True
+) -> Column:
+    raise NotImplementedError("namespace.unique_indices not implemented for pandas yet")