Fixes #16, handle the case when read_json returns an iterator

Author: sdpython

_unittests/ut_df/test_dataframe_io_helpers.py

@@ -153,11 +153,13 @@ def test_read_json_rows(self):
  def test_read_json_rows2(self):
  data = b'''{"a": 1, "b": 2}
  {"a": 3, "b": 4}'''
+ dfs = pandas.read_json(BytesIO(data), lines=True)
+ self.assertEqual(dfs.shape, (2, 2))
  it = StreamingDataFrame.read_json(BytesIO(data), lines="stream")
  dfs = list(it)
  self.assertEqual(len(dfs), 1)
  js = dfs[0].to_json(orient='records')
- self.assertEqual(js, '[{"a":1,"b":2},{"a":3,"b":4}]')
+ self.assertEqual('[{"a":1,"b":2},{"a":3,"b":4}]', js)
 
  def test_read_json_ijson(self):
  it = StreamingDataFrame.read_json(
@@ -207,14 +209,14 @@ def test_read_json_file2(self):
  it = StreamingDataFrame.read_json(
  BytesIO(data), lines="stream", flatten=True)
  dfs = list(it)
- self.assertEqual(list(sorted(dfs[0].columns)), [
- 'a_a', 'a_c', 'b_0', 'b_1', 'b_2'])
+ self.assertEqual(['a_a', 'a_c', 'b_0', 'b_1', 'b_2'],
+ list(sorted(dfs[0].columns)), )
  self.assertEqual(len(dfs), 1)
  js = dfs[0].to_json(orient='records', lines=True)
  jsjson = loads('[' + js.replace("\n", ",") + ']')
  exp = [{'a_a': None, 'a_c': 1.0, 'b_0': 2, 'b_1': 3, 'b_2': None},
  {'a_a': 3.0, 'a_c': None, 'b_0': 4, 'b_1': 5, 'b_2': 'r'}]
- self.assertEqual(jsjson, exp)
+ self.assertEqual(exp, jsjson)
 
  def test_read_json_item(self):
  text = TestDataFrameIOHelpers.text_json

pandas_streaming/df/dataframe.py

@@ -204,30 +204,64 @@ def read_json(*args, chunksize=100000, flatten=False, **kwargs) -> 'StreamingDat
  print(dfs)
  """
  if not isinstance(chunksize, int) or chunksize <= 0:
- raise ValueError(
- 'chunksize must be a positive integer') # pragma: no cover
+ raise ValueError( # pragma: no cover
+ 'chunksize must be a positive integer')
  kwargs_create = StreamingDataFrame._process_kwargs(kwargs)
+
  if isinstance(args[0], (list, dict)):
  if flatten:
  return StreamingDataFrame.read_df(json_normalize(args[0]), **kwargs_create)
  return StreamingDataFrame.read_df(args[0], **kwargs_create)
+
  if kwargs.get('lines', None) == 'stream':
  del kwargs['lines']
  st = JsonIterator2Stream(enumerate_json_items(
  args[0], encoding=kwargs.get('encoding', None), lines=True, flatten=flatten))
  args = args[1:]
- return StreamingDataFrame(lambda: pandas.read_json(st, *args, chunksize=chunksize, lines=True, **kwargs), **kwargs_create)
+
+ if chunksize is None:
+ return StreamingDataFrame(
+ lambda: pandas.read_json(
+ st, *args, chunksize=None, lines=True, **kwargs),
+ **kwargs_create)
+
+ def fct1(st=st, args=args, chunksize=chunksize, kw=kwargs.copy()):
+ for r in pandas.read_json(st, *args, chunksize=chunksize, nrows=chunksize,
+ lines=True, **kw):
+ yield r
+ return StreamingDataFrame(fct1, **kwargs_create)
+
  if kwargs.get('lines', False):
  if flatten:
  raise NotImplementedError(
  "flatten==True is implemented with option lines='stream'")
- return StreamingDataFrame(lambda: pandas.read_json(*args, chunksize=chunksize, **kwargs), **kwargs_create)
+ if chunksize is None:
+ return StreamingDataFrame(
+ lambda: pandas.read_json(*args, chunksize=None, **kwargs),
+ **kwargs_create)
+
+ def fct2(args=args, chunksize=chunksize, kw=kwargs.copy()):
+ for r in pandas.read_json(*args, chunksize=chunksize, nrows=chunksize, **kw):
+ yield r
+ return StreamingDataFrame(fct2, **kwargs_create)
+
  st = JsonIterator2Stream(enumerate_json_items(
  args[0], encoding=kwargs.get('encoding', None), flatten=flatten))
  args = args[1:]
  if 'lines' in kwargs:
  del kwargs['lines']
- return StreamingDataFrame(lambda: pandas.read_json(st, *args, chunksize=chunksize, lines=True, **kwargs), **kwargs_create)
+
+ if chunksize is None:
+ return StreamingDataFrame(
+ lambda: pandas.read_json(
+ st, *args, chunksize=chunksize, lines=True, **kwargs),
+ **kwargs_create)
+
+ def fct3(st=st, args=args, chunksize=chunksize, kw=kwargs.copy()):
+ for r in pandas.read_json(st, *args, chunksize=chunksize, nrows=chunksize,
+ lines=True, **kw):
+ yield r
+ return StreamingDataFrame(fct3, **kwargs_create)
 
  @staticmethod
  def read_csv(*args, **kwargs) -> 'StreamingDataFrame':