fix: have unproject_layout modify the input layout again (#327)

src/dask_awkward/lib/unproject_layout.py

@@ -403,7 +403,6 @@ def unproject_layout(form: Form | None, layout: Content) -> Content:
         not appear in the projected layout will be PlaceholderArrays).
 
     """
-    return layout
-    # if form is None:
-    #     return layout
-    # return _unproject_layout(form, layout, layout.length, layout.backend)
+    if form is None:
+        return layout
+    return _unproject_layout(form, layout, layout.length, layout.backend)

tests/test_core.py

@@ -2,10 +2,11 @@
 
 import copy
 from collections import namedtuple
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
 import awkward as ak
 import dask.array as da
+import dask.config
 import fsspec
 import numpy as np
 import pytest
@@ -826,3 +827,14 @@ def test_map_partitions_no_dask_collections_passed(caa):
         match="map_partitions expects at least one Dask collection",
     ):
         dak.num(caa.points.x, axis=1)
+
+
+@pytest.mark.parametrize("fn", [dak.count, dak.zeros_like, dak.ones_like])
+def test_shape_only_ops(fn: Callable, tmp_path_factory: pytest.TempPathFactory) -> None:
+    a = ak.Array([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+    p = tmp_path_factory.mktemp("zeros-like-flat")
+    ak.to_parquet(a, str(p / "file.parquet"))
+    lazy = dak.from_parquet(str(p))
+    result = fn(lazy.b)
+    with dask.config.set({"awkward.optimization.enabled": True}):
+        result.compute()