diff --git a/src/dask_awkward/__init__.py b/src/dask_awkward/__init__.py index a9726941..9b205092 100644 --- a/src/dask_awkward/__init__.py +++ b/src/dask_awkward/__init__.py @@ -18,7 +18,7 @@ map_partitions, partition_compatibility, ) -from dask_awkward.lib.describe import fields +from dask_awkward.lib.describe import backend, fields from dask_awkward.lib.inspect import ( report_necessary_buffers, report_necessary_columns, @@ -90,6 +90,7 @@ singletons, sort, strings_astype, + to_list, to_packed, to_regular, unflatten, diff --git a/src/dask_awkward/lib/__init__.py b/src/dask_awkward/lib/__init__.py index a3d0703f..90835ff7 100644 --- a/src/dask_awkward/lib/__init__.py +++ b/src/dask_awkward/lib/__init__.py @@ -7,7 +7,7 @@ map_partitions, partition_compatibility, ) -from dask_awkward.lib.describe import fields +from dask_awkward.lib.describe import backend, fields from dask_awkward.lib.inspect import ( report_necessary_buffers, report_necessary_columns, @@ -76,6 +76,7 @@ singletons, sort, strings_astype, + to_list, to_packed, to_regular, unflatten, diff --git a/src/dask_awkward/lib/describe.py b/src/dask_awkward/lib/describe.py index c879febc..57e14a20 100644 --- a/src/dask_awkward/lib/describe.py +++ b/src/dask_awkward/lib/describe.py @@ -1,6 +1,8 @@ from __future__ import annotations -from dask_awkward.lib.core import Array, Record +import awkward as ak + +from dask_awkward.lib.core import Array, Record, Scalar def fields(collection: Array | Record) -> list[str] | None: @@ -19,3 +21,22 @@ def fields(collection: Array | Record) -> list[str] | None: """ return collection.fields + + +def backend(*arrays: Array | Record) -> str: + """Get the name of the backend used by `arrays`. + + Parameters + ---------- + arrays : dask_awkward.Array or dask_awkward.Record + Array or Record collection + + Returns + ------- + str + The backend name, which is always `"typetracer"` for + dask-awkward arrays. + """ + return ak.backend( + *[x._meta if isinstance(x, (Array, Record, Scalar)) else x for x in arrays] + ) diff --git a/src/dask_awkward/lib/structure.py b/src/dask_awkward/lib/structure.py index 6656b1e7..6b9fafad 100644 --- a/src/dask_awkward/lib/structure.py +++ b/src/dask_awkward/lib/structure.py @@ -62,6 +62,7 @@ "singletons", "sort", "strings_astype", + "to_list", "to_packed", "to_regular", "unflatten", @@ -691,6 +692,15 @@ def ones_like( ) +@borrow_docstring(ak.to_list) +def to_list(array: Array) -> list: + """Return list/dict version of the data + + Unlike most functions, this one requires a compute() of the data. + """ + return array.compute().to_list() + + @borrow_docstring(ak.to_packed) def to_packed( array: Array, diff --git a/tests/test_behavior.py b/tests/test_behavior.py index 3a744874..e14add9c 100644 --- a/tests/test_behavior.py +++ b/tests/test_behavior.py @@ -65,6 +65,7 @@ def test_distance_behavior( daa2 = dak.with_name(daa_p2.points, name="Point", behavior=behaviors) caa1 = ak.Array(caa_p1.points, with_name="Point", behavior=behaviors) caa2 = ak.Array(caa_p2.points) + assert_eq(daa1.distance(daa2), caa1.distance(caa2)) assert_eq(np.abs(daa1), np.abs(caa1)) diff --git a/tests/test_describe.py b/tests/test_describe.py index cbcf8a17..63eef665 100644 --- a/tests/test_describe.py +++ b/tests/test_describe.py @@ -1,14 +1,24 @@ +from typing import Any + import awkward as ak +import pytest import dask_awkward as dak -def test_fields(ndjson_points_file: str) -> None: +@pytest.mark.parametrize("quak", [ak, dak]) +def test_fields(ndjson_points_file: str, quak: Any) -> None: daa = dak.from_json([ndjson_points_file] * 2) # records fields same as array of records fields - assert dak.fields(daa[0].points) == dak.fields(daa.points) + assert quak.fields(daa[0].points) == quak.fields(daa.points) # computed is same as collection - assert dak.fields(daa) == ak.fields(daa.compute()) + assert quak.fields(daa) == ak.fields(daa.compute()) daa.reset_meta() # removed meta gives None fields - assert dak.fields(daa) == [] + assert quak.fields(daa) == [] + + +@pytest.mark.parametrize("quak", [ak, dak]) +def test_backend(ndjson_points_file: str, quak: Any) -> None: + daa = dak.from_json([ndjson_points_file] * 2) + assert quak.backend(daa) == "typetracer" diff --git a/tests/test_structure.py b/tests/test_structure.py index 851459ca..180e922e 100644 --- a/tests/test_structure.py +++ b/tests/test_structure.py @@ -458,6 +458,10 @@ def test_to_packed(daa, caa): ) +def test_to_list(daa): + assert dak.to_list(daa) == daa.compute().to_list() + + def test_ravel(daa, caa): assert_eq( dak.ravel(daa.points.x),