Replace global schema state with schema adapter dataclass

[effigies]

This repository's current structure makes it awkward to work with multiple layouts in a single process. In the fMRIPrep use case, we need the following layouts:

1. templateflow: Not quite BIDS
2. Input dataset (BIDS, but look for some extra files for backwards compatibility reasons)
3. Output dataset: Utilizes unmerged BEPs

To me, the natural way to do this would be something like:

tf_schema: Namespace = patch_schema(...)
tf_tab = index_dataset('~/.cache/templateflow', schema=tf_schema)

ds_tab = index_dataset(input_path)

deriv_tab = index_dataset(output_path, schema='https://bids-specification.readthedocs.io/en/bepXYZ/schema.json')

I think at the moment, we will need to do something like

tf_schema = patch_schema(...)
tf_schema_path = Path(tempfile.NamedTemporaryFile(suffix='.json', delete=False))
tf_schema_path.write_text(tf_schema.to_json())
set_bids_schema(tf_schema_path)
tf_tab = index_dataset('~/.cache/templateflow')
set_bids_schema()  # Reset

How would you feel about abstracting this global state into an object like:

class Schema:
    # Default, can be given a URL or path
    schema_path: PathLike | None
    # Alternate: Can be passed an instantiated schema
    bids_schema: Namespace

    def set_bids_schema(path: str | Path | None = None) -> None: ...
    def get_bids_schema() -> Namespace: ...
    def get_bids_entity_arrow_schema() -> pa.Schema: ...

You could then preserve the current API with:

_GLOBAL_SCHEMA = Schema()
set_bids_schema = _GLOBAL_SCHEMA.set_bids_schema
get_bids_schema = _GLOBAL_SCHEMA.get_bids_schema
get_bids_entity_arrow_schema = _GLOBAL_SCHEMA.get_bids_entity_arrow_schema

Functions that call any of these functions somewhere in their call stack (particularly index_dataset) could be given a schema=_GLOBAL_SCHEMA default argument.

LMK what you think and I'd be happy to put together a PR.

[effigies]

After a brainstorming session with Claude, instead of encapsulating global state, I'm now thinking that we want something that knows how to load various schema sources and always provides access to a pa.Schema:

