feat(api): add Table.sample method for sampling rows from a table

ibis/expr/format.py

@@ -291,6 +291,7 @@ def _join(op, left, right, predicates, **kwargs):
 
 
 @fmt.register(ops.Limit)
+@fmt.register(ops.Sample)
 def _limit(op, table, **kwargs):
     params = inline_args(kwargs)
     return f"{op.__class__.__name__}[{table}, {params}]"

ibis/expr/operations/relations.py

@@ -3,7 +3,7 @@
 import abc
 import itertools
 from abc import abstractmethod
-from typing import TYPE_CHECKING, Any, Literal, Optional
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional
 from typing import Union as UnionType
 
 from public import public
@@ -17,7 +17,7 @@
 from ibis.common.collections import FrozenDict  # noqa: TCH001
 from ibis.common.deferred import Deferred
 from ibis.common.grounds import Immutable
-from ibis.common.patterns import Coercible, Eq
+from ibis.common.patterns import Between, Coercible, Eq
 from ibis.common.typing import VarTuple  # noqa: TCH001
 from ibis.expr.operations.core import Column, Named, Node, Scalar, Value
 from ibis.expr.operations.sortkeys import SortKey  # noqa: TCH001
@@ -580,6 +580,20 @@ def schema(self):
         return self.table.schema
 
 
+@public
+class Sample(Relation):
+    """Sample performs random sampling of records in a table."""
+
+    table: Relation
+    fraction: Annotated[float, Between(0, 1)]
+    method: Literal["row", "block"] = "row"
+    seed: UnionType[int, None] = None
+
+    @attribute
+    def schema(self):
+        return self.table.schema
+
+
 # TODO(kszucs): split it into two operations, one working with a single replacement
 # value and the other with a mapping
 # TODO(kszucs): the single value case was limited to numeric and string types

ibis/expr/types/relations.py

@@ -1191,6 +1191,80 @@ def distinct(
             return res.select(self.columns)
         return res
 
+    def sample(
+        self,
+        fraction: float,
+        *,
+        method: Literal["row", "block"] = "row",
+        seed: int | None = None,
+    ) -> Table:
+        """Sample a fraction of rows from a table.
+
+        Parameters
+        ----------
+        fraction
+            The percentage of rows to include in the sample, expressed as a
+            float between 0 and 1.
+        method
+            The sampling method to use. The default is "row", which includes
+            each row with a probability of ``fraction``. If method is "block",
+            some backends may instead perform sampling a fraction of blocks of
+            rows (where "block" is a backend dependent definition). This is
+            identical to "row" for backends lacking a blockwise sampling
+            implementation. For those coming from SQL, "row" and "block"
+            correspond to "bernoulli" and "system" respectively in a
+            TABLESAMPLE clause.
+        seed
+            An optional random seed to use, for repeatable sampling. Backends
+            that never support specifying a seed for repeatable sampling will
+            error appropriately. Note that some backends (like DuckDB) do
+            support specifying a seed, but may still not have repeatable
+            results in all cases.
+
+        Returns
+        -------
+        Table
+            The input table, with `fraction` of rows selected.
+
+        Examples
+        --------
+        >>> import ibis
+        >>> ibis.options.interactive = True
+        >>> t = ibis.memtable({"x": [1, 2, 3, 4], "y": ["a", "b", "c", "d"]})
+        >>> t
+        ┏━━━━━━━┳━━━━━━━━┓
+        ┃ x     ┃ y      ┃
+        ┡━━━━━━━╇━━━━━━━━┩
+        │ int64 │ string │
+        ├───────┼────────┤
+        │     1 │ a      │
+        │     2 │ b      │
+        │     3 │ c      │
+        │     4 │ d      │
+        └───────┴────────┘
+
+        Sample approximately half the rows, with a seed specified for
+        reproducibility.
+
+        >>> t.sample(0.5, seed=1234)
+        ┏━━━━━━━┳━━━━━━━━┓
+        ┃ x     ┃ y      ┃
+        ┡━━━━━━━╇━━━━━━━━┩
+        │ int64 │ string │
+        ├───────┼────────┤
+        │     2 │ b      │
+        │     3 │ c      │
+        └───────┴────────┘
+        """
+        if fraction == 1:
+            return self
+        elif fraction == 0:
+            return self.limit(0)
+        else:
+            return ops.Sample(
+                self, fraction=fraction, method=method, seed=seed
+            ).to_expr()
+
     def limit(self, n: int | None, offset: int = 0) -> Table:
         """Select `n` rows from `self` starting at `offset`.
 

ibis/tests/expr/test_value_exprs.py

@@ -1666,3 +1666,20 @@ def test_quantile_shape():
     (b1,) = expr.op().selections
 
     assert b1.shape.is_columnar()
+
+
+def test_sample():
+    t = ibis.table({"x": "int64", "y": "string"})
+
+    expr = t.sample(1)
+    assert expr.equals(t)
+
+    expr = t.sample(0)
+    assert expr.equals(t.limit(0))
+
+    expr = t.sample(0.5, method="block", seed=1234)
+    assert expr.schema() == t.schema()
+    op = expr.op()
+    assert op.fraction == 0.5
+    assert op.method == "block"
+    assert op.seed == 1234