codekingpro/portable-devtools
115k
1import base642import collections3import io4import itertools5import logging6import math7import os8from functools import lru_cache9from itertools import chain10from typing import TYPE_CHECKING, Literal11 12import fsspec.core13from fsspec.spec import AbstractBufferedFile14 15try:16 import ujson as json17except ImportError:18 if not TYPE_CHECKING:19 import json20 21from fsspec.asyn import AsyncFileSystem22from fsspec.callbacks import DEFAULT_CALLBACK23from fsspec.core import filesystem, open, split_protocol24from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper25from fsspec.utils import (26 isfilelike,27 merge_offset_ranges,28 other_paths,29)30 31logger = logging.getLogger("fsspec.reference")32 33 34class ReferenceNotReachable(RuntimeError):35 def __init__(self, reference, target, *args):36 super().__init__(*args)37 self.reference = reference38 self.target = target39 40 def __str__(self):41 return f'Reference "{self.reference}" failed to fetch target {self.target}'42 43 44def _first(d):45 return next(iter(d.values()))46 47 48def _prot_in_references(path, references):49 ref = references.get(path)50 if isinstance(ref, (list, tuple)) and isinstance(ref[0], str):51 return split_protocol(ref[0])[0] if ref[0] else ref[0]52 53 54def _protocol_groups(paths, references):55 if isinstance(paths, str):56 return {_prot_in_references(paths, references): [paths]}57 out = {}58 for path in paths:59 protocol = _prot_in_references(path, references)60 out.setdefault(protocol, []).append(path)61 return out62 63 64class RefsValuesView(collections.abc.ValuesView):65 def __iter__(self):66 for val in self._mapping.zmetadata.values():67 yield json.dumps(val).encode()68 yield from self._mapping._items.values()69 for field in self._mapping.listdir():70 chunk_sizes = self._mapping._get_chunk_sizes(field)71 if len(chunk_sizes) == 0:72 yield self._mapping[field + "/0"]73 continue74 yield from self._mapping._generate_all_records(field)75 76 77class RefsItemsView(collections.abc.ItemsView):78 def __iter__(self):79 return zip(self._mapping.keys(), self._mapping.values())80 81 82def ravel_multi_index(idx, sizes):83 val = 084 mult = 185 for i, s in zip(idx[::-1], sizes[::-1]):86 val += i * mult87 mult *= s88 return val89 90 91class LazyReferenceMapper(collections.abc.MutableMapping):92 """This interface can be used to read/write references from Parquet stores.93 It is not intended for other types of references.94 It can be used with Kerchunk's MultiZarrToZarr method to combine95 references into a parquet store.96 Examples of this use-case can be found here:97 https://fsspec.github.io/kerchunk/advanced.html?highlight=parquet#parquet-storage"""98 99 # import is class level to prevent numpy dep requirement for fsspec100 @property101 def np(self):102 import numpy as np103 104 return np105 106 @property107 def pd(self):108 import pandas as pd109 110 return pd111 112 def __init__(113 self,114 root,115 fs=None,116 out_root=None,117 cache_size=128,118 categorical_threshold=10,119 engine: Literal["fastparquet", "pyarrow"] = "fastparquet",120 ):121 """122 123 This instance will be writable, storing changes in memory until full partitions124 are accumulated or .flush() is called.125 126 To create an empty lazy store, use .create()127 128 Parameters129 ----------130 root : str131 Root of parquet store132 fs : fsspec.AbstractFileSystem133 fsspec filesystem object, default is local filesystem.134 cache_size : int, default=128135 Maximum size of LRU cache, where cache_size*record_size denotes136 the total number of references that can be loaded in memory at once.137 categorical_threshold : int138 Encode urls as pandas.Categorical to reduce memory footprint if the ratio139 of the number of unique urls to total number of refs for each variable140 is greater than or equal to this number. (default 10)141 engine: Literal["fastparquet","pyarrow"]142 Engine choice for reading parquet files. (default is "fastparquet")143 """144 145 self.root = root146 self.chunk_sizes = {}147 self.cat_thresh = categorical_threshold148 self.engine = engine149 self.cache_size = cache_size150 self.url = self.root + "/{field}/refs.{record}.parq"151 # TODO: derive fs from `root`152 self.fs = fsspec.filesystem("file") if fs is None else fs153 self.out_root = self.fs.unstrip_protocol(out_root or self.root)154 155 from importlib.util import find_spec156 157 if self.engine == "pyarrow" and find_spec("pyarrow") is None:158 raise ImportError("engine choice `pyarrow` is not installed.")159 160 # Apply `lru_cache` decorator manually per instance.161 # This way `self` reference is not held on class level.162 # WARNING: However, this means that self and its members are not reflected163 # in the cache key, so we expect they won't be mutated once a value is cached.164 self.listdir = lru_cache()(self.listdir)165 self._key_to_record = lru_cache(maxsize=4096)(self._key_to_record)166 167 def __getattr__(self, item):168 if item in ("_items", "record_size", "zmetadata"):169 self.setup()170 # avoid possible recursion if setup fails somehow171 return self.__dict__[item]172 raise AttributeError(item)173 174 def setup(self):175 self._items = {}176 self._items[".zmetadata"] = self.fs.cat_file(177 "/".join([self.root, ".zmetadata"])178 )179 met = json.loads(self._items[".zmetadata"])180 self.record_size = met["record_size"]181 self.zmetadata = met["metadata"]182 183 # Define function to open and decompress refs184 @lru_cache(maxsize=self.cache_size)185 def open_refs(field, record):186 """cached parquet file loader"""187 path = self.url.format(field=field, record=record)188 data = io.BytesIO(self.fs.cat_file(path))189 try:190 df = self.pd.read_parquet(data, engine=self.engine)191 refs = {c: df[c].to_numpy() for c in df.columns}192 except OSError:193 refs = None194 return refs195 196 self.open_refs = open_refs197 198 @staticmethod199 def create(root, storage_options=None, fs=None, record_size=10000, **kwargs):200 """Make empty parquet reference set201 202 First deletes the contents of the given directory, if it exists.203 204 Parameters205 ----------206 root: str207 Directory to contain the output; will be created208 storage_options: dict | None209 For making the filesystem to use for writing is fs is None210 fs: FileSystem | None211 Filesystem for writing212 record_size: int213 Number of references per parquet file214 kwargs: passed to __init__215 216 Returns217 -------218 LazyReferenceMapper instance219 """220 met = {"metadata": {}, "record_size": record_size}221 if fs is None:222 fs, root = fsspec.core.url_to_fs(root, **(storage_options or {}))223 if fs.exists(root):224 fs.rm(root, recursive=True)225 fs.makedirs(root, exist_ok=True)226 fs.pipe("/".join([root, ".zmetadata"]), json.dumps(met).encode())227 return LazyReferenceMapper(root, fs, **kwargs)228 229 def listdir(self):230 """List top-level directories"""231 dirs = (p.rsplit("/", 1)[0] for p in self.zmetadata if not p.startswith(".z"))232 return set(dirs)233 234 def ls(self, path="", detail=True):235 """Shortcut file listings"""236 path = path.rstrip("/")237 pathdash = path + "/" if path else ""238 dirnames = self.listdir()239 dirs = [240 d241 for d in dirnames242 if d.startswith(pathdash) and "/" not in d.lstrip(pathdash)243 ]244 if dirs:245 others = {246 f247 for f in chain(248 [".zmetadata"],249 (name for name in self.zmetadata),250 (name for name in self._items),251 )252 if f.startswith(pathdash) and "/" not in f.lstrip(pathdash)253 }254 if detail is False:255 others.update(dirs)256 return sorted(others)257 dirinfo = [{"name": name, "type": "directory", "size": 0} for name in dirs]258 fileinfo = [259 {260 "name": name,261 "type": "file",262 "size": len(263 json.dumps(self.zmetadata[name])264 if name in self.zmetadata265 else self._items[name]266 ),267 }268 for name in others269 ]270 return sorted(dirinfo + fileinfo, key=lambda s: s["name"])271 field = path272 others = set(273 [name for name in self.zmetadata if name.startswith(f"{path}/")]274 + [name for name in self._items if name.startswith(f"{path}/")]275 )276 fileinfo = [277 {278 "name": name,279 "type": "file",280 "size": len(281 json.dumps(self.zmetadata[name])282 if name in self.zmetadata283 else self._items[name]284 ),285 }286 for name in others287 ]288 keys = self._keys_in_field(field)289 290 if detail is False:291 return list(others) + list(keys)292 recs = self._generate_all_records(field)293 recinfo = [294 {"name": name, "type": "file", "size": rec[-1]}295 for name, rec in zip(keys, recs)296 if rec[0] # filters out path==None, deleted/missing297 ]298 return fileinfo + recinfo299 300 def _load_one_key(self, key):301 """Get the reference for one key302 303 Returns bytes, one-element list or three-element list.304 """305 if key in self._items:306 return self._items[key]307 elif key in self.zmetadata:308 return json.dumps(self.zmetadata[key]).encode()309 elif "/" not in key or self._is_meta(key):310 raise KeyError(key)311 field, _ = key.rsplit("/", 1)312 record, ri, chunk_size = self._key_to_record(key)313 maybe = self._items.get((field, record), {}).get(ri, False)314 if maybe is None:315 # explicitly deleted316 raise KeyError317 elif maybe:318 return maybe319 elif chunk_size == 0:320 return b""321 322 # Chunk keys can be loaded from row group and cached in LRU cache323 try:324 refs = self.open_refs(field, record)325 except (ValueError, TypeError, FileNotFoundError) as exc:326 raise KeyError(key) from exc327 columns = ["path", "offset", "size", "raw"]328 selection = [refs[c][ri] if c in refs else None for c in columns]329 raw = selection[-1]330 if raw is not None:331 return raw332 if selection[0] is None:333 raise KeyError("This reference does not exist or has been deleted")334 if selection[1:3] == [0, 0]:335 # URL only336 return selection[:1]337 # URL, offset, size338 return selection[:3]339 340 def _key_to_record(self, key):341 """Details needed to construct a reference for one key"""342 field, chunk = key.rsplit("/", 1)343 chunk_sizes = self._get_chunk_sizes(field)344 if len(chunk_sizes) == 0:345 return 0, 0, 0346 chunk_idx = [int(c) for c in chunk.split(".")]347 chunk_number = ravel_multi_index(chunk_idx, chunk_sizes)348 record = chunk_number // self.record_size349 ri = chunk_number % self.record_size350 return record, ri, len(chunk_sizes)351 352 def _get_chunk_sizes(self, field):353 """The number of chunks along each axis for a given field"""354 if field not in self.chunk_sizes:355 zarray = self.zmetadata[f"{field}/.zarray"]356 size_ratio = [357 math.ceil(s / c) for s, c in zip(zarray["shape"], zarray["chunks"])358 ]359 self.chunk_sizes[field] = size_ratio or [1]360 return self.chunk_sizes[field]361 362 def _generate_record(self, field, record):363 """The references for a given parquet file of a given field"""364 refs = self.open_refs(field, record)365 it = iter(zip(*refs.values()))366 if len(refs) == 3:367 # All urls368 return (list(t) for t in it)369 elif len(refs) == 1:370 # All raws371 return refs["raw"]372 else:373 # Mix of urls and raws374 return (list(t[:3]) if not t[3] else t[3] for t in it)375 376 def _generate_all_records(self, field):377 """Load all the references within a field by iterating over the parquet files"""378 nrec = 1379 for ch in self._get_chunk_sizes(field):380 nrec *= ch381 nrec = math.ceil(nrec / self.record_size)382 for record in range(nrec):383 yield from self._generate_record(field, record)384 385 def values(self):386 return RefsValuesView(self)387 388 def items(self):389 return RefsItemsView(self)390 391 def __hash__(self):392 return id(self)393 394 def __getitem__(self, key):395 return self._load_one_key(key)396 397 def __setitem__(self, key, value):398 if "/" in key and not self._is_meta(key):399 field, chunk = key.rsplit("/", 1)400 record, i, _ = self._key_to_record(key)401 subdict = self._items.setdefault((field, record), {})402 subdict[i] = value403 if len(subdict) == self.record_size:404 self.write(field, record)405 else:406 # metadata or top-level407 if hasattr(value, "to_bytes"):408 val = value.to_bytes().decode()409 elif isinstance(value, bytes):410 val = value.decode()411 else:412 val = value413 self._items[key] = val414 new_value = json.loads(val)415 self.zmetadata[key] = {**self.zmetadata.get(key, {}), **new_value}416 417 @staticmethod418 def _is_meta(key):419 return key.startswith(".z") or "/.z" in key420 421 def __delitem__(self, key):422 if key in self._items:423 del self._items[key]424 elif key in self.zmetadata:425 del self.zmetadata[key]426 else:427 if "/" in key and not self._is_meta(key):428 field, _ = key.rsplit("/", 1)429 record, i, _ = self._key_to_record(key)430 subdict = self._items.setdefault((field, record), {})431 subdict[i] = None432 if len(subdict) == self.record_size:433 self.write(field, record)434 else:435 # metadata or top-level436 self._items[key] = None437 438 def write(self, field, record, base_url=None, storage_options=None):439 # extra requirements if writing440 import kerchunk.df441 import numpy as np442 import pandas as pd443 444 partition = self._items[(field, record)]445 original = False446 if len(partition) < self.record_size:447 try:448 original = self.open_refs(field, record)449 except OSError:450 pass451 452 if original:453 paths = original["path"]454 offsets = original["offset"]455 sizes = original["size"]456 raws = original["raw"]457 else:458 paths = np.full(self.record_size, np.nan, dtype="O")459 offsets = np.zeros(self.record_size, dtype="int64")460 sizes = np.zeros(self.record_size, dtype="int64")461 raws = np.full(self.record_size, np.nan, dtype="O")462 for j, data in partition.items():463 if isinstance(data, list):464 if (465 str(paths.dtype) == "category"466 and data[0] not in paths.dtype.categories467 ):468 paths = paths.add_categories(data[0])469 paths[j] = data[0]470 if len(data) > 1:471 offsets[j] = data[1]472 sizes[j] = data[2]473 elif data is None:474 # delete475 paths[j] = None476 offsets[j] = 0477 sizes[j] = 0478 raws[j] = None479 else:480 # this is the only call into kerchunk, could remove481 raws[j] = kerchunk.df._proc_raw(data)482 # TODO: only save needed columns483 df = pd.DataFrame(484 {485 "path": paths,486 "offset": offsets,487 "size": sizes,488 "raw": raws,489 },490 copy=False,491 )492 if df.path.count() / (df.path.nunique() or 1) > self.cat_thresh:493 df["path"] = df["path"].astype("category")494 object_encoding = {"raw": "bytes", "path": "utf8"}495 has_nulls = ["path", "raw"]496 497 fn = f"{base_url or self.out_root}/{field}/refs.{record}.parq"498 self.fs.mkdirs(f"{base_url or self.out_root}/{field}", exist_ok=True)499 500 if self.engine == "pyarrow":501 df_backend_kwargs = {"write_statistics": False}502 elif self.engine == "fastparquet":503 df_backend_kwargs = {504 "stats": False,505 "object_encoding": object_encoding,506 "has_nulls": has_nulls,507 }508 else:509 raise NotImplementedError(f"{self.engine} not supported")510 df.to_parquet(511 fn,512 engine=self.engine,513 storage_options=storage_options514 or getattr(self.fs, "storage_options", None),515 compression="zstd",516 index=False,517 **df_backend_kwargs,518 )519 520 partition.clear()521 self._items.pop((field, record))522 523 def flush(self, base_url=None, storage_options=None):524 """Output any modified or deleted keys525 526 Parameters527 ----------528 base_url: str529 Location of the output530 """531 532 # write what we have so far and clear sub chunks533 for thing in list(self._items):534 if isinstance(thing, tuple):535 field, record = thing536 self.write(537 field,538 record,539 base_url=base_url,540 storage_options=storage_options,541 )542 543 # gather .zmetadata from self._items and write that too544 for k in list(self._items):545 if k != ".zmetadata" and ".z" in k:546 self.zmetadata[k] = json.loads(self._items.pop(k))547 met = {"metadata": self.zmetadata, "record_size": self.record_size}548 self._items.clear()549 self._items[".zmetadata"] = json.dumps(met).encode()550 self.fs.pipe(551 "/".join([base_url or self.out_root, ".zmetadata"]),552 self._items[".zmetadata"],553 )554 555 # TODO: only clear those that we wrote to?556 self.open_refs.cache_clear()557 558 def __len__(self):559 # Caveat: This counts expected references, not actual - but is fast560 count = 0561 for field in self.listdir():562 if field.startswith("."):563 count += 1564 else:565 count += math.prod(self._get_chunk_sizes(field))566 count += len(self.zmetadata) # all metadata keys567 # any other files not in reference partitions568 count += sum(1 for _ in self._items if not isinstance(_, tuple))569 return count570 571 def __iter__(self):572 # Caveat: returns only existing keys, so the number of these does not573 # match len(self)574 metas = set(self.zmetadata)575 metas.update(self._items)576 for bit in metas:577 if isinstance(bit, str):578 yield bit579 for field in self.listdir():580 for k in self._keys_in_field(field):581 if k in self:582 yield k583 584 def __contains__(self, item):585 try:586 self._load_one_key(item)587 return True588 except KeyError:589 return False590 591 def _keys_in_field(self, field):592 """List key names in given field593 594 Produces strings like "field/x.y" appropriate from the chunking of the array595 """596 chunk_sizes = self._get_chunk_sizes(field)597 if len(chunk_sizes) == 0:598 yield field + "/0"599 return600 inds = itertools.product(*(range(i) for i in chunk_sizes))601 for ind in inds:602 yield field + "/" + ".".join([str(c) for c in ind])603 604 605class ReferenceFileSystem(AsyncFileSystem):606 """View byte ranges of some other file as a file system607 Initial version: single file system target, which must support608 async, and must allow start and end args in _cat_file. Later versions609 may allow multiple arbitrary URLs for the targets.610 This FileSystem is read-only. It is designed to be used with async611 targets (for now). We do not get original file details from the target FS.612 Configuration is by passing a dict of references at init, or a URL to613 a JSON file containing the same; this dict614 can also contain concrete data for some set of paths.615 Reference dict format:616 {path0: bytes_data, path1: (target_url, offset, size)}617 https://github.com/fsspec/kerchunk/blob/main/README.md618 """619 620 protocol = "reference"621 cachable = False622 623 def __init__(624 self,625 fo,626 target=None,627 ref_storage_args=None,628 target_protocol=None,629 target_options=None,630 remote_protocol=None,631 remote_options=None,632 fs=None,633 template_overrides=None,634 simple_templates=True,635 max_gap=64_000,636 max_block=256_000_000,637 cache_size=128,638 **kwargs,639 ):640 """641 Parameters642 ----------643 fo : dict or str644 The set of references to use for this instance, with a structure as above.645 If str referencing a JSON file, will use fsspec.open, in conjunction646 with target_options and target_protocol to open and parse JSON at this647 location. If a directory, then assume references are a set of parquet648 files to be loaded lazily.649 target : str650 For any references having target_url as None, this is the default file651 target to use652 ref_storage_args : dict653 If references is a str, use these kwargs for loading the JSON file.654 Deprecated: use target_options instead.655 target_protocol : str656 Used for loading the reference file, if it is a path. If None, protocol657 will be derived from the given path658 target_options : dict659 Extra FS options for loading the reference file ``fo``, if given as a path660 remote_protocol : str661 The protocol of the filesystem on which the references will be evaluated662 (unless fs is provided). If not given, will be derived from the first663 URL that has a protocol in the templates or in the references, in that664 order.665 remote_options : dict666 kwargs to go with remote_protocol667 fs : AbstractFileSystem | dict(str, (AbstractFileSystem | dict))668 Directly provide a file system(s):669 - a single filesystem instance670 - a dict of protocol:filesystem, where each value is either a filesystem671 instance, or a dict of kwargs that can be used to create in672 instance for the given protocol673 674 If this is given, remote_options and remote_protocol are ignored.675 template_overrides : dict676 Swap out any templates in the references file with these - useful for677 testing.678 simple_templates: bool679 Whether templates can be processed with simple replace (True) or if680 jinja is needed (False, much slower). All reference sets produced by681 ``kerchunk`` are simple in this sense, but the spec allows for complex.682 max_gap, max_block: int683 For merging multiple concurrent requests to the same remote file.684 Neighboring byte ranges will only be merged when their685 inter-range gap is <= ``max_gap``. Default is 64KB. Set to 0686 to only merge when it requires no extra bytes. Pass a negative687 number to disable merging, appropriate for local target files.688 Neighboring byte ranges will only be merged when the size of689 the aggregated range is <= ``max_block``. Default is 256MB.690 cache_size : int691 Maximum size of LRU cache, where cache_size*record_size denotes692 the total number of references that can be loaded in memory at once.693 Only used for lazily loaded references.694 kwargs : passed to parent class695 """696 super().__init__(**kwargs)697 self.target = target698 self.template_overrides = template_overrides699 self.simple_templates = simple_templates700 self.templates = {}701 self.fss = {}702 self._dircache = {}703 self.max_gap = max_gap704 self.max_block = max_block705 if isinstance(fo, str):706 dic = dict(707 **(ref_storage_args or target_options or {}), protocol=target_protocol708 )709 ref_fs, fo2 = fsspec.core.url_to_fs(fo, **dic)710 if ".json" not in fo2 and (711 fo.endswith(("parq", "parquet", "/")) or ref_fs.isdir(fo2)712 ):713 # Lazy parquet refs714 logger.info("Open lazy reference dict from URL %s", fo)715 self.references = LazyReferenceMapper(716 fo2,717 fs=ref_fs,718 cache_size=cache_size,719 )720 else:721 # text JSON722 with fsspec.open(fo, "rb", **dic) as f:723 logger.info("Read reference from URL %s", fo)724 text = json.load(f)725 self._process_references(text, template_overrides)726 else:727 # dictionaries728 self._process_references(fo, template_overrides)729 if isinstance(fs, dict):730 self.fss = {731 k: (732 fsspec.filesystem(k.split(":", 1)[0], **opts)733 if isinstance(opts, dict)734 else opts735 )736 for k, opts in fs.items()737 }738 if None not in self.fss:739 self.fss[None] = filesystem("file")740 return741 if fs is not None:742 # single remote FS743 remote_protocol = (744 fs.protocol[0] if isinstance(fs.protocol, tuple) else fs.protocol745 )746 self.fss[remote_protocol] = fs747 748 if remote_protocol is None:749 # get single protocol from any templates750 for ref in self.templates.values():751 if callable(ref):752 ref = ref()753 protocol, _ = fsspec.core.split_protocol(ref)754 if protocol and protocol not in self.fss:755 fs = filesystem(protocol, **(remote_options or {}))756 self.fss[protocol] = fs757 if remote_protocol is None:758 # get single protocol from references759 # TODO: warning here, since this can be very expensive?760 for ref in self.references.values():761 if callable(ref):762 ref = ref()763 if isinstance(ref, list) and ref[0]:764 protocol, _ = fsspec.core.split_protocol(ref[0])765 if protocol not in self.fss:766 fs = filesystem(protocol, **(remote_options or {}))767 self.fss[protocol] = fs768 # only use first remote URL769 break770 771 if remote_protocol and remote_protocol not in self.fss:772 fs = filesystem(remote_protocol, **(remote_options or {}))773 self.fss[remote_protocol] = fs774 775 self.fss[None] = fs or filesystem("file") # default one776 # Wrap any non-async filesystems to ensure async methods are available below777 for k, f in self.fss.items():778 if not f.async_impl:779 self.fss[k] = AsyncFileSystemWrapper(f, asynchronous=self.asynchronous)780 elif self.asynchronous ^ f.asynchronous:781 raise ValueError(782 "Reference-FS's target filesystem must have same value "783 "of asynchronous"784 )785 786 def _cat_common(self, path, start=None, end=None):787 path = self._strip_protocol(path)788 logger.debug(f"cat: {path}")789 try:790 part = self.references[path]791 except KeyError as exc:792 raise FileNotFoundError(path) from exc793 if isinstance(part, str):794 part = part.encode()795 if hasattr(part, "to_bytes"):796 part = part.to_bytes()797 if isinstance(part, bytes):798 logger.debug(f"Reference: {path}, type bytes")799 if part.startswith(b"base64:"):800 part = base64.b64decode(part[7:])801 return part, None, None802 803 if len(part) == 1:804 logger.debug(f"Reference: {path}, whole file => {part}")805 url = part[0]806 start1, end1 = start, end807 else:808 url, start0, size = part809 logger.debug(f"Reference: {path} => {url}, offset {start0}, size {size}")810 end0 = start0 + size811 812 if start is not None:813 if start >= 0:814 start1 = start0 + start815 else:816 start1 = end0 + start817 else:818 start1 = start0819 if end is not None:820 if end >= 0:821 end1 = start0 + end822 else:823 end1 = end0 + end824 else:825 end1 = end0826 if url is None:827 url = self.target828 return url, start1, end1829 830 async def _cat_file(self, path, start=None, end=None, **kwargs):831 part_or_url, start0, end0 = self._cat_common(path, start=start, end=end)832 if isinstance(part_or_url, bytes):833 return part_or_url[start:end]834 protocol, _ = split_protocol(part_or_url)835 try:836 return await self.fss[protocol]._cat_file(837 part_or_url, start=start0, end=end0838 )839 except Exception as e:840 raise ReferenceNotReachable(path, part_or_url) from e841 842 def cat_file(self, path, start=None, end=None, **kwargs):843 part_or_url, start0, end0 = self._cat_common(path, start=start, end=end)844 if isinstance(part_or_url, bytes):845 return part_or_url[start:end]846 protocol, _ = split_protocol(part_or_url)847 try:848 return self.fss[protocol].cat_file(part_or_url, start=start0, end=end0)849 except Exception as e:850 raise ReferenceNotReachable(path, part_or_url) from e851 852 def pipe_file(self, path, value, **_):853 """Temporarily add binary data or reference as a file"""854 self.references[path] = value855 856 async def _get_file(self, rpath, lpath, **kwargs):857 if self.isdir(rpath):858 return os.makedirs(lpath, exist_ok=True)859 data = await self._cat_file(rpath)860 with open(lpath, "wb") as f:861 f.write(data)862 863 def get_file(self, rpath, lpath, callback=DEFAULT_CALLBACK, **kwargs):864 if self.isdir(rpath):865 return os.makedirs(lpath, exist_ok=True)866 data = self.cat_file(rpath, **kwargs)867 callback.set_size(len(data))868 if isfilelike(lpath):869 lpath.write(data)870 else:871 with open(lpath, "wb") as f:872 f.write(data)873 callback.absolute_update(len(data))874 875 def get(self, rpath, lpath, recursive=False, **kwargs):876 if recursive:877 # trigger directory build878 self.ls("")879 rpath = self.expand_path(rpath, recursive=recursive)880 fs = fsspec.filesystem("file", auto_mkdir=True)881 targets = other_paths(rpath, lpath)882 if recursive:883 data = self.cat([r for r in rpath if not self.isdir(r)])884 else:885 data = self.cat(rpath)886 for remote, local in zip(rpath, targets):887 if remote in data:888 fs.pipe_file(local, data[remote])889 890 def cat(self, path, recursive=False, on_error="raise", **kwargs):891 if isinstance(path, str) and recursive:892 raise NotImplementedError893 if isinstance(path, list) and (recursive or any("*" in p for p in path)):894 raise NotImplementedError895 # TODO: if references is lazy, pre-fetch all paths in batch before access896 proto_dict = _protocol_groups(path, self.references)897 out = {}898 for proto, paths in proto_dict.items():899 fs = self.fss[proto]900 urls, starts, ends, valid_paths = [], [], [], []901 for p in paths:902 # find references or label not-found. Early exit if any not903 # found and on_error is "raise"904 try:905 u, s, e = self._cat_common(p)906 if not isinstance(u, (bytes, str)):907 # nan/None from parquet908 continue909 except FileNotFoundError as err:910 if on_error == "raise":911 raise912 if on_error != "omit":913 out[p] = err914 else:915 urls.append(u)916 starts.append(s)917 ends.append(e)918 valid_paths.append(p)919 920 # process references into form for merging921 urls2 = []922 starts2 = []923 ends2 = []924 paths2 = []925 whole_files = set()926 for u, s, e, p in zip(urls, starts, ends, valid_paths):927 if isinstance(u, bytes):928 # data929 out[p] = u930 elif s is None:931 # whole file - limits are None, None, but no further932 # entries take for this file933 whole_files.add(u)934 urls2.append(u)935 starts2.append(s)936 ends2.append(e)937 paths2.append(p)938 for u, s, e, p in zip(urls, starts, ends, valid_paths):939 # second run to account for files that are to be loaded whole940 if s is not None and u not in whole_files:941 urls2.append(u)942 starts2.append(s)943 ends2.append(e)944 paths2.append(p)945 946 # merge and fetch consolidated ranges947 new_paths, new_starts, new_ends = merge_offset_ranges(948 list(urls2),949 list(starts2),950 list(ends2),951 sort=True,952 max_gap=self.max_gap,953 max_block=self.max_block,954 )955 bytes_out = fs.cat_ranges(new_paths, new_starts, new_ends)956 957 # unbundle from merged bytes - simple approach958 for u, s, e, p in zip(urls, starts, ends, valid_paths):959 if p in out:960 continue # was bytes, already handled961 for np, ns, ne, b in zip(new_paths, new_starts, new_ends, bytes_out):962 if np == u and (ns is None or ne is None):963 if isinstance(b, Exception):964 out[p] = b965 else:966 out[p] = b[s:e]967 elif np == u and s >= ns and e <= ne:968 if isinstance(b, Exception):969 out[p] = b970 else:971 out[p] = b[s - ns : (e - ne) or None]972 973 for k, v in out.copy().items():974 # these were valid references, but fetch failed, so transform exc975 if isinstance(v, Exception) and k in self.references:976 ex = out[k]977 new_ex = ReferenceNotReachable(k, self.references[k])978 new_ex.__cause__ = ex979 if on_error == "raise":980 raise new_ex981 elif on_error != "omit":982 out[k] = new_ex983 984 if len(out) == 1 and isinstance(path, str) and "*" not in path:985 return _first(out)986 return out987 988 def _process_references(self, references, template_overrides=None):989 vers = references.get("version", None)990 if vers is None:991 self._process_references0(references)992 elif vers == 1:993 self._process_references1(references, template_overrides=template_overrides)994 else:995 raise ValueError(f"Unknown reference spec version: {vers}")996 # TODO: we make dircache by iterating over all entries, but for Spec >= 1,997 # can replace with programmatic. Is it even needed for mapper interface?998 999 def _process_references0(self, references):1000 """Make reference dict for Spec Version 0"""1001 if isinstance(references, dict):1002 # do not do this for lazy/parquet backend, which will not make dicts,1003 # but must remain writable in the original object1004 references = {1005 key: json.dumps(val) if isinstance(val, dict) else val1006 for key, val in references.items()1007 }1008 self.references = references1009 1010 def _process_references1(self, references, template_overrides=None):1011 if not self.simple_templates or self.templates:1012 import jinja21013 self.references = {}1014 self._process_templates(references.get("templates", {}))1015 1016 @lru_cache(1000)1017 def _render_jinja(u):1018 return jinja2.Template(u).render(**self.templates)1019 1020 for k, v in references.get("refs", {}).items():1021 if isinstance(v, str):1022 if v.startswith("base64:"):1023 self.references[k] = base64.b64decode(v[7:])1024 self.references[k] = v1025 elif isinstance(v, dict):1026 self.references[k] = json.dumps(v)1027 elif self.templates:1028 u = v[0]1029 if "{{" in u:1030 if self.simple_templates:1031 u = (1032 u.replace("{{", "{")1033 .replace("}}", "}")1034 .format(**self.templates)1035 )1036 else:1037 u = _render_jinja(u)1038 self.references[k] = [u] if len(v) == 1 else [u, v[1], v[2]]1039 else:1040 self.references[k] = v1041 self.references.update(self._process_gen(references.get("gen", [])))1042 1043 def _process_templates(self, tmp):1044 self.templates = {}1045 if self.template_overrides is not None:1046 tmp.update(self.template_overrides)1047 for k, v in tmp.items():1048 if "{{" in v:1049 import jinja21050 1051 self.templates[k] = lambda temp=v, **kwargs: jinja2.Template(1052 temp1053 ).render(**kwargs)1054 else:1055 self.templates[k] = v1056 1057 def _process_gen(self, gens):1058 out = {}1059 for gen in gens:1060 dimension = {1061 k: (1062 v1063 if isinstance(v, list)1064 else range(v.get("start", 0), v["stop"], v.get("step", 1))1065 )1066 for k, v in gen["dimensions"].items()1067 }1068 products = (1069 dict(zip(dimension.keys(), values))1070 for values in itertools.product(*dimension.values())1071 )1072 for pr in products:1073 import jinja21074 1075 key = jinja2.Template(gen["key"]).render(**pr, **self.templates)1076 url = jinja2.Template(gen["url"]).render(**pr, **self.templates)1077 if ("offset" in gen) and ("length" in gen):1078 offset = int(1079 jinja2.Template(gen["offset"]).render(**pr, **self.templates)1080 )1081 length = int(1082 jinja2.Template(gen["length"]).render(**pr, **self.templates)1083 )1084 out[key] = [url, offset, length]1085 elif ("offset" in gen) ^ ("length" in gen):1086 raise ValueError(1087 "Both 'offset' and 'length' are required for a "1088 "reference generator entry if either is provided."1089 )1090 else:1091 out[key] = [url]1092 return out1093 1094 def _dircache_from_items(self):1095 self.dircache = {"": []}1096 it = self.references.items()1097 for path, part in it:1098 if isinstance(part, (bytes, str)) or hasattr(part, "to_bytes"):1099 size = len(part)1100 elif len(part) == 1:1101 size = None1102 else:1103 _, _, size = part1104 par = path.rsplit("/", 1)[0] if "/" in path else ""1105 par0 = par1106 subdirs = [par0]1107 while par0 and par0 not in self.dircache:1108 # collect parent directories1109 par0 = self._parent(par0)1110 subdirs.append(par0)1111 1112 subdirs.reverse()1113 for parent, child in zip(subdirs, subdirs[1:]):1114 # register newly discovered directories1115 assert child not in self.dircache1116 assert parent in self.dircache1117 self.dircache[parent].append(1118 {"name": child, "type": "directory", "size": 0}1119 )1120 self.dircache[child] = []1121 1122 self.dircache[par].append({"name": path, "type": "file", "size": size})1123 1124 def _open(self, path, mode="rb", block_size=None, cache_options=None, **kwargs):1125 part_or_url, start0, end0 = self._cat_common(path)1126 # This logic is kept outside `ReferenceFile` to avoid unnecessary redirection.1127 # That does mean `_cat_common` gets called twice if it eventually reaches `ReferenceFile`.1128 if isinstance(part_or_url, bytes):1129 return io.BytesIO(part_or_url[start0:end0])1130 1131 protocol, _ = split_protocol(part_or_url)1132 if start0 is None and end0 is None:1133 return self.fss[protocol]._open(1134 part_or_url,1135 mode,1136 block_size=block_size,1137 cache_options=cache_options,1138 **kwargs,1139 )1140 1141 return ReferenceFile(1142 self,1143 path,1144 mode,1145 block_size=block_size,1146 cache_options=cache_options,1147 **kwargs,1148 )1149 1150 def ls(self, path, detail=True, **kwargs):1151 logger.debug("list %s", path)1152 path = self._strip_protocol(path)1153 if isinstance(self.references, LazyReferenceMapper):1154 try:1155 return self.references.ls(path, detail)1156 except KeyError:1157 pass1158 raise FileNotFoundError(f"'{path}' is not a known key")1159 if not self.dircache:1160 self._dircache_from_items()1161 out = self._ls_from_cache(path)1162 if out is None:1163 raise FileNotFoundError(path)1164 if detail:1165 return out1166 return [o["name"] for o in out]1167 1168 def exists(self, path, **kwargs): # overwrite auto-sync version1169 return self.isdir(path) or self.isfile(path)1170 1171 def isdir(self, path): # overwrite auto-sync version1172 if self.dircache:1173 return path in self.dircache1174 elif isinstance(self.references, LazyReferenceMapper):1175 return path in self.references.listdir()1176 else:1177 # this may be faster than building dircache for single calls, but1178 # by looping will be slow for many calls; could cache it?1179 return any(_.startswith(f"{path}/") for _ in self.references)1180 1181 def isfile(self, path): # overwrite auto-sync version1182 return path in self.references1183 1184 async def _ls(self, path, detail=True, **kwargs): # calls fast sync code1185 return self.ls(path, detail, **kwargs)1186 1187 def find(self, path, maxdepth=None, withdirs=False, detail=False, **kwargs):1188 if withdirs:1189 return super().find(1190 path, maxdepth=maxdepth, withdirs=withdirs, detail=detail, **kwargs1191 )1192 if path:1193 path = self._strip_protocol(path)1194 r = sorted(k for k in self.references if k.startswith(path))1195 else:1196 r = sorted(self.references)1197 if detail:1198 if not self.dircache:1199 self._dircache_from_items()1200 return {k: self._ls_from_cache(k)[0] for k in r}