Source code for kerchunk.df

import base64
import logging

import ujson
import pandas as pd
import fsspec
import zarr

# example from preffs's README'
df = pd.DataFrame(
        "key": ["a/b", "a/b", "b"],
        "path": ["a.dat", "b.dat", None],
        "offset": [123, 12, 0],
        "size": [12, 17, 0],
        "raw": [None, None, b"data"],

logger = logging.getLogger("kerchunk.df")

def _proc_raw(r):
    if not isinstance(r, bytes):
        r = r.encode()
    if r.startswith(b"base64:"):
        return base64.b64decode(r[7:])
    return r

def get_variables(refs, consolidated=True):
    """Get list of variable names from references.

    Finds the top-level prefixes in a reference set, corresponding to
    the directory listing of the root for zarr.

    refs : dict
        kerchunk references keys
    consolidated : bool
        Whether or not to add consolidated metadata key to references. (default True)

    fields : list of str
        List of variable names.
    fields = []
    meta = {}
    for k in refs:
        if ".z" in k and consolidated:
            meta[k] = refs[k]
        if "/" in k:
            name, chunk = k.split("/")
            if name not in fields:
    if consolidated and ".zmetadata" not in fields:
        refs[".zmetadata"] = ujson.loads(meta[".zmetadata"])
    return fields

def _normalize_json(json_obj):
    """Normalize json representation as bytes

    json_obj : str, bytes, dict, list
        JSON data for parquet file to be written.
    if not isinstance(json_obj, str) and not isinstance(json_obj, bytes):
        json_obj = ujson.dumps(json_obj)
    if not isinstance(json_obj, bytes):
        json_obj = json_obj.encode()
    return json_obj

def _write_json(fname, json_obj, storage_options=None):
    """Write references into a parquet file.

    fname : str
        Output filename.
    json_obj : str, bytes, dict, list
        JSON data for parquet file to be written.
    storage_options = {} if storage_options is None else storage_options
    json_obj = _normalize_json(json_obj)
    with, "wb", **storage_options) as f:

[docs]def refs_to_dataframe( fo, url, target_protocol=None, target_options=None, storage_options=None, record_size=100_000, categorical_threshold=10, ): """Write references as a parquet files store. The directory structure should mimic a normal zarr store but instead of standard chunk keys, references are saved as parquet dataframes. Parameters ---------- fo : str | dict Location of a JSON file containing references or a reference set already loaded into memory. url : str Location for the output, together with protocol. This must be a writable directory. target_protocol : str Used for loading the reference file, if it is a path. If None, protocol will be derived from the given path target_options : dict Extra FS options for loading the reference file ``fo``, if given as a path storage_options: dict | None Passed to fsspec when for writing the parquet. record_size : int Number of references to store in each reference file (default 10000). Bigger values mean fewer read requests but larger memory footprint. categorical_threshold : int Encode urls as pandas.Categorical to reduce memory footprint if the ratio of the number of unique urls to total number of refs for each variable is greater than or equal to this number. (default 10) """ from fsspec.implementations.reference import LazyReferenceMapper if isinstance(fo, str): # JSON file dic = dict(**(target_options or {}), protocol=target_protocol) with, "rb", **dic) as f:"Read reference from URL %s", fo) refs = ujson.load(f) else: # Mapping object refs = fo if "refs" in refs: refs = refs["refs"] fs, _ = fsspec.core.url_to_fs(url, **(storage_options or {})) out = LazyReferenceMapper.create( record_size=record_size, root=url, fs=fs, categorical_threshold=categorical_threshold, ) for k in sorted(refs): out[k] = refs[k] out.flush()