fix: table pagination for next page (#2790)

Regression when moving to narwhals, pagination offset logic was
incorrect. This fixes it and adds a test

marimo/_output/data/data.py

@@ -1,4 +1,6 @@
 # Copyright 2024 Marimo. All rights reserved.
+from __future__ import annotations
+
 import base64
 import io
 from typing import Union
@@ -171,3 +173,13 @@ def any_data(data: Union[str, bytes, io.BytesIO], ext: str) -> VirtualFile:
         return item.virtual_file
 
     raise ValueError(f"Unsupported data type: {type(data)}")
+
+
+# Format: data:mime_type;base64,data
+def from_data_uri(data: str) -> tuple[str, bytes]:
+    assert isinstance(data, str)
+    assert data.startswith("data:")
+    mime_type, data = data.split(",", 1)
+    # strip data: and ;base64
+    mime_type = mime_type.split(";")[0][5:]
+    return mime_type, base64.b64decode(data)

marimo/_plugins/ui/_impl/tables/narwhals_table.py

@@ -127,7 +127,7 @@ def take(self, count: int, offset: int) -> TableManager[Any]:
             raise ValueError("Count must be a positive integer")
         if offset < 0:
             raise ValueError("Offset must be a non-negative integer")
-        return self.with_new_data(self.data[offset:count])
+        return self.with_new_data(self.data[offset : offset + count])
 
     def search(self, query: str) -> TableManager[Any]:
         query = query.lower()

tests/_plugins/ui/_impl/tables/test_ibis_table.py

@@ -137,6 +137,44 @@ def test_limit(self) -> None:
             expected_data.to_pandas()
         )
 
