Implement __len__ and leverage parquet statistics

dask_expr/collection.py

@@ -17,6 +17,7 @@
 from dask_expr import expr
 from dask_expr.expr import no_default
 from dask_expr.merge import Merge
+from dask_expr.reductions import Len
 from dask_expr.repartition import Repartition
 
 #
@@ -73,6 +74,13 @@ def _meta(self):
     def size(self):
         return new_collection(self.expr.size)
 
+    def __len__(self):
+        return self._len
+
+    @functools.cached_property
+    def _len(self):
+        return new_collection(Len(self.expr)).compute()
+
     @property
     def nbytes(self):
         raise NotImplementedError("nbytes is not implemented on DataFrame")
@@ -559,7 +567,7 @@ def read_parquet(
     index=None,
     storage_options=None,
     dtype_backend=None,
-    calculate_divisions=False,
+    calculate_divisions=True,
     ignore_metadata_file=False,
     metadata_task_size=None,
     split_row_groups="infer",

dask_expr/expr.py

@@ -19,6 +19,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
@@ -39,6 +40,7 @@ class Expr:
     associative = False
     _parameters = []
     _defaults = {}
+    _lengths = None
 
     def __init__(self, *args, **kwargs):
         operands = list(args)
@@ -616,6 +618,61 @@ def visualize(self, filename="dask-expr.svg", format=None, **kwargs):
         graphviz_to_file(g, filename, format)
         return g
 
+    def _partitioning(self, columns: list) -> dict:
+        """Known partitioning information
+
+        Return known-partitioning information for the specified
+        list of columns. This information should be formatted
+        as a dict containing "columns" and "how" keys, where
+        the `"columns"` value should be a tuple of column names,
+        and the `"how"` value should be a tuple that uniquely
+        identifies the partitioning.
+
+        If no partitioning information is known, an empty
+        dictionary will be returned.
+
+        Examples of `"how"`:
+
+          - Sorted data: `("increasing", <divisions>)`
+          - Reverse-sorted data: `("decreasing", <divisions>)`
+          - Shuffled data: `("hash", <npartitions>)`
+
+        Note that un-named index columns must be specified as
+        `"__index__"` (`None` is not supported).
+
+        Return
+        ------
+        partitioning: dict
+        """
+        assert isinstance(columns, list), "columns must be list"
+
+        # By default, we only know about partitioning from known divisions
+        index_name = self._meta.index.name or "__index__"
+        if self.known_divisions and columns[0] == index_name:
+            return {
+                "columns": (index_name,),
+                "how": self.divisions,
+            }
+
+        return {}
+
+
+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
@@ -765,7 +822,9 @@ class Elemwise(Blockwise):
     optimizations, like `len` will care about which operations preserve length
     """
 
-    pass
+    @property
+    def _lengths(self):
+        return self.dependencies()[0]._lengths
 
 
 class AsType(Elemwise):
@@ -857,6 +916,9 @@ def __str__(self):
             base = "(" + base + ")"
         return f"{base}[{repr(self.columns)}]"
 
+    def _partitioning(self, columns: list) -> dict:
+        return self.frame._partitioning(columns)
+
     def _simplify_down(self):
         if isinstance(self.frame, Projection):
             # df[a][b]
@@ -1085,6 +1147,12 @@ def _simplify_down(self):
     def _node_label_args(self):
         return [self.frame, self.partitions]
 
+    @property
+    def _lengths(self):
+        lengths = self.frame._lengths
+        if lengths:
+            return tuple(lengths[i] for i in self.partitions)
+
 
 class PartitionsFiltered(Expr):
     """Mixin class for partition filtering

dask_expr/io/io.py

@@ -68,6 +68,15 @@ def _divisions_and_locations(self):
             divisions = (None,) * len(locations)
         return divisions, locations
 
+    @functools.cached_property
+    def _lengths(self):
+        locations = self._locations()
+        return tuple(
+            offset - locations[i]
+            for i, offset in enumerate(locations[1:])
+            if not self._filtered or i in self._partitions
+        )
+
     def _divisions(self):
         return self._divisions_and_locations[0]