Update json.py

Author: ArjunJagdale

src/datasets/packaged_modules/json/json.py

@@ -1,7 +1,7 @@
 import io
 import itertools
 from dataclasses import dataclass
-from typing import Optional
+from typing import List, Optional
 
 import pandas as pd
 import pyarrow as pa
@@ -70,12 +70,18 @@ def _info(self):
             raise ValueError("The JSON loader parameter `newlines_in_values` is no longer supported")
         return datasets.DatasetInfo(features=self.config.features)
 
-    def _split_generators(self, dl_manager):
+    def _available_splits(self) -> Optional[List[str]]:
+        return [str(split) for split in self.config.data_files] if isinstance(self.config.data_files, dict) else None
+
+    def _split_generators(self, dl_manager, splits: Optional[List[str]] = None):
         """We handle string, list and dicts in datafiles"""
         if not self.config.data_files:
             raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
         dl_manager.download_config.extract_on_the_fly = True
-        data_files = dl_manager.download_and_extract(self.config.data_files)
+        data_files = self.config.data_files
+        if splits and isinstance(data_files, dict):
+            data_files = {split: data_files[split] for split in splits}
+        data_files = dl_manager.download_and_extract(data_files)
         splits = []
         for split_name, files in data_files.items():
             if isinstance(files, str):