diff --git a/src/datasets/io/json.py b/src/datasets/io/json.py index eb347aea7aa..32bc6c292ef 100644 --- a/src/datasets/io/json.py +++ b/src/datasets/io/json.py @@ -104,9 +104,12 @@ def write(self) -> int: if compression not in [None, "infer", "gzip", "bz2", "xz"]: raise NotImplementedError(f"`datasets` currently does not support {compression} compression") - if not lines and self.batch_size < self.dataset.num_rows: + # Array-producing orientations ("records", "values") support batching when lines=False; + # dict-producing orientations ("split", "index", "columns", "table") do not. + if not lines and orient not in ("records", "values") and self.batch_size < self.dataset.num_rows: raise NotImplementedError( - "Output JSON will not be formatted correctly when lines = False and batch_size < number of rows in the dataset. Use pandas.DataFrame.to_json() instead." + f"Output JSON will not be formatted correctly when lines=False, orient={orient!r}, " + "and batch_size < number of rows in the dataset. Use pandas.DataFrame.to_json() instead." ) if isinstance(self.path_or_buf, (str, bytes, os.PathLike)): @@ -147,32 +150,61 @@ def _write( lines, **to_json_kwargs, ) -> int: - """Writes the pyarrow table as JSON lines to a binary file handle. + """Writes the pyarrow table as JSON to a binary file handle. - Caller is responsible for opening and closing the handle. + When ``lines=False`` and the dataset spans multiple batches, each batch's + JSON array (``[{...}, ...]``) is merged into a single outer array so the + output remains valid JSON. Caller is responsible for opening and closing + the handle. """ + import contextlib + written = 0 + num_rows, batch_size = len(self.dataset), self.batch_size + offsets = range(0, num_rows, batch_size) + total_batches = (num_rows // batch_size) + (1 if num_rows % batch_size else 0) + + pool_ctx = ( + contextlib.nullcontext() + if self.num_proc is None or self.num_proc == 1 + else multiprocessing.Pool(self.num_proc) + ) - if self.num_proc is None or self.num_proc == 1: - for offset in hf_tqdm( - range(0, len(self.dataset), self.batch_size), - unit="ba", - desc="Creating json from Arrow format", - ): - json_str = self._batch_json((offset, orient, lines, to_json_kwargs)) - written += file_obj.write(json_str) - else: - num_rows, batch_size = len(self.dataset), self.batch_size - with multiprocessing.Pool(self.num_proc) as pool: - for json_str in hf_tqdm( + with pool_ctx as pool: + if pool is None: + batch_iter = ( + self._batch_json((offset, orient, lines, to_json_kwargs)) + for offset in hf_tqdm(offsets, unit="ba", desc="Creating json from Arrow format") + ) + else: + batch_iter = hf_tqdm( pool.imap( self._batch_json, - [(offset, orient, lines, to_json_kwargs) for offset in range(0, num_rows, batch_size)], + [(offset, orient, lines, to_json_kwargs) for offset in offsets], ), - total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size, + total=total_batches, unit="ba", desc="Creating json from Arrow format", - ): - written += file_obj.write(json_str) + ) + + if not lines and orient in ("records", "values"): + # These orientations produce a JSON array per batch (e.g. `[{...}, ...]`). + # Strip the outer brackets from each batch and merge them into a single + # valid JSON array across all batches. + written += file_obj.write(b"[") + first = True + for json_bytes in batch_iter: + inner = json_bytes.rstrip(b"\n") + if inner.startswith(b"[") and inner.endswith(b"]"): + inner = inner[1:-1] + if inner: + if not first: + written += file_obj.write(b",") + written += file_obj.write(inner) + first = False + written += file_obj.write(b"]") + else: + for json_bytes in batch_iter: + written += file_obj.write(json_bytes) return written diff --git a/src/datasets/iterable_dataset.py b/src/datasets/iterable_dataset.py index 17b9a2020fc..1c00a2cba54 100644 --- a/src/datasets/iterable_dataset.py +++ b/src/datasets/iterable_dataset.py @@ -4036,7 +4036,14 @@ def add_column(self, name: str, column: Union[list, np.array]) -> "IterableDatas Returns: `IterableDataset` """ - return self.map(partial(add_column_fn, name=name, column=column), with_indices=True) + original_features = self._info.features.copy() if self._info.features else None + ds_iterable = self.map(partial(add_column_fn, name=name, column=column), with_indices=True) + if original_features is not None: + col_table = pa.table({name: column}) + new_features = original_features.copy() + new_features[name] = Features.from_arrow_schema(col_table.schema)[name] + ds_iterable._info.features = new_features + return ds_iterable def rename_column(self, original_column_name: str, new_column_name: str) -> "IterableDataset": """ diff --git a/tests/io/test_json.py b/tests/io/test_json.py index 0acde7d567b..ae16e5cb635 100644 --- a/tests/io/test_json.py +++ b/tests/io/test_json.py @@ -280,6 +280,17 @@ def test_dataset_to_json_orient_multiproc(self, orient, container, keys, len_at, else: assert len(exported_content) == 10 + @pytest.mark.parametrize("orient", ["records", "values"]) + def test_dataset_to_json_lines_false_batched(self, orient, dataset): + """lines=False with batch_size < num_rows should produce a single valid JSON array (issue #7037).""" + # Use batch_size=3 so the 10-row dataset spans multiple batches + with io.BytesIO() as buffer: + JsonDatasetWriter(dataset, buffer, lines=False, orient=orient, batch_size=3).write() + buffer.seek(0) + exported_content = load_json(buffer) + assert isinstance(exported_content, list) + assert len(exported_content) == 10 + def test_dataset_to_json_orient_invalidproc(self, dataset): with pytest.raises(ValueError): with io.BytesIO() as buffer: