ergonomics

Author: MarcoGorelli

spec/API_specification/dataframe_api/__init__.py

@@ -6,13 +6,18 @@
 from typing import Mapping, Sequence, Any, Literal
 
 from .eagercolumn_object import EagerColumn
+from .eagerframe_object import EagerFrame
 from .expression_object import Expression
 from .dataframe_object import DataFrame
+from .groupby_object import GroupBy
 
 __all__ = [
     "__dataframe_api_version__",
     "DataFrame",
+    "EagerFrame",
     "EagerColumn",
+    "Expression",
+    "GroupBy",
     "column_from_sequence",
     "column_from_1d_array",
     "col",
@@ -35,6 +40,8 @@
     "Float32",
     "Bool",
     "is_dtype",
+    "any_rowwise",
+    "all_rowwise",
 ]
 
 
@@ -202,7 +209,7 @@ def dataframe_from_2d_array(array: Any, *, names: Sequence[str], dtypes: Mapping
     """
     ...
 
-def any_rowwise(keys: list[str] | None = None, *, skip_nulls: bool = True) -> Expression:
+def any_rowwise(*keys: str, skip_nulls: bool = True) -> Expression:
     """
     Reduction returns an Expression.
 
@@ -211,8 +218,8 @@ def any_rowwise(keys: list[str] | None = None, *, skip_nulls: bool = True) -> Ex
 
     Parameters
     ----------
-    keys : list[str]
-        Column names to consider. If `None`, all columns are considered.
+    keys : str
+        Column names to consider.
 
     Raises
     ------
@@ -221,7 +228,7 @@ def any_rowwise(keys: list[str] | None = None, *, skip_nulls: bool = True) -> Ex
     """
     ...
 
-def all_rowwise(keys: list[str] | None = None, *, skip_nulls: bool = True) -> Expression:
+def all_rowwise(*keys: str, skip_nulls: bool = True) -> Expression:
     """
     Reduction returns an Expression.
 
@@ -230,8 +237,8 @@ def all_rowwise(keys: list[str] | None = None, *, skip_nulls: bool = True) -> Ex
 
     Parameters
     ----------
-    keys : list[str]
-        Column names to consider. If `None`, all columns are considered.
+    keys : str
+        Column names to consider.
 
     Raises
     ------
@@ -241,8 +248,7 @@ def all_rowwise(keys: list[str] | None = None, *, skip_nulls: bool = True) -> Ex
     ...
 
 def sorted_indices(
-    keys: str | list[str] | None = None,
-    *,
+    *keys: str,
     ascending: Sequence[bool] | bool = True,
     nulls_position: Literal['first', 'last'] = 'last',
 ) -> Expression:
@@ -253,9 +259,8 @@ def sorted_indices(
 
     Parameters
     ----------
-    keys : str | list[str], optional
+    keys : str
         Names of columns to sort by.
-        If `None`, sort by all columns.
     ascending : Sequence[bool] or bool
         If `True`, sort by all keys in ascending order.
         If `False`, sort by all keys in descending order.
@@ -280,15 +285,14 @@ def sorted_indices(
     ...
 
 
-def unique_indices(keys: str | list[str] | None = None, *, skip_nulls: bool = True) -> Expression:
+def unique_indices(*keys: str, skip_nulls: bool = True) -> Expression:
     """
     Return indices corresponding to unique values across selected columns.
 
     Parameters
     ----------
-    keys : str | list[str], optional
+    keys : str
         Column names to consider when finding unique values.
-        If `None`, all columns are considered.
 
     Returns
     -------

spec/API_specification/dataframe_api/_types.py

@@ -17,9 +17,17 @@
     TypeVar,
     Union,
     Protocol,
+    TYPE_CHECKING,
+    TypeAlias
 )
 from enum import Enum
 
+if TYPE_CHECKING:
+    from .expression_object import Expression
+    from .eagercolumn_object import EagerColumn
+
+    IntoExpression: TypeAlias = Expression | EagerColumn
+
 # Type alias: Mypy needs Any, but for readability we need to make clear this
 # is a Python scalar (i.e., an instance of `bool`, `int`, `float`, `str`, etc.)
 Scalar = Any

spec/API_specification/dataframe_api/dataframe_object.py

@@ -8,12 +8,11 @@
     from .eagerframe_object import EagerFrame
     from .eagercolumn_object import EagerColumn
     from .groupby_object import GroupBy
-    from ._types import NullType, Scalar
+    from ._types import NullType, Scalar, IntoExpression
 
 
 __all__ = ["DataFrame"]
 
-IntoExpression = EagerColumn[Any] | Expression
 
 
 class DataFrame:
@@ -68,13 +67,13 @@ def dataframe(self) -> object:
         """
         ...
 
-    def groupby(self, keys: str | list[str], /) -> GroupBy:
+    def groupby(self, *keys: str) -> GroupBy:
         """
         Group the DataFrame by the given columns.
 
         Parameters
         ----------
-        keys : str | list[str]
+        keys : str
 
         Returns
         -------
@@ -93,7 +92,7 @@ def groupby(self, keys: str | list[str], /) -> GroupBy:
         """
         ...
 
-    def select(self, names: str | Expression | Sequence[str | Expression], /) -> DataFrame:
+    def select(self, *names: str | Expression) -> DataFrame:
         """
         Select multiple columns, either by name or by expressions.
 
@@ -107,16 +106,9 @@ def select(self, names: str | Expression | Sequence[str | Expression], /) -> Dat
 
         Examples
         --------
-        Select columns 'a' and 'b':
-
-        >>> df: DataFrame
-        >>> df.select(['a', 'b'])
-
-        You can also pass expressions:
-
         >>> df: DataFrame
-        >>> namespace = df.__dataframe_namespace__()
-        >>> df.select(['a', namespace.col('b')+1])
+        >>> col = df.__dataframe_namespace__().col
+        >>> df = df.select('a', col('b'), (col('c')+col('d')+1).rename('e'))
 
         Raises
         ------
@@ -145,13 +137,13 @@ def slice_rows(
         """
         ...
 
-    def filter(self, mask: Expression) -> DataFrame:
+    def filter(self, mask: IntoExpression) -> DataFrame:
         """
         Select a subset of rows corresponding to a mask.
 
         Parameters
         ----------
-        mask : Expression
+        mask : Expression or EagerColumn
 
         Returns
         -------
@@ -170,7 +162,7 @@ def filter(self, mask: Expression) -> DataFrame:
         """
         ...
 
-    def insert_columns(self, columns: IntoExpression | Sequence[IntoExpression]) -> DataFrame:
+    def insert_columns(self, *columns: Expression | EagerColumn[Any]) -> DataFrame:
         """
         Insert column into DataFrame at rightmost location.
 
@@ -184,7 +176,7 @@ def insert_columns(self, columns: IntoExpression | Sequence[IntoExpression]) ->
             namespace = df.__dataframe_namespace__()
             col = namespace.col
             new_column = namespace.col('a') + 1
-            df = df.insert_column(new_column.rename('a_plus_1'))
+            df = df.insert_columns(new_column.rename('a_plus_1'))
         
         If you need to insert the column at a different location, combine with
         :meth:`select`, e.g.:
@@ -196,7 +188,7 @@ def insert_columns(self, columns: IntoExpression | Sequence[IntoExpression]) ->
             col = namespace.col
             new_column = namespace.col('a') + 1
             new_columns_names = ['a_plus_1'] + df.get_column_names()
-            df = df.insert_column(new_column.rename('a_plus_1'))
+            df = df.insert_columns(new_column.rename('a_plus_1'))
             df = df.select(new_column_names)
 
         Parameters
@@ -206,7 +198,7 @@ def insert_columns(self, columns: IntoExpression | Sequence[IntoExpression]) ->
         """
         ...
 
-    def update_columns(self, columns: Expression | EagerColumn | Sequence[Expression | EagerColumn], /) -> DataFrame:
+    def update_columns(self, *columns: Expression | EagerColumn[Any]) -> DataFrame:
         """
         Update values in existing column(s) from Dataframe.
 
@@ -224,7 +216,7 @@ def update_columns(self, columns: Expression | EagerColumn | Sequence[Expression
 
         Parameters
         ----------
-        columns : Expression | Sequence[Expression]
+        columns : Expression, EagerColumn, or sequence of either
             Column(s) to update. If updating multiple columns, they must all have
             different names.
 
@@ -268,7 +260,8 @@ def rename_columns(self, mapping: Mapping[str, str]) -> DataFrame:
         """
         ...
 
-    def get_column_names(self) -> list[str]:
+    @property
+    def column_names(self) -> list[str]:
         """
         Get column names.
 
@@ -280,8 +273,7 @@ def get_column_names(self) -> list[str]:
 
     def sort(
         self,
-        keys: str | Expression | list[str | Expression] | None = None,
-        *,
+        *keys: str | Expression,
         ascending: Sequence[bool] | bool = True,
         nulls_position: Literal['first', 'last'] = 'last',
     ) -> DataFrame:
@@ -293,9 +285,9 @@ def sort(
 
         Parameters
         ----------
-        keys : str | list[str], optional
+        keys : str | Expression
             Names of columns to sort by.
-            If `None`, sort by all columns.
+            If not passed, will sort by all columns.
         ascending : Sequence[bool] or bool
             If `True`, sort by all keys in ascending order.
             If `False`, sort by all keys in descending order.
@@ -759,38 +751,6 @@ def fill_nan(self, value: float | NullType, /) -> DataFrame:
         """
         ...
 
-    def fill_null(
-        self, value: Scalar, /, *, column_names : list[str] | None = None
-    ) -> DataFrame:
-        """
-        Fill null values with the given fill value.
-
-        This method can only be used if all columns that are to be filled are
-        of the same dtype (e.g., all of ``Float64`` or all of string dtype).
-        If that is not the case, it is not possible to use a single Python
-        scalar type that matches the dtype of all columns to which
-        ``fill_null`` is being applied, and hence an exception will be raised.
-
-        Parameters
-        ----------
-        value : Scalar
-            Value used to replace any ``null`` values in the dataframe with.
-            Must be of the Python scalar type matching the dtype(s) of the dataframe.
-        column_names : list[str] | None
-            A list of column names for which to replace nulls with the given
-            scalar value. If ``None``, nulls will be replaced in all columns.
-
-        Raises
-        ------
-        TypeError
-            If the columns of the dataframe are not all of the same kind.
-        KeyError
-            If ``column_names`` contains a column name that is not present in
-            the dataframe.
-
-        """
-        ...
-    
     def join(
         self,
         other: DataFrame,