Repository: dask/dask

Issue:
Add include_path_column to read_json
Similar to #6575, I ran into a situation where the `include_path_column` argument from `dd.read_csv` would have been very useful on `dd.read_json`. 


Discussion:
Hi @rfilmyer, thanks for opening this feature request. If you are interested in working on this, please feel free to open a Pull Request.
I think @bryanwweber mentioned he is planning on picking this issue up 

Resolving patch:
diff --git a/dask/dataframe/io/json.py b/dask/dataframe/io/json.py
--- a/dask/dataframe/io/json.py
+++ b/dask/dataframe/io/json.py
@@ -1,5 +1,6 @@
 import io
 import os
+from itertools import zip_longest
 
 import pandas as pd
 from fsspec.core import open_files
@@ -115,6 +116,8 @@ def read_json(
     compression="infer",
     meta=None,
     engine=pd.read_json,
+    include_path_column=False,
+    path_converter=None,
     **kwargs,
 ):
     """Create a dataframe from a set of JSON files
@@ -149,7 +152,7 @@ def read_json(
         newline character.
     sample: int
         Number of bytes to pre-load, to provide an empty dataframe structure
-        to any blocks without data. Only relevant is using blocksize.
+        to any blocks without data. Only relevant when using blocksize.
     encoding, errors:
         Text conversion, ``see bytes.decode()``
     compression : string or None
@@ -157,6 +160,14 @@ def read_json(
     engine : function object, default ``pd.read_json``
         The underlying function that dask will use to read JSON files. By
         default, this will be the pandas JSON reader (``pd.read_json``).
+    include_path_column : bool or str, optional
+        Include a column with the file path where each row in the dataframe
+        originated. If ``True``, a new column is added to the dataframe called
+        ``path``. If ``str``, sets new column name. Default is ``False``.
+    path_converter : function or None, optional
+        A function that takes one argument and returns a string. Used to convert
+        paths in the ``path`` column, for instance, to strip a common prefix from
+        all the paths.
     $META
 
     Returns
@@ -192,24 +203,62 @@ def read_json(
             "input (orient='records', lines=True)."
         )
     storage_options = storage_options or {}
+    if include_path_column is True:
+        include_path_column = "path"
+
+    if path_converter is None:
+        path_converter = lambda x: x
+
     if blocksize:
-        first, chunks = read_bytes(
+        b_out = read_bytes(
             url_path,
             b"\n",
             blocksize=blocksize,
             sample=sample,
             compression=compression,
+            include_path=include_path_column,
             **storage_options,
         )
-        chunks = list(flatten(chunks))
+        if include_path_column:
+            first, chunks, paths = b_out
+            first_path = path_converter(paths[0])
+            path_dtype = pd.CategoricalDtype(path_converter(p) for p in paths)
+            flat_paths = flatten(
+                [path_converter(p)] * len(chunk) for p, chunk in zip(paths, chunks)
+            )
+        else:
+            first, chunks = b_out
+            first_path = None
+            flat_paths = (None,)
+            path_dtype = None
+
+        flat_chunks = flatten(chunks)
         if meta is None:
-            meta = read_json_chunk(first, encoding, errors, engine, kwargs)
+            meta = read_json_chunk(
+                first,
+                encoding,
+                errors,
+                engine,
+                include_path_column,
+                first_path,
+                path_dtype,
+                kwargs,
+            )
         meta = make_meta(meta)
         parts = [
-            delayed(read_json_chunk)(chunk, encoding, errors, engine, kwargs, meta=meta)
-            for chunk in chunks
+            delayed(read_json_chunk)(
+                chunk,
+                encoding,
+                errors,
+                engine,
+                include_path_column,
+                path,
+                path_dtype,
+                kwargs,
+                meta=meta,
+            )
+            for chunk, path in zip_longest(flat_chunks, flat_paths)
         ]
-        return from_delayed(parts, meta=meta)
     else:
         files = open_files(
             url_path,
@@ -219,22 +268,52 @@ def read_json(
             compression=compression,
             **storage_options,
         )
+        path_dtype = pd.CategoricalDtype(path_converter(f.path) for f in files)
         parts = [
-            delayed(read_json_file)(f, orient, lines, engine, kwargs) for f in files
+            delayed(read_json_file)(
+                f,
+                orient,
+                lines,
+                engine,
+                include_path_column,
+                path_converter(f.path),
+                path_dtype,
+                kwargs,
+            )
+            for f in files
         ]
-        return from_delayed(parts, meta=meta)
 
+    return from_delayed(parts, meta=meta)
 
-def read_json_chunk(chunk, encoding, errors, engine, kwargs, meta=None):
+
+def read_json_chunk(
+    chunk, encoding, errors, engine, column_name, path, path_dtype, kwargs, meta=None
+):
     s = io.StringIO(chunk.decode(encoding, errors))
     s.seek(0)
     df = engine(s, orient="records", lines=True, **kwargs)
     if meta is not None and df.empty:
         return meta
-    else:
-        return df
 
+    if column_name:
+        df = add_path_column(df, column_name, path, path_dtype)
 
-def read_json_file(f, orient, lines, engine, kwargs):
-    with f as f:
-        return engine(f, orient=orient, lines=lines, **kwargs)
+    return df
+
+
+def read_json_file(f, orient, lines, engine, column_name, path, path_dtype, kwargs):
+    with f as open_file:
+        df = engine(open_file, orient=orient, lines=lines, **kwargs)
+    if column_name:
+        df = add_path_column(df, column_name, path, path_dtype)
+    return df
+
+
+def add_path_column(df, column_name, path, dtype):
+    if column_name in df.columns:
+        raise ValueError(
+            f"Files already contain the column name: '{column_name}', so the path "
+            "column cannot use this name. Please set `include_path_column` to a "
+            "unique name."
+        )
+    return df.assign(**{column_name: pd.Series([path] * len(df), dtype=dtype)})


Test patch:
diff --git a/dask/dataframe/io/tests/test_json.py b/dask/dataframe/io/tests/test_json.py
--- a/dask/dataframe/io/tests/test_json.py
+++ b/dask/dataframe/io/tests/test_json.py
@@ -13,6 +13,86 @@
 ddf = dd.from_pandas(df, npartitions=2)
 
 
+@pytest.mark.parametrize("orient", ["split", "records", "index", "columns", "values"])
+def test_read_json_with_path_column(orient):
+    with tmpfile("json") as f:
+        df.to_json(f, orient=orient, lines=False)
+        actual = dd.read_json(f, orient=orient, lines=False, include_path_column=True)
+        actual_pd = pd.read_json(f, orient=orient, lines=False)
+        # The default column name when include_path_colum is True is "path"
+        # The paths on Windows are converted to forward slash somewhere in the file
+        # reading chain in Dask, so we have to do the same here.
+        actual_pd["path"] = pd.Series(
+            (f.replace(os.sep, "/"),) * len(actual_pd), dtype="category"
+        )
+        assert actual.path.dtype == "category"
+        assert_eq(actual, actual_pd)
+
+
+def test_read_json_path_column_with_duplicate_name_is_error():
+    with tmpfile("json") as f:
+        df.to_json(f, orient="records", lines=False)
+        with pytest.raises(ValueError, match="Files already contain"):
+            dd.read_json(f, orient="records", lines=False, include_path_column="x")
+
+
+def test_read_json_with_path_converter():
+    path_column_name = "filenames"
+
+    def path_converter(x):
+        return "asdf.json"
+
+    with tmpfile("json") as f:
+        df.to_json(f, orient="records", lines=False)
+        actual = dd.read_json(
+            f,
+            orient="records",
+            lines=False,
+            include_path_column=path_column_name,
+            path_converter=path_converter,
+        )
+        actual_pd = pd.read_json(f, orient="records", lines=False)
+        actual_pd[path_column_name] = pd.Series(
+            (path_converter(f),) * len(actual_pd), dtype="category"
+        )
+        assert_eq(actual, actual_pd)
+
+
+def test_read_orient_not_records_and_lines():
+    with pytest.raises(ValueError, match="Line-delimited JSON"):
+        dd.read_json("nofile.json", orient="split", lines=True)
+
+
+def test_write_orient_not_records_and_lines():
+    with tmpfile("json") as f:
+        with pytest.raises(ValueError, match="Line-delimited JSON"):
+            dd.to_json(ddf, f, orient="split", lines=True)
+
+
+@pytest.mark.parametrize("blocksize", [5, 15, 33, 200, 90000])
+def test_read_json_multiple_files_with_path_column(blocksize, tmpdir):
+    fil1 = str(tmpdir.join("fil1.json")).replace(os.sep, "/")
+    fil2 = str(tmpdir.join("fil2.json")).replace(os.sep, "/")
+    df = pd.DataFrame({"x": range(5), "y": ["a", "b", "c", "d", "e"]})
+    df2 = df.assign(x=df.x + 0.5)
+    orient = "records"
+    lines = True
+    df.to_json(fil1, orient=orient, lines=lines)
+    df2.to_json(fil2, orient=orient, lines=lines)
+    path_dtype = pd.CategoricalDtype((fil1, fil2))
+    df["path"] = pd.Series((fil1,) * len(df), dtype=path_dtype)
+    df2["path"] = pd.Series((fil2,) * len(df2), dtype=path_dtype)
+    sol = pd.concat([df, df2])
+    res = dd.read_json(
+        str(tmpdir.join("fil*.json")),
+        orient=orient,
+        lines=lines,
+        include_path_column=True,
+        blocksize=blocksize,
+    )
+    assert_eq(res, sol, check_index=False)
+
+
 @pytest.mark.parametrize("orient", ["split", "records", "index", "columns", "values"])
 def test_read_json_basic(orient):
     with tmpfile("json") as f:
@@ -20,11 +100,10 @@ def test_read_json_basic(orient):
         actual = dd.read_json(f, orient=orient, lines=False)
         actual_pd = pd.read_json(f, orient=orient, lines=False)
 
-        out = actual.compute()
-        assert_eq(out, actual_pd)
+        assert_eq(actual, actual_pd)
         if orient == "values":
-            out.columns = list(df.columns)
-        assert_eq(out, df)
+            actual.columns = list(df.columns)
+        assert_eq(actual, df)
 
 
 @pytest.mark.parametrize("fkeyword", ["pandas", "json"])
@@ -78,10 +157,9 @@ def test_write_json_basic(orient):
         fn = os.path.join(path, "1.json")
         df.to_json(fn, orient=orient, lines=False)
         actual = dd.read_json(fn, orient=orient, lines=False)
-        out = actual.compute()
         if orient == "values":
-            out.columns = list(df.columns)
-        assert_eq(out, df)
+            actual.columns = list(df.columns)
+        assert_eq(actual, df)
 
 
 def test_to_json_with_get():
@@ -127,7 +205,7 @@ def test_json_compressed(compression):
     with tmpdir() as path:
         dd.to_json(ddf, path, compression=compression)
         actual = dd.read_json(os.path.join(path, "*"), compression=compression)
-        assert_eq(df, actual.compute(), check_index=False)
+        assert_eq(df, actual, check_index=False)
 
 
 def test_read_json_inferred_compression():
@@ -135,7 +213,7 @@ def test_read_json_inferred_compression():
         fn = os.path.join(path, "*.json.gz")
         dd.to_json(ddf, fn, compression="gzip")
         actual = dd.read_json(fn)
-        assert_eq(df, actual.compute(), check_index=False)
+        assert_eq(df, actual, check_index=False)
 
 
 def test_to_json_results():