Repository: narwhals-dev/narwhals

Issue:
[Enh]: Add `estimated_size`
### We would like to learn about your use case. For example, if this feature is needed  to adopt Narwhals in an open source project, could you please enter the link to it below?


https://github.com/mwouts/itables

### Please describe the purpose of the new feature or describe the problem to solve.

We need it for a downsampling function. Please see comment:
https://github.com/mwouts/itables/pull/339#:~:text=how%20to%20estimate%20the%20size%20of%20a%20dataframe

### Suggest a solution if possible.

_No response_

### If you have tried alternatives, please describe them below.

_No response_

### Additional information that may help us understand your needs.

_No response_

Discussion:
DeaMariaLeon: @FBruzzesi I can't get rid of a typing error. And I have a coverage error too that has nothing to do with my changes (I think). 

I haven't added the pizza yet. I would need to add it to typing, is that OK?


Patch:
diff --git a/docs/api-reference/dataframe.md b/docs/api-reference/dataframe.md
index 883fb789..c93862f2 100644
--- a/docs/api-reference/dataframe.md
+++ b/docs/api-reference/dataframe.md
@@ -11,6 +11,7 @@
         - columns
         - drop
         - drop_nulls
+        - estimated_size
         - filter
         - gather_every
         - get_column
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 68f7ab53..fe0201a9 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -21,6 +21,7 @@ from narwhals.utils import flatten
 from narwhals.utils import generate_temporary_column_name
 from narwhals.utils import is_sequence_but_not_str
 from narwhals.utils import parse_columns_to_drop
+from narwhals.utils import scale_bytes
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -35,6 +36,7 @@ if TYPE_CHECKING:
     from narwhals._arrow.series import ArrowSeries
     from narwhals._arrow.typing import IntoArrowExpr
     from narwhals.dtypes import DType
+    from narwhals.typing import SizeUnit
     from narwhals.utils import Version
 
 
@@ -285,6 +287,10 @@ class ArrowDataFrame:
     def collect_schema(self: Self) -> dict[str, DType]:
         return self.schema
 
+    def estimated_size(self: Self, unit: SizeUnit) -> int | float:
+        sz = self._native_frame.nbytes
+        return scale_bytes(sz, unit)
+
     @property
     def columns(self: Self) -> list[str]:
         return self._native_frame.schema.names  # type: ignore[no-any-return]
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 0593a141..c91e3ff4 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -372,7 +372,9 @@ class ArrowSeries:
 
     def __getitem__(self: Self, idx: int | slice | Sequence[int]) -> Any | Self:
         if isinstance(idx, int):
-            return self._native_series[idx]
+            return maybe_extract_py_scalar(
+                self._native_series[idx], return_py_scalar=True
+            )
         if isinstance(idx, Sequence):
             return self._from_native_series(self._native_series.take(idx))
         return self._from_native_series(self._native_series[idx])
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 32ba74f2..a2b739d3 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -25,6 +25,7 @@ from narwhals.utils import generate_temporary_column_name
 from narwhals.utils import import_dtypes_module
 from narwhals.utils import is_sequence_but_not_str
 from narwhals.utils import parse_columns_to_drop
+from narwhals.utils import scale_bytes
 
 if TYPE_CHECKING:
     from types import ModuleType
@@ -38,6 +39,7 @@ if TYPE_CHECKING:
     from narwhals._pandas_like.series import PandasLikeSeries
     from narwhals._pandas_like.typing import IntoPandasLikeExpr
     from narwhals.dtypes import DType
+    from narwhals.typing import SizeUnit
     from narwhals.utils import Version
 
 
@@ -371,6 +373,10 @@ class PandasLikeDataFrame:
         plx = self.__narwhals_namespace__()
         return self.filter(~plx.any_horizontal(plx.col(*subset).is_null()))
 
+    def estimated_size(self, unit: SizeUnit) -> int | float:
+        sz = self._native_frame.memory_usage(deep=True).sum()
+        return scale_bytes(sz, unit=unit)
+
     def with_row_index(self, name: str) -> Self:
         row_index = create_compliant_series(
             range(len(self._native_frame)),
diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py
index c057b722..04cd43da 100644
--- a/narwhals/dataframe.py
+++ b/narwhals/dataframe.py
@@ -36,6 +36,7 @@ if TYPE_CHECKING:
     from narwhals.typing import IntoDataFrame
     from narwhals.typing import IntoExpr
     from narwhals.typing import IntoFrame
+    from narwhals.typing import SizeUnit
     from narwhals.utils import Implementation
 
 FrameT = TypeVar("FrameT", bound="IntoFrame")
@@ -764,6 +765,50 @@ class DataFrame(BaseFrame[DataFrameT]):
             level=self._level,
         )
 
+    def estimated_size(self, unit: SizeUnit = "b") -> int | float:
+        """Return an estimation of the total (heap) allocated size of the `DataFrame`.
+
+        Estimated size is given in the specified unit (bytes by default).
+
+        Arguments:
+            unit: 'b', 'kb', 'mb', 'gb', 'tb', 'bytes', 'kilobytes', 'megabytes',
+                    'gigabytes', or 'terabytes'.
+
+        Returns:
+            Integer or Float.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> from narwhals.typing import IntoDataFrameT
+            >>> data = {
+            ...     "foo": [1, 2, 3],
+            ...     "bar": [6.0, 7.0, 8.0],
+            ...     "ham": ["a", "b", "c"],
+            ... }
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+            >>> df_pa = pa.table(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> def agnostic_estimated_size(df_native: IntoDataFrameT) -> int | float:
+            ...     df = nw.from_native(df_native)
+            ...     return df.estimated_size()
+
+            We can then pass either pandas, Polars or PyArrow to `agnostic_estimated_size`:
+
+            >>> agnostic_estimated_size(df_pd)
+            np.int64(330)
+            >>> agnostic_estimated_size(df_pl)
+            51
+            >>> agnostic_estimated_size(df_pa)
+            63
+        """
+        return self._compliant_frame.estimated_size(unit=unit)  # type: ignore[no-any-return]
+
     @overload
     def __getitem__(self, item: tuple[Sequence[int], slice]) -> Self: ...
     @overload
diff --git a/narwhals/dependencies.py b/narwhals/dependencies.py
index 463a64a6..61937f08 100644
--- a/narwhals/dependencies.py
+++ b/narwhals/dependencies.py
@@ -24,6 +24,9 @@ if TYPE_CHECKING:
     import pyarrow as pa
     import pyspark.sql as pyspark_sql
 
+    from narwhals.dataframe import DataFrame
+    from narwhals.dataframe import LazyFrame
+    from narwhals.series import Series
     from narwhals.typing import IntoSeries
 
 # We silently allow these but - given that they claim
@@ -320,6 +323,42 @@ def is_into_dataframe(native_dataframe: Any) -> bool:
     )
 
 
+def is_narwhals_dataframe(df: Any) -> TypeGuard[DataFrame[Any]]:
+    """Check whether `df` is a Narwhals DataFrame.
+
+    This is useful if you expect a user to pass in a Narwhals
+    DataFrame directly, and you want to catch both ``narwhals.DataFrame``
+    and ``narwhals.stable.v1.DataFrame`.
+    """
+    from narwhals.dataframe import DataFrame
+
+    return isinstance(df, DataFrame)
+
+
+def is_narwhals_lazyframe(lf: Any) -> TypeGuard[LazyFrame[Any]]:
+    """Check whether `lf` is a Narwhals LazyFrame.
+
+    This is useful if you expect a user to pass in a Narwhals
+    LazyFrame directly, and you want to catch both ``narwhals.LazyFrame``
+    and ``narwhals.stable.v1.LazyFrame`.
+    """
+    from narwhals.dataframe import LazyFrame
+
+    return isinstance(lf, LazyFrame)
+
+
+def is_narwhals_series(ser: Any) -> TypeGuard[Series[Any]]:
+    """Check whether `ser` is a Narwhals Series.
+
+    This is useful if you expect a user to pass in a Narwhals
+    Series directly, and you want to catch both ``narwhals.Series``
+    and ``narwhals.stable.v1.Series`.
+    """
+    from narwhals.series import Series
+
+    return isinstance(ser, Series)
+
+
 __all__ = [
     "get_cudf",
     "get_ibis",
@@ -336,6 +375,9 @@ __all__ = [
     "is_into_series",
     "is_modin_dataframe",
     "is_modin_series",
+    "is_narwhals_dataframe",
+    "is_narwhals_lazyframe",
+    "is_narwhals_series",
     "is_numpy_array",
     "is_pandas_dataframe",
     "is_pandas_index",
diff --git a/narwhals/dtypes.py b/narwhals/dtypes.py
index b0434125..b086d216 100644
--- a/narwhals/dtypes.py
+++ b/narwhals/dtypes.py
@@ -386,6 +386,33 @@ class Datetime(TemporalType):
 
     Notes:
         Adapted from [Polars implementation](https://github.com/pola-rs/polars/blob/py-1.7.1/py-polars/polars/datatypes/classes.py#L398-L457)
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import pyarrow as pa
+        >>> import pyarrow.compute as pc
+        >>> import narwhals as nw
+        >>> from datetime import datetime, timedelta
+        >>> data = [datetime(2024, 12, 9) + timedelta(days=n) for n in range(5)]
+        >>> ser_pd = (
+        ...     pd.Series(data)
+        ...     .dt.tz_localize("Africa/Accra")
+        ...     .astype("datetime64[ms, Africa/Accra]")
+        ... )
+        >>> ser_pl = (
+        ...     pl.Series(data).cast(pl.Datetime("ms")).dt.replace_time_zone("Africa/Accra")
+        ... )
+        >>> ser_pa = pc.assume_timezone(
+        ...     pa.chunked_array([data], type=pa.timestamp("ms")), "Africa/Accra"
+        ... )
+
+        >>> nw.from_native(ser_pd, series_only=True).dtype
+        Datetime(time_unit='ms', time_zone='Africa/Accra')
+        >>> nw.from_native(ser_pl, series_only=True).dtype
+        Datetime(time_unit='ms', time_zone='Africa/Accra')
+        >>> nw.from_native(ser_pa, series_only=True).dtype
+        Datetime(time_unit='ms', time_zone='Africa/Accra')
     """
 
     def __init__(
diff --git a/narwhals/stable/v1/_dtypes.py b/narwhals/stable/v1/_dtypes.py
index 22cf05db..41dc854f 100644
--- a/narwhals/stable/v1/_dtypes.py
+++ b/narwhals/stable/v1/_dtypes.py
@@ -37,6 +37,33 @@ class Datetime(NwDatetime):
 
     Notes:
         Adapted from [Polars implementation](https://github.com/pola-rs/polars/blob/py-1.7.1/py-polars/polars/datatypes/classes.py#L398-L457)
+
+    Examples:
+        >>> import pandas as pd
+        >>> import polars as pl
+        >>> import pyarrow as pa
+        >>> import pyarrow.compute as pc
+        >>> import narwhals as nw
+        >>> from datetime import datetime, timedelta
+        >>> data = [datetime(2024, 12, 9) + timedelta(days=n) for n in range(5)]
+        >>> ser_pd = (
+        ...     pd.Series(data)
+        ...     .dt.tz_localize("Africa/Accra")
+        ...     .astype("datetime64[ms, Africa/Accra]")
+        ... )
+        >>> ser_pl = (
+        ...     pl.Series(data).cast(pl.Datetime("ms")).dt.replace_time_zone("Africa/Accra")
+        ... )
+        >>> ser_pa = pc.assume_timezone(
+        ...     pa.chunked_array([data], type=pa.timestamp("ms")), "Africa/Accra"
+        ... )
+
+        >>> nw.from_native(ser_pd, series_only=True).dtype
+        Datetime(time_unit='ms', time_zone='Africa/Accra')
+        >>> nw.from_native(ser_pl, series_only=True).dtype
+        Datetime(time_unit='ms', time_zone='Africa/Accra')
+        >>> nw.from_native(ser_pa, series_only=True).dtype
+        Datetime(time_unit='ms', time_zone='Africa/Accra')
     """
 
     def __hash__(self) -> int:
diff --git a/narwhals/stable/v1/dependencies.py b/narwhals/stable/v1/dependencies.py
index 6a020622..8ad1ad20 100644
--- a/narwhals/stable/v1/dependencies.py
+++ b/narwhals/stable/v1/dependencies.py
@@ -15,6 +15,9 @@ from narwhals.dependencies import is_into_dataframe
 from narwhals.dependencies import is_into_series
 from narwhals.dependencies import is_modin_dataframe
 from narwhals.dependencies import is_modin_series
+from narwhals.dependencies import is_narwhals_dataframe
+from narwhals.dependencies import is_narwhals_lazyframe
+from narwhals.dependencies import is_narwhals_series
 from narwhals.dependencies import is_numpy_array
 from narwhals.dependencies import is_pandas_dataframe
 from narwhals.dependencies import is_pandas_index
@@ -43,6 +46,9 @@ __all__ = [
     "is_into_series",
     "is_modin_dataframe",
     "is_modin_series",
+    "is_narwhals_dataframe",
+    "is_narwhals_lazyframe",
+    "is_narwhals_series",
     "is_numpy_array",
     "is_pandas_dataframe",
     "is_pandas_index",
diff --git a/narwhals/typing.py b/narwhals/typing.py
index 8121c33d..329512de 100644
--- a/narwhals/typing.py
+++ b/narwhals/typing.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Literal
 from typing import Protocol
 from typing import TypeVar
 from typing import Union
@@ -173,6 +174,19 @@ Examples:
     ...     return s.abs().to_native()
 """
 
+SizeUnit: TypeAlias = Literal[
+    "b",
+    "kb",
+    "mb",
+    "gb",
+    "tb",
+    "bytes",
+    "kilobytes",
+    "megabytes",
+    "gigabytes",
+    "terabytes",
+]
+
 
 class DTypes:
     Int64: type[dtypes.Int64]
diff --git a/narwhals/utils.py b/narwhals/utils.py
index f69eb566..58f56cdf 100644
--- a/narwhals/utils.py
+++ b/narwhals/utils.py
@@ -43,6 +43,7 @@ if TYPE_CHECKING:
     from narwhals.series import Series
     from narwhals.typing import DTypes
     from narwhals.typing import IntoSeriesT
+    from narwhals.typing import SizeUnit
 
     FrameOrSeriesT = TypeVar(
         "FrameOrSeriesT", bound=Union[LazyFrame[Any], DataFrame[Any], Series[Any]]
@@ -681,6 +682,31 @@ def maybe_convert_dtypes(
     return obj_any  # type: ignore[no-any-return]
 
 
+def scale_bytes(sz: int, unit: SizeUnit) -> int | float:
+    """Scale size in bytes to other size units (eg: "kb", "mb", "gb", "tb").
+
+    Arguments:
+        sz: original size in bytes
+        unit: size unit to convert into
+
+    Returns:
+        Integer or float.
+    """
+    if unit in {"b", "bytes"}:
+        return sz
+    elif unit in {"kb", "kilobytes"}:
+        return sz / 1024
+    elif unit in {"mb", "megabytes"}:
+        return sz / 1024**2
+    elif unit in {"gb", "gigabytes"}:
+        return sz / 1024**3
+    elif unit in {"tb", "terabytes"}:
+        return sz / 1024**4
+    else:
+        msg = f"`unit` must be one of {{'b', 'kb', 'mb', 'gb', 'tb'}}, got {unit!r}"
+        raise ValueError(msg)
+
+
 def is_ordered_categorical(series: Series[Any]) -> bool:
     """Return whether indices of categories are semantically meaningful.
 


Test patch:
diff --git a/.github/workflows/downstream_tests.yml b/.github/workflows/downstream_tests.yml
index ec167034..046947a5 100644
--- a/.github/workflows/downstream_tests.yml
+++ b/.github/workflows/downstream_tests.yml
@@ -35,6 +35,8 @@ jobs:
         run: |
             cd altair
             uv pip install -e ".[dev, all]" --system
+            # temporary pin to get ci green
+            uv pip install "numpy<2.2" --system
       - name: install-narwhals-dev
         run: |
             uv pip uninstall narwhals --system
@@ -140,6 +142,8 @@ jobs:
         run: |
             cd scikit-lego
             uv pip install -e ".[test]" --system
+            # temporary pin to get ci green
+            uv pip install "scikit-learn<1.6" --system
       - name: install-narwhals-dev
         run: |
             uv pip uninstall narwhals --system
@@ -184,8 +188,6 @@ jobs:
             make narwhals-install-shiny
       - name: install-narwhals-dev
         run: |
-            # temporary pin to get CI green
-            uv pip install "ollama < 0.4.2" --system
             uv pip uninstall narwhals --system
             uv pip install -e . --system
       - name: show-deps
diff --git a/tests/dependencies/is_narwhals_dataframe_test.py b/tests/dependencies/is_narwhals_dataframe_test.py
new file mode 100644
index 00000000..0a25bc8b
--- /dev/null
+++ b/tests/dependencies/is_narwhals_dataframe_test.py
@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import narwhals as nw
+import narwhals.stable.v1 as nws
+from narwhals.stable.v1.dependencies import is_narwhals_dataframe
+
+if TYPE_CHECKING:
+    from tests.utils import ConstructorEager
+
+
+def test_is_narwhals_dataframe(constructor_eager: ConstructorEager) -> None:
+    df = constructor_eager({"col1": [1, 2], "col2": [3, 4]})
+
+    assert is_narwhals_dataframe(nw.from_native(df))
+    assert is_narwhals_dataframe(nws.from_native(df))
+    assert not is_narwhals_dataframe(df)
diff --git a/tests/dependencies/is_narwhals_lazyframe_test.py b/tests/dependencies/is_narwhals_lazyframe_test.py
new file mode 100644
index 00000000..1dcbe5fc
--- /dev/null
+++ b/tests/dependencies/is_narwhals_lazyframe_test.py
@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import narwhals as nw
+import narwhals.stable.v1 as nws
+from narwhals.stable.v1.dependencies import is_narwhals_lazyframe
+from tests.utils import Constructor
+
+if TYPE_CHECKING:
+    from tests.utils import Constructor
+
+
+def test_is_narwhals_lazyframe(constructor: Constructor) -> None:
+    lf = constructor({"a": [1, 2, 3]})
+
+    assert is_narwhals_lazyframe(nw.from_native(lf).lazy())
+    assert is_narwhals_lazyframe(nws.from_native(lf).lazy())
+    assert not is_narwhals_lazyframe(lf)
diff --git a/tests/dependencies/is_narwhals_series_test.py b/tests/dependencies/is_narwhals_series_test.py
new file mode 100644
index 00000000..0beb3fc1
--- /dev/null
+++ b/tests/dependencies/is_narwhals_series_test.py
@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import narwhals as nw
+import narwhals.stable.v1 as nws
+from narwhals.stable.v1.dependencies import is_narwhals_series
+
+if TYPE_CHECKING:
+    from tests.utils import ConstructorEager
+
+
+def test_is_narwhals_series(constructor_eager: ConstructorEager) -> None:
+    df = constructor_eager({"col1": [1, 2], "col2": [3, 4]})
+
+    assert is_narwhals_series(nw.from_native(df, eager_only=True)["col1"])
+    assert is_narwhals_series(nws.from_native(df, eager_only=True)["col1"])
+    assert not is_narwhals_series(nw.from_native(df, eager_only=True)["col1"].to_native())
diff --git a/tests/frame/estimated_size_test.py b/tests/frame/estimated_size_test.py
new file mode 100644
index 00000000..499a154f
--- /dev/null
+++ b/tests/frame/estimated_size_test.py
@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pytest
+
+import narwhals.stable.v1 as nw
+
+if TYPE_CHECKING:
+    from tests.utils import ConstructorEager
+
+data = {"a": list(range(100))}
+
+
+def test_estimated_size(constructor_eager: ConstructorEager) -> None:
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+
+    assert df.estimated_size("b") > 0
+    assert df.estimated_size("kb") == (df.estimated_size("b") / 1024)
+    assert df.estimated_size("mb") == (df.estimated_size("kb") / 1024)
+    assert df.estimated_size("gb") == (df.estimated_size("mb") / 1024)
+    assert df.estimated_size("tb") == (df.estimated_size("gb") / 1024)
+
+    with pytest.raises(
+        ValueError,
+        match="`unit` must be one of {'b', 'kb', 'mb', 'gb', 'tb'}, got 'pizza'",
+    ):
+        df.estimated_size("pizza")  # type: ignore[arg-type]
diff --git a/tests/series_only/slice_test.py b/tests/series_only/slice_test.py
index 5f8a0a0a..1eb5e0ba 100644
--- a/tests/series_only/slice_test.py
+++ b/tests/series_only/slice_test.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import pyarrow as pa
+
 import narwhals.stable.v1 as nw
 from tests.utils import ConstructorEager
 from tests.utils import assert_equal_data
@@ -32,3 +34,8 @@ def test_slice(constructor_eager: ConstructorEager) -> None:
     result = {"b": df[[], 1]}
     expected = {"b": []}
     assert_equal_data(result, expected)
+
+
+def test_getitem_arrow_scalar() -> None:
+    result = nw.from_native(pa.chunked_array([[1]]), series_only=True)[0]
+    assert isinstance(result, int)