Use pandas ujson in JSON loader to improve performance

src/datasets/packaged_modules/json/json.py

@@ -1,9 +1,9 @@
 import io
 import itertools
-import json
 from dataclasses import dataclass
 from typing import Optional
 
+import pandas as pd
 import pyarrow as pa
 import pyarrow.json as paj
 
@@ -15,6 +15,14 @@
 logger = datasets.utils.logging.get_logger(__name__)
 
 
+def ujson_loads(*args, **kwargs):
+    try:
+        return pd.io.json.ujson_loads(*args, **kwargs)
+    except AttributeError:
+        # Before pandas-2.2.0, ujson_loads was renamed to loads: import ujson_loads as loads
+        return pd.io.json.loads(*args, **kwargs)
+
+
 @dataclass
 class JsonConfig(datasets.BuilderConfig):
     """BuilderConfig for JSON."""
@@ -80,7 +88,7 @@ def _generate_tables(self, files):
             # If the file is one json object and if we need to look at the list of items in one specific field
             if self.config.field is not None:
                 with open(file, encoding=self.config.encoding, errors=self.config.encoding_errors) as f:
-                    dataset = json.load(f)
+                    dataset = ujson_loads(f.read())
 
                 # We keep only the field we are interested in
                 dataset = dataset[self.config.field]
@@ -142,8 +150,8 @@ def _generate_tables(self, files):
                                 with open(
                                     file, encoding=self.config.encoding, errors=self.config.encoding_errors
                                 ) as f:
-                                    dataset = json.load(f)
-                            except json.JSONDecodeError:
+                                    dataset = ujson_loads(f.read())
+                            except ValueError:
                                 logger.error(f"Failed to read file '{file}' with error {type(e)}: {e}")
                                 raise e
                             # If possible, parse the file as a list of json objects/strings and exit the loop