Implement __len__ and leverage parquet statistics

dask_expr/collection.py

@@ -22,6 +22,7 @@
 from dask_expr.merge import Merge
 from dask_expr.reductions import (
     DropDuplicates,
+    Len,
     MemoryUsageFrame,
     MemoryUsageIndex,
     NLargest,
@@ -85,6 +86,9 @@ def _meta(self):
     def size(self):
         return new_collection(self.expr.size)
 
+    def __len__(self):
+        return new_collection(Len(self.expr)).compute()
+
     @property
     def nbytes(self):
         raise NotImplementedError("nbytes is not implemented on DataFrame")

dask_expr/expr.py

@@ -20,6 +20,7 @@
     is_dataframe_like,
     is_index_like,
     is_series_like,
+    make_meta,
 )
 from dask.dataframe.dispatch import meta_nonempty
 from dask.utils import M, apply, funcname, import_required, is_arraylike
@@ -636,6 +637,23 @@ def visualize(self, filename="dask-expr.svg", format=None, **kwargs):
         return g
 
 
+class Literal(Expr):
+    """Represent a literal (known) value as an `Expr`"""
+
+    _parameters = ["value"]
+
+    def _divisions(self):
+        return (None, None)
+
+    @property
+    def _meta(self):
+        return make_meta(self.value)
+
+    def _task(self, index: int):
+        assert index == 0
+        return self.value
+
+
 class Blockwise(Expr):
     """Super-class for block-wise operations
 
@@ -1031,6 +1049,33 @@ def _task(self, index: int):
         )
 
 
+class Lengths(Expr):
+    """Returns a tuple of partition lengths"""
+
+    _parameters = ["frame"]
+
+    @property
+    def _meta(self):
+        return tuple()
+
+    def _divisions(self):
+        return (None, None)
+
+    def _simplify_down(self):
+        if isinstance(self.frame, Elemwise):
+            child = max(self.frame.dependencies(), key=lambda expr: expr.npartitions)
+            return Lengths(child)
+
+    def _layer(self):
+        name = "part-" + self._name
+        dsk = {
+            (name, i): (len, (self.frame._name, i))
+            for i in range(self.frame.npartitions)
+        }
+        dsk[(self._name, 0)] = (tuple, list(dsk.keys()))
+        return dsk
+
+
 class ResetIndex(Elemwise):
     """Reset the index of a Series or DataFrame"""
 

dask_expr/io/io.py

@@ -1,9 +1,12 @@
+from __future__ import annotations
+
 import functools
 import math
 
 from dask.dataframe.io.io import sorted_division_locations
 
-from dask_expr.expr import Blockwise, Expr, PartitionsFiltered
+from dask_expr.expr import Blockwise, Expr, Lengths, Literal, PartitionsFiltered
+from dask_expr.reductions import Len
 
 
 class IO(Expr):
@@ -44,6 +47,7 @@ class FromPandas(PartitionsFiltered, BlockwiseIO):
 
     _parameters = ["frame", "npartitions", "sort", "_partitions"]
     _defaults = {"npartitions": 1, "sort": True, "_partitions": None}
+    _pd_length_stats = None
 
     @property
     def _meta(self):
@@ -68,6 +72,27 @@ def _divisions_and_locations(self):
             divisions = (None,) * len(locations)
         return divisions, locations
 
+    def _get_lengths(self) -> tuple | None:
+        if self._pd_length_stats is None:
+            locations = self._locations()
+            self._pd_length_stats = tuple(
+                offset - locations[i]
+                for i, offset in enumerate(locations[1:])
+                if not self._filtered or i in self._partitions
+            )
+        return self._pd_length_stats
+
+    def _simplify_up(self, parent):
+        if isinstance(parent, Lengths):
+            _lengths = self._get_lengths()
+            if _lengths:
+                return Literal(_lengths)
+
+        if isinstance(parent, Len):
+            _lengths = self._get_lengths()
+            if _lengths:
+                return Literal(sum(_lengths))
+
     def _divisions(self):
         return self._divisions_and_locations[0]