Expose engine argument in dask_cudf.read_json

python/dask_cudf/dask_cudf/backends.py

@@ -2,6 +2,7 @@
 
 import warnings
 from collections.abc import Iterator
+from functools import partial
 
 import cupy as cp
 import numpy as np
@@ -493,11 +494,15 @@ def read_parquet(*args, engine=None, **kwargs):
             )
 
         @staticmethod
-        def read_json(*args, engine=None, **kwargs):
+        def read_json(*args, engine="cudf", **kwargs):
             return _default_backend(
                 dd.read_json,
                 *args,
-                engine=cudf.read_json,
+                engine=(
+                    partial(cudf.read_json, engine=engine)
+                    if isinstance(engine, str)
+                    else engine
+                ),
                 **kwargs,
             )
 

python/dask_cudf/dask_cudf/io/json.py

@@ -6,4 +6,17 @@
 
 import cudf
 
-read_json = partial(dask.dataframe.read_json, engine=cudf.read_json)
+
+def read_json(path, engine="auto", **kwargs):
+    # Wrap `dd.read_json` with special engine handling
+
+    # TODO: Add optimized code path to leverage the
+    # `byte_range` argument in `cudf.read_json` for
+    # local storage (see `dask_cudf.read_csv`)
+    if isinstance(engine, str):
+        # Pass `str` engine argument to `cudf.read_json``
+        engine = partial(cudf.read_json, engine=engine)
+    elif not callable(engine):
+        raise ValueError(f"Unsupported engine option: {engine}")
+    # Pass `callable` engine argument to `dd.read_json`
+    return dask.dataframe.read_json(path, engine=engine, **kwargs)

python/dask_cudf/dask_cudf/io/tests/test_json.py

@@ -71,3 +71,22 @@ def test_read_json_lines(lines):
         actual = dask_cudf.read_json(f, orient="records", lines=lines)
         actual_pd = pd.read_json(f, orient="records", lines=lines)
         dd.assert_eq(actual, actual_pd)
+
+
+@pytest.mark.filterwarnings("ignore:Using CPU")
+def test_read_json_nested_experimental():
+    # Check that `engine="cudf_experimental"` can
+    # be used to support nested data
+    df = pd.DataFrame(
+        {
+            "a": [{"y": 2}, {"y": 4}, {"y": 6}, {"y": 8}],
+            "b": [[1, 2, 3], [4, 5], [6], [7]],
+            "c": [1, 3, 5, 7],
+        }
+    )
+    lines = dict(orient="records", lines=True)
+    with tmpfile("json") as f:
+        df.to_json(f, **lines)
+        actual = dask_cudf.read_json(f, engine="cudf_experimental", **lines)
+        actual_pd = pd.read_json(f, **lines)
+        dd.assert_eq(actual, actual_pd)