@dataclass(frozen=True)
class BIDSSchema:
    arrow_schema: pa.Schema
    _entity_schema: dict[str, dict[str, Any]]
    _name_entity_map: dict[str, str]
    _source: str | Path | None = None

    @cached_property
    def bids_schema(self) -> Namespace | None:
        """Lazily reload the bidsschematools Namespace from `_source`.

        Returns None if the instance was constructed directly from a
        `pa.Schema` (no source available).
        """

    @classmethod
    def from_path(cls, path: str | Path | None) -> "BIDSSchema":
        """Load via bidsschematools.schema.load_schema(path)."""

    @classmethod
    def from_namespace(cls, ns: Namespace) -> "BIDSSchema":
        """Build from an already-loaded Namespace."""

    @classmethod
    def from_arrow(cls, arrow_schema: pa.Schema) -> "BIDSSchema":
        """Reconstruct from a pa.Schema previously emitted by bids2table.

        `bids_schema` will return None for instances built this way.
        """

    @classmethod
    def prepare(cls, obj) -> "BIDSSchema":
        """Polymorphic constructor.

        Accepts: BIDSSchema | pa.Schema | Namespace | str | Path | None.
        Raises TypeError for any other type.

For backwards compatiblity, we set a default BIDSSchema, but instead of mutating its contents on set_bids_schema, we just swap it out with a new immutable object:

_DEFAULT_SCHEMA: BIDSSchema = BIDSSchema.from_path(None)

def set_bids_schema(path: str | Path | None = None) -> None:
    global _DEFAULT_SCHEMA
    _DEFAULT_SCHEMA = BIDSSchema.from_path(path)

def get_bids_schema() -> Namespace | None:
    return _DEFAULT_SCHEMA.bids_schema

def get_bids_entity_arrow_schema() -> pa.Schema:
    return _DEFAULT_SCHEMA.arrow_schema

The following functions can be passed anything that can be passed to BIDSSchema.prepare():

• validate_bids_entities
• get_arrow_schema
• get_column_names
• index_dataset
• batch_index_dataset

Only arrow_schema is passed to subprocesses, to avoid any overhead of pickling namespaces.

[gkiar]

@kaitj @nx10 let me know if you have any issue with this; I think this makes sense and is a subtle change from what we currently do.

[kaitj]

I think this makes sense too and would be a little more flexible.

[effigies]

After discussion with @nx10, there is no need to preserve backwards compatibility here, so I would like to make a more sweeping proposal that I think is fundamentally simpler.

The publicly visible piece here is a SchemaSpec, which is either a loaded schema or something bidsschematools can load:

type SchemaSpec = Namespace | PathT | str | None

From here, bids2table internally creates an adapter which is a dataclass with only the components that bids2table cares about. This can be used to derive a pyarrow.Schema as needed:

flowchart LR
  SchemaSpec --> BIDSSchemaAdapter --> pyarrow.Schema

Loading

Finally, all public functions (where it could be relevant) accept SchemaSpecs. Internal functions will generally accept adapters, and some may work with only pyarrow schemas. The call tree should generally have the form:

flowchart LR
  public["public(schema)"] --> internal["_private(adapter)"] --> pyarrow["_pyarrow(arrow_schema)"]
  public --> pyarrow

Loading

More detail follows.

1. SchemaAdapter class

This is just a class that extracts the things bids2table cares about.

@dataclass(frozen=True)
class BIDSSchemaAdapter:
    bids_version: str
    schema_version: str
    entity_schema: dict[str, dict[str, Any]]  # Corresponds to _BIDS_ENTITY_SCHEMA

From here, we can now compute what are currently _BIDS_NAME_ENTITY_MAP and _BIDS_ENTITY_ARROW_SCHEMA with no further reference to the original schema:

class BIDSSchemaAdapter:
    ...
    @property
    def name_entity_map(self) -> dict[str, str]: ...
    @property
    def entity_arrow_schema(self) -> pa.Schema: ...

The primary ingestion point would be a classmethod or a function, which takes either a loaded schema (Namespace) or something bidsschematools can use to load a schema:

class BIDSSchemaAdapter:
    ...
    @classmethod
    def load(cls, spec: SchemaSpec = None) -> BIDSSchemaAdapter:
        schema: Namespace = spec if isinstance(spec, Namespace) else bst.schema.load(spec)
        return cls(
            bids_version=schema.bids_version,
            schema_version=schema.schema_version,
            entity_schema={ent: schema.objects.entities[ent] for ent in schema.rules.entities},
        )

As additional schema components become useful to extract and transform, this class will grow.

I am used to using classmethods and properties instead of factories and functions, but given that bids2table is overall more functional/procedural, I'm happy to switch.

I also haven't considered whether any memoization of either the classmethod or properties are desirable. Both are possible, and I would probably key on (bids_version, schema_version) and require anybody passing a modified schema to update those fields.

2. Public and private API

The BIDSSchemaAdapter class is purely internal. The change to the public API would be for all functions that could plausibly change behavior as BIDS evolves to acquire a keyword-only argument:

def index_dataset(..., *, schema: SchemaSpec = None): ...

The general approach would be:

def public_function(*args, schema: SchemaSpec = None) -> ReturnType:
    return _private_function(*args, adapter=BIDSSchemaAdapter.load(schema))

def _private_function(*args, adapter: BIDSSchemaAdapter) -> ReturnType: ...

For functions that use the pyarrow schema, you go a step further:

def _pyarrow_function(*args, pa_schema: pa.Schema) -> ReturnType: ...

These become one-way gates, enforced by the type system. If you have a schema spec, you call a public function. If you have already created an adapter, you call a private function, and if you have a pyarrow schema, you can only call pyarrow-schema-accepting functions.

I have not analyzed whether the current call stack alternates between the pyarrow schema and the derived dictionaries. It's possible that enforcing this scheme will require some refactoring, but I think that would be better in the long run than cycling back and forth between the two representations (which is definitely doable).

Questions

1. Functions or classmethods/properties? I'm a bit leaning toward functions based on the near total absence of object-oriented programming in bids2table so far.
2. Should the schema adapter and (some) functions that take it be public? I would lean toward no; easier to add to the public API than remove from it. That said, if a user wants to call multiple functions with a schema, you would get redundant extractions and transformations.
3. Is memoization desirable? Premature optimization? These things are quite cheap so far, but memoization is also an easy way to avoid redundancy without doing (2).

I have rejected the following, and these can be revisited if desired:

1. Maintaining a reference to the BIDS schema Namespace object. You can always extract more from it, and then it's explicit.
2. Allowing a schema adapter to be generated from other schema adapters or pyarrow schemas. While it is convenient to be able to retrieve this representation, I think it led to less clear separation of concerns. I think functions should take exactly one.

[kaitj]

I generally agree with @nx10, I don't think we need to preserve backwards compatibility for this.

Functions or classmethods/properties? I'm a bit leaning toward functions based on the near total absence of object-oriented programming in bids2table so far.

I think functions is probably the way to go here, but also not opposed to start the transition towards a more object oriented codebase. I can't recall at the moment if there was reason for the more functional/procedural development.

Should the schema adapter and (some) functions that take it be public? I would lean toward no; easier to add to the public API than remove from it. That said, if a user wants to call multiple functions with a schema, you would get redundant extractions and transformations.

Agree with your leaning of "no" for the same reasons you noted. Something to revisit in the future if needed.

Is memoization desirable? Premature optimization? These things are quite cheap so far, but memoization is also an easy way to avoid redundancy without doing (2).

The memoization is interesting and I think generally welcome. Would be happy to hear more thoughts about this.