Repository: pandas-dev/pandas
Pull request: Add nrows to read json.
- [x] closes #33916
- [x] tests added / passed
- [x] passes `black pandas`
- [x] passes `git diff upstream/master -u -- "*.py" | flake8 --diff`

Add the nrows to read_json parameter that returns only the required number of json from the line delimited json 

Problem statement:
ENH: Add nrows parameter to pd.read_json
#### Is your feature request related to a problem?

Let's say I have a huge `jsonlines` file and I want to read only the first `n` lines of the file.

#### Describe the solution you'd like

This problem could be fixed by adding `nrows` parameter in method `pd.read_json`. This parameter should be applicable if and only if `lines=True`

#### API breaking implications

Simply add and implement `nrows` parameter in `pd.read_json` (like `pd.read_csv`).

#### Additional context

How this enhancement could work:

```python
import pandas as pd

df = pd.read_json('/path/to/file.jsonlines', lines=True, nrows=1000)
```

Hints from issue discussion:
Hi @sp1thas ,
I checked the code and I feel, I should be able to implement this feature.

Commit message: ENH solve linting via black8 for Add nrow parameter for line delimited json for read_json #33916

Patch to review:
diff --git a/pandas/io/json/_json.py b/pandas/io/json/_json.py
index ac6f9ff372601..f4940b9fc26ac 100644
--- a/pandas/io/json/_json.py
+++ b/pandas/io/json/_json.py
@@ -363,6 +363,7 @@ def read_json(
     lines=False,
     chunksize=None,
     compression="infer",
+    nrows=None,
 ):
     """
     Convert a JSON string to pandas object.
@@ -493,6 +494,12 @@ def read_json(
         for more information on ``chunksize``.
         This can only be passed if `lines=True`.
         If this is None, the file will be read into memory all at once.
+
+    chunksize : int, optional
+        The number of lines from the line-delimited jsonfile that has to be read.
+        This can only be passed if `lines=True`.
+        If this is None, all the rows will be returned.
+
     compression : {'infer', 'gzip', 'bz2', 'zip', 'xz', None}, default 'infer'
         For on-the-fly decompression of on-disk data. If 'infer', then use
         gzip, bz2, zip or xz if path_or_buf is a string ending in
@@ -600,6 +607,7 @@ def read_json(
         lines=lines,
         chunksize=chunksize,
         compression=compression,
+        nrows=nrows,
     )
 
     if chunksize:
@@ -637,6 +645,7 @@ def __init__(
         lines,
         chunksize,
         compression,
+        nrows,
     ):
 
         self.path_or_buf = filepath_or_buffer
@@ -655,11 +664,16 @@ def __init__(
         self.chunksize = chunksize
         self.nrows_seen = 0
         self.should_close = False
+        self.nrows = nrows
 
         if self.chunksize is not None:
             self.chunksize = _validate_integer("chunksize", self.chunksize, 1)
             if not self.lines:
                 raise ValueError("chunksize can only be passed if lines=True")
+        if self.nrows is not None:
+            self.chunksize = _validate_integer("nrows", self.nrows, 0)
+            if not self.lines:
+                raise ValueError("nrows can only be passed if lines=True")
 
         data = self._get_data_from_filepath(filepath_or_buffer)
         self.data = self._preprocess_data(data)
@@ -726,7 +740,10 @@ def read(self):
             obj = concat(self)
         elif self.lines:
             data = ensure_str(self.data)
-            obj = self._get_object_parser(self._combine_lines(data.split("\n")))
+            data = data.split("\n")
+            if self.nrows:
+                data = data[: self.nrows]
+            obj = self._get_object_parser(self._combine_lines(data))
         else:
             obj = self._get_object_parser(self.data)
         self.close()
@@ -773,6 +790,11 @@ def close(self):
                 pass
 
     def __next__(self):
+        if self.nrows:
+            if self.nrows_seen >= self.nrows:
+                self.close()
+                raise StopIteration
+
         lines = list(islice(self.data, self.chunksize))
         if lines:
             lines_json = self._combine_lines(lines)
diff --git a/pandas/tests/io/json/test_readlines.py b/pandas/tests/io/json/test_readlines.py
index e531457627342..a86070c1614d2 100644
--- a/pandas/tests/io/json/test_readlines.py
+++ b/pandas/tests/io/json/test_readlines.py
@@ -130,6 +130,7 @@ def test_readjson_chunks_closes(chunksize):
             lines=True,
             chunksize=chunksize,
             compression=None,
+            nrows=None,
         )
         reader.read()
         assert (
@@ -179,3 +180,28 @@ def test_readjson_unicode(monkeypatch):
         result = read_json(path)
         expected = pd.DataFrame({"£©µÀÆÖÞßéöÿ": ["АБВГДабвгд가"]})
         tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("nrows", [1, 2])
+def test_readjson_nrows(nrows):
+    # Test reading line-format JSON to Series with nrows param
+    jsonl = """{"a": 1, "b": 2}
+        {"a": 3, "b": 4}
+        {"a": 5, "b": 6}
+        {"a": 7, "b": 8}"""
+    result = pd.read_json(jsonl, lines=True, nrows=nrows)
+    expected = pd.DataFrame({"a": [1, 3, 5, 7], "b": [2, 4, 6, 8]}).iloc[:nrows]
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("nrows,chunksize", [(2, 2), (4, 2)])
+def test_readjson_nrows_chunks(nrows, chunksize):
+    # Test reading line-format JSON to Series with nrows and chunksize param
+    jsonl = """{"a": 1, "b": 2}
+        {"a": 3, "b": 4}
+        {"a": 5, "b": 6}
+        {"a": 7, "b": 8}"""
+    reader = read_json(jsonl, lines=True, nrows=nrows, chunksize=chunksize)
+    chunked = pd.concat(reader)
+    expected = pd.DataFrame({"a": [1, 3, 5, 7], "b": [2, 4, 6, 8]}).iloc[:nrows]
+    tm.assert_frame_equal(chunked, expected)


Code review comments:
diff_hunk: @@ -363,6 +363,7 @@ def read_json(
     lines=False,
     chunksize=None,
     compression="infer",
+    nrows=None,
path: pandas/io/json/_json.py
text: @user1:
add typing to new args (existing ones if you can / want)

@author:
Okay

diff_hunk: @@ -726,7 +740,10 @@ def read(self):
             obj = concat(self)
         elif self.lines:
             data = ensure_str(self.data)
-            obj = self._get_object_parser(self._combine_lines(data.split("\n")))
+            data = data.split("\n")
+            if self.nrows:
path: pandas/io/json/_json.py
text: @user1:
Can you implement this at a lower level? I think should be done in the C extension to match behavior of other IO methods

@author:
I have not worked with lower-level extension till now. If it's to be done there, I will look into it
Also, it would be very nice to get some direction on how to achieve it.

@user2:
I actually don't think we need to implement this on a lower level  (lines is not implemented that way currently).

rather you need to slice this differently to avoid putting all of this into memory

e.g. you can do this: https://stackoverflow.com/questions/3862010/is-there-a-generator-version-of-string-split-in-python

to create a generator on the split rows, then islice them, process and repeat