+    def test_take(self) -> None:
+        assert (
+            self.manager.take(1, 0).select_columns(["A"]).to_json()
+            == b'[{"A":1}]'
+        )
+        assert (
+            self.manager.take(2, 0).select_columns(["A"]).to_json()
+            == b'[{"A":1},{"A":2}]'
+        )
+        assert (
+            self.manager.take(2, 1).select_columns(["A"]).to_json()
+            == b'[{"A":2},{"A":3}]'
+        )
+        assert (
+            self.manager.take(2, 2).select_columns(["A"]).to_json()
+            == b'[{"A":3}]'
+        )
+
+    def test_take_zero(self) -> None:
+        limited_manager = self.manager.take(0, 0)
+        assert limited_manager.data.count().execute() == 0
+
+    def test_take_negative(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(-1, 0)
+
+    def test_take_negative_offset(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(1, -1)
+
+    def test_take_out_of_bounds(self) -> None:
+        # Too large of page
+        assert self.manager.take(10, 0).data.count().execute() == 3
+        assert self.data.count().execute() == 3
+
+        # Too large of page and offset
+        assert self.manager.take(10, 10).data.count().execute() == 0
+
     def test_summary_integer(self) -> None:
         column = "A"
         summary = self.manager.get_summary(column)

tests/_plugins/ui/_impl/tables/test_narwhals.py

@@ -240,6 +240,24 @@ def test_limit(self) -> None:
         expected_data = self.data.head(1)
         assert_frame_equal(limited_manager.data, expected_data)
 
+    def test_take(self) -> None:
+        assert self.manager.take(1, 0).data["A"].to_list() == [1]
+        assert self.manager.take(2, 0).data["A"].to_list() == [1, 2]
+        assert self.manager.take(2, 1).data["A"].to_list() == [2, 3]
+        assert self.manager.take(2, 2).data["A"].to_list() == [3]
+
+    def test_take_zero(self) -> None:
+        limited_manager = self.manager.take(0, 0)
+        assert limited_manager.data.is_empty()
+
+    def test_take_negative(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(-1, 0)
+
+    def test_take_negative_offset(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(1, -1)
+
     def test_take_out_of_bounds(self) -> None:
         # Too large of page
         assert len(self.manager.take(10, 0).data) == 3

tests/_plugins/ui/_impl/tables/test_pandas_table.py

@@ -318,6 +318,24 @@ def test_limit(self) -> None:
         expected_data = self.data.head(limit)
         assert_frame_equal(limited_manager.data, expected_data)
 
+    def test_take(self) -> None:
+        assert self.manager.take(1, 0).data["A"].to_list() == [1]
+        assert self.manager.take(2, 0).data["A"].to_list() == [1, 2]
+        assert self.manager.take(2, 1).data["A"].to_list() == [2, 3]
+        assert self.manager.take(2, 2).data["A"].to_list() == [3]
+
+    def test_take_zero(self) -> None:
+        limited_manager = self.manager.take(0, 0)
+        assert limited_manager.data.is_empty()
+
+    def test_take_negative(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(-1, 0)
+
+    def test_take_negative_offset(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(1, -1)
+
     def test_take_out_of_bounds(self) -> None:
         # Too large of page
         assert len(self.manager.take(10, 0).data) == 3

tests/_plugins/ui/_impl/tables/test_polars_table.py

@@ -265,6 +265,24 @@ def test_limit(self) -> None:
         assert "PolarsTableManager" in str(type(limited_manager))
         assert assert_frame_equal(limited_manager.data, expected_data)
 
+    def test_take(self) -> None:
+        assert self.manager.take(1, 0).data["A"].to_list() == [1]
+        assert self.manager.take(2, 0).data["A"].to_list() == [1, 2]
+        assert self.manager.take(2, 1).data["A"].to_list() == [2, 3]
+        assert self.manager.take(2, 2).data["A"].to_list() == [3]
+
+    def test_take_zero(self) -> None:
+        limited_manager = self.manager.take(0, 0)
+        assert limited_manager.data.is_empty()
+
+    def test_take_negative(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(-1, 0)
+
+    def test_take_negative_offset(self) -> None:
+        with pytest.raises(ValueError):
+            self.manager.take(1, -1)
+
     def test_take_out_of_bounds(self) -> None:
         # Too large of page
         assert len(self.manager.take(10, 0).data) == 3

tests/_plugins/ui/_impl/test_table.py

@@ -7,13 +7,15 @@
 import pytest
 
 from marimo._dependencies.dependencies import DependencyManager
+from marimo._output.data.data import from_data_uri
 from marimo._plugins import ui
 from marimo._plugins.ui._impl.dataframes.transforms.types import Condition
 from marimo._plugins.ui._impl.table import SearchTableArgs, SortArgs
 from marimo._plugins.ui._impl.tables.default_table import DefaultTableManager
 from marimo._plugins.ui._impl.utils.dataframe import TableData
 from marimo._runtime.functions import EmptyArgs
 from marimo._runtime.runtime import Kernel
+from tests._data.mocks import create_dataframes
 
 
 @pytest.fixture
@@ -211,24 +213,18 @@ def test_sort_dict_of_tuples(dtm: DefaultTableManager) -> None:
 
 def test_value() -> None:
     data = ["banana", "apple", "cherry", "date", "elderberry"]
-    data = _normalize_data(data)
     table = ui.table(data)
     assert list(table.value) == []
 
 
 def test_value_with_selection() -> None:
     data = ["banana", "apple", "cherry", "date", "elderberry"]
-    data = _normalize_data(data)
     table = ui.table(data)
-    assert list(table._convert_value(["0", "2"])) == [
-        {"value": "banana"},
-        {"value": "cherry"},
-    ]
+    assert list(table._convert_value(["0", "2"])) == ["banana", "cherry"]
 
 
 def test_value_with_sorting_then_selection() -> None:
     data = ["banana", "apple", "cherry", "date", "elderberry"]
-    data = _normalize_data(data)
     table = ui.table(data)
 
     table.search(
@@ -257,9 +253,37 @@ def test_value_with_sorting_then_selection() -> None:
     ]
 
 
+@pytest.mark.parametrize(
+    "df",
+    create_dataframes(
+        {"a": ["x", "z", "y"]}, exclude=["ibis", "duckdb", "pyarrow"]
+    ),
+)
+def test_value_with_sorting_then_selection_dfs(df: Any) -> None:
+    import narwhals as nw
+
+    table = ui.table(df)
+    table.search(
+        SearchTableArgs(
+            sort=SortArgs("a", descending=True),
+            page_size=10,
+            page_number=0,
+        )
+    )
+    assert nw.from_native(table._convert_value(["0"]))["a"][0] == "z"
+
+    table.search(
+        SearchTableArgs(
+            sort=SortArgs("a", descending=False),
+            page_size=10,
+            page_number=0,
+        )
+    )
+    assert nw.from_native(table._convert_value(["0"]))["a"][0] == "x"
+
+
 def test_value_with_search_then_selection() -> None:
     data = ["banana", "apple", "cherry", "date", "elderberry"]
-    data = _normalize_data(data)
     table = ui.table(data)
 
     table.search(
@@ -291,7 +315,45 @@ def test_value_with_search_then_selection() -> None:
             page_number=0,
         )
     )
-    assert list(table._convert_value(["2"])) == [{"value": "cherry"}]
+    assert list(table._convert_value(["2"])) == ["cherry"]
+
+
+@pytest.mark.parametrize(
+    "df",
+    create_dataframes(
+        {"a": ["foo", "bar", "baz"]}, exclude=["ibis", "duckdb", "pyarrow"]
+    ),
+)
+def test_value_with_search_then_selection_dfs(df: Any) -> None:
+    import narwhals as nw
+
+    table = ui.table(df)
+    table.search(
+        SearchTableArgs(
+            query="bar",
+            page_size=10,
+            page_number=0,
+        )
+    )
+    assert nw.from_native(table._convert_value(["0"]))["a"][0] == "bar"
+
+    table.search(
+        SearchTableArgs(
+            query="foo",
+            page_size=10,
+            page_number=0,
+        )
+    )
+    assert nw.from_native(table._convert_value(["0"]))["a"][0] == "foo"
+
+    # empty search
+    table.search(
+        SearchTableArgs(
+            page_size=10,
+            page_number=0,
+        )
+    )
+    assert nw.from_native(table._convert_value(["2"]))["a"][0] == "baz"
 
 
 def test_table_with_too_many_columns_passes() -> None:
@@ -337,6 +399,29 @@ def test_can_get_second_page_with_search() -> None:
     assert result.data[-1]["a"] == 27
 
 
+@pytest.mark.parametrize(
+    "df",
+    create_dataframes({"a": list(range(40))}, ["ibis"]),
+)
+def test_can_get_second_page_with_search_df(df: Any) -> None:
+    import polars as pl
+
+    table = ui.table(df)
+    result = table.search(
+        SearchTableArgs(
+            query="2",
+            page_size=5,
+            page_number=1,
+        )
+    )
+    mime_type, data = from_data_uri(result.data)
+    assert mime_type == "text/csv"
+    data = pl.read_csv(data)
+    assert len(data) == 5
+    assert int(data["a"][0]) == 23
+    assert int(data["a"][-1]) == 27
+
+
 def test_with_no_pagination() -> None:
     data = {"a": list(range(20))}
     table = ui.table(data, pagination=False)