diff --git a/packages/evo-blockmodels/pyproject.toml b/packages/evo-blockmodels/pyproject.toml index ee2f5e3d..2ddc1561 100644 --- a/packages/evo-blockmodels/pyproject.toml +++ b/packages/evo-blockmodels/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "evo-blockmodels" description = "Python SDK for using the Seequent Evo Geoscience Block Model API" -version = "0.6.0" +version = "0.7.0" requires-python = ">=3.10" license-files = ["LICENSE.md"] dynamic = ["readme"] diff --git a/packages/evo-blockmodels/src/evo/blockmodels/__init__.py b/packages/evo-blockmodels/src/evo/blockmodels/__init__.py index d0d4d103..9a0f5fa9 100644 --- a/packages/evo-blockmodels/src/evo/blockmodels/__init__.py +++ b/packages/evo-blockmodels/src/evo/blockmodels/__init__.py @@ -10,7 +10,11 @@ # limitations under the License. from .client import BlockModelAPIClient +from .data import QUALIFIED_TITLE_SEPARATOR, get_qualified_title, qualify_column_titles __all__ = [ + "QUALIFIED_TITLE_SEPARATOR", "BlockModelAPIClient", + "get_qualified_title", + "qualify_column_titles", ] diff --git a/packages/evo-blockmodels/src/evo/blockmodels/_types.py b/packages/evo-blockmodels/src/evo/blockmodels/_types.py index 00e5cc55..6a13687e 100644 --- a/packages/evo-blockmodels/src/evo/blockmodels/_types.py +++ b/packages/evo-blockmodels/src/evo/blockmodels/_types.py @@ -74,6 +74,10 @@ def num_rows(self) -> int: def to_pandas(self) -> DataFrame: """Convert to a pandas-compatible NumPy array or DataFrame, as appropriate""" + def rename_columns(self, names: list[str]) -> "Table": + """Return a copy of this table with its columns renamed to ``names`` (positional).""" + ... + class DataFrame(Protocol): """Pandas DataFrame. diff --git a/packages/evo-blockmodels/src/evo/blockmodels/client.py b/packages/evo-blockmodels/src/evo/blockmodels/client.py index 32dec066..7e5f0d79 100644 --- a/packages/evo-blockmodels/src/evo/blockmodels/client.py +++ b/packages/evo-blockmodels/src/evo/blockmodels/client.py @@ -25,12 +25,17 @@ from ._types import Table from ._utils import convert_dtype, extract_payload +from .data import ( + QUALIFIED_TITLE_SEPARATOR as _QUALIFIED_TITLE_SEPARATOR, +) from .data import ( BaseGridDefinition, BlockModel, ColumnMetadataUpdate, FlexibleGridDefinition, FullySubBlockedGridDefinition, + GroupDefinition, + GroupMetadataUpdate, ListingVersion, OctreeGridDefinition, RegularGridDefinition, @@ -136,6 +141,58 @@ def _version_listing_from_model(version: models.ListingVersion) -> ListingVersio } +def _group_lite_from_definition(definition: GroupDefinition) -> models.GroupLite: + """Convert a public :class:`GroupDefinition` to the generated title-addressed ``GroupLite``.""" + return models.GroupLite(**definition.model_dump(exclude_unset=True)) + + +def _group_values_from_update(update: GroupMetadataUpdate) -> models.GroupUpdateMetadataValuesLite: + """Convert a public :class:`GroupMetadataUpdate` to the generated ``GroupUpdateMetadataValuesLite``. + + Only fields the caller explicitly set are forwarded, so untouched fields are omitted on the wire. + The public ``new_title`` is mapped onto the wire field ``title`` (a group rename). + """ + values = update.model_dump(exclude_unset=True) + if "new_title" in values: + values["title"] = values.pop("new_title") + return models.GroupUpdateMetadataValuesLite(**values) + + +def _build_update_groups_lite( + new: list[GroupDefinition] | None, + update: dict[str, GroupMetadataUpdate] | None, + delete: list[str] | None, +) -> models.UpdateGroupsLite: + """Build the title-addressed ``UpdateGroupsLite`` payload from the public group arguments.""" + return models.UpdateGroupsLite( + new=[_group_lite_from_definition(definition) for definition in (new or [])], + update_metadata=[ + models.GroupUpdateMetadataLite(title=title, values=_group_values_from_update(values)) + for title, values in (update or {}).items() + ], + delete=list(delete or []), + ) + + +def _title_from_column_title(column_title: str, group: str | None) -> str: + """Recover a column's plain title from its (possibly qualified) column title. + + ``column_title`` is the column's title in the data table (e.g. ``Assays▸Cu``); ``group`` is the + qualified group path it should belong to (e.g. ``Assays``). Stripping the ``group▸`` prefix yields + the title the service stores. An ungrouped column (no group) keeps its plain title, so it is + returned as-is. + """ + if not group: + return column_title + prefix = f"{group}{_QUALIFIED_TITLE_SEPARATOR}" + if not column_title.startswith(prefix): + raise MissingColumnInTable( + f"column '{column_title}' is declared in group '{group}' but its column title is not the qualified " + f"title '{prefix}'. Key the data by each column's exact title (see qualify_column_titles)." + ) + return column_title.removeprefix(prefix) + + class BlockModelAPIClient(BaseAPIClient): def __init__( self, @@ -399,21 +456,24 @@ async def _upload_data(self, bm_id: uuid.UUID, job_id: uuid.UUID, upload_url: st return await self.upload_block_model(bm_id, job_id, upload_url, cache_location) async def _update_model_no_data( - self, bm_id: UUID, columns: models.UpdateColumnsLite, comment: str | None = None + self, + bm_id: UUID, + columns: models.UpdateColumnsLite, + comment: str | None = None, + groups: models.UpdateGroupsLite | None = None, ) -> Version: """Helper to apply an UpdateColumnsLite and return the resulting Version. This is for column operations where new data is not required. """ + # Only set ``groups`` when provided so it stays unset (and off the wire) otherwise. + update_data = models.UpdateDataLite1(columns=columns, comment=comment) + if groups is not None: + update_data.groups = groups update_response = await self._column_operations_api.update_block_model_from_latest_version( org_id=str(self._environment.org_id), workspace_id=str(self._environment.workspace_id), bm_id=str(bm_id), - update_data_lite=models.UpdateDataLite( - models.UpdateDataLite1( - columns=columns, - comment=comment, - ) - ), + update_data_lite=models.UpdateDataLite(update_data), additional_headers=self._preview_headers(), ) @@ -636,6 +696,11 @@ async def create_block_model( fully sub-blocked model with ``update_type``=``merge`` and ``geometry_change``=``True`` will fill any missing sub-blocks with data from the parent block. Defaults to ``False``. :return: A tuple containing the created block model and the version of the block model. + + .. note:: + To place columns in a group, first create the model, define the groups with :meth:`update_groups`, + then add the columns with ``column_groups`` on :meth:`add_new_columns` / + :meth:`update_block_model_columns`. Groups cannot be referenced during creation because none exist yet. """ if units is not None and initial_data is None: raise ValueError("units can only be provided if initial_data is provided") @@ -670,6 +735,7 @@ async def add_new_subblocked_columns( data: Table, units: dict[str, str] | None = None, tags: dict[str, dict[str, Any]] | None = None, + column_groups: dict[str, str] | None = None, ) -> Version: """Add new columns to an existing sub-blocked block model. This will not change the sub-blocking structure, thus the provided data must match existing sub-blocks in the model. @@ -678,14 +744,23 @@ async def add_new_subblocked_columns( This method requires the `pyarrow` package to be installed, and the 'cache' parameter to be set in the constructor. :param bm_id: The ID of the block model to add columns to. - :param data: The data containing the new columns to add. + :param data: The data containing the new columns to add, keyed by each column's title + (a plain title for an ungrouped column, or the qualified ``group▸title`` for a grouped one). :param units: A dictionary mapping column names within `data` to units. :param tags: A dictionary mapping column names within `data` to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param column_groups: A dictionary mapping a grouped column's qualified title (its key in + `data`, e.g. ``"Assays▸Cu"``) to the qualified title of the group it belongs to (e.g. ``"Assays"``). + Ungrouped columns are keyed by their plain title in `data` and omitted here. `data` must be keyed by each + column's exact title; :func:`~evo.blockmodels.data.qualify_column_titles` can build that from + plain-titled data. Column groups are a preview feature; the client must be constructed with + ``preview=True`` to use them. :raises CacheNotConfiguredException: If the cache is not configured. :return: The new version of the block model with the added columns. """ - return await self._add_new_columns(bm_id, data, units, geometry_change=False, tags=tags) + return await self._add_new_columns( + bm_id, data, units, geometry_change=False, tags=tags, column_groups=column_groups + ) async def _add_new_columns( self, @@ -694,6 +769,7 @@ async def _add_new_columns( units: dict[str, str] | None = None, geometry_change: bool | None = None, tags: dict[str, dict[str, Any]] | None = None, + column_groups: dict[str, str] | None = None, ) -> Version: """Add new columns to an existing block model. @@ -704,11 +780,18 @@ async def _add_new_columns( This method requires the `pyarrow` package to be installed, and the 'cache' parameter to be set in the constructor. :param bm_id: The ID of the block model to add columns to. - :param data: The data containing the new columns to add. + :param data: The data containing the new columns to add, keyed by each column's title + (a plain title for an ungrouped column, or the qualified ``group▸title`` for a grouped one). :param units: A dictionary mapping column names within `data` to units. :param geometry_change: Whether the geometry of the block model is changing. :param tags: A dictionary mapping column names within `data` to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param column_groups: A dictionary mapping a grouped column's qualified title (its key in + `data`, e.g. ``"Assays▸Cu"``) to the qualified title of the group it belongs to (e.g. ``"Assays"``). + Ungrouped columns are keyed by their plain title in `data` and omitted here. `data` must be keyed by each + column's exact title; :func:`~evo.blockmodels.data.qualify_column_titles` can build that from + plain-titled data. Column groups are a preview feature; the client must be constructed with + ``preview=True`` to use them. :raises CacheNotConfiguredException: If the cache is not configured. :return: The new version of the block model with the added columns. """ @@ -722,6 +805,8 @@ async def _add_new_columns( units = {} if tags is None: tags = {} + if column_groups is None: + column_groups = {} new_column_names = {name for name in schema.names if name not in _GEOMETRY_COLUMNS} unknown_unit_columns = set(units) - new_column_names if unknown_unit_columns: @@ -729,13 +814,19 @@ async def _add_new_columns( unknown_tag_columns = set(tags) - new_column_names if unknown_tag_columns: raise MissingColumnInTable(f"tags reference columns that are not being added: {unknown_tag_columns}") + unknown_group_columns = set(column_groups) - new_column_names + if unknown_group_columns: + raise MissingColumnInTable( + f"column_groups reference columns that are not being added: {unknown_group_columns}" + ) columns = models.UpdateColumnsLite( new=[ models.ColumnLite( - title=name, + title=_title_from_column_title(name, column_groups.get(name)), data_type=convert_dtype(data_type), unit_id=units.get(name), **({"tags": tags[name]} if name in tags else {}), + **({"group": column_groups[name]} if name in column_groups else {}), ) for name, data_type in zip(schema.names, schema.types) if name not in _GEOMETRY_COLUMNS @@ -765,6 +856,7 @@ async def add_new_columns( data: Table, units: dict[str, str] | None = None, tags: dict[str, dict[str, Any]] | None = None, + column_groups: dict[str, str] | None = None, ) -> Version: """Add new columns to an existing regular block model. @@ -773,14 +865,23 @@ async def add_new_columns( This method requires the `pyarrow` package to be installed, and the 'cache' parameter to be set in the constructor. :param bm_id: The ID of the block model to add columns to. - :param data: The data containing the new columns to add. + :param data: The data containing the new columns to add, keyed by each column's title + (a plain title for an ungrouped column, or the qualified ``group▸title`` for a grouped one). :param units: A dictionary mapping column names within `data` to units. :param tags: A dictionary mapping column names within `data` to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param column_groups: A dictionary mapping a grouped column's qualified title (its key in + `data`, e.g. ``"Assays▸Cu"``) to the qualified title of the group it belongs to (e.g. ``"Assays"``). + Ungrouped columns are keyed by their plain title in `data` and omitted here. `data` must be keyed by each + column's exact title; :func:`~evo.blockmodels.data.qualify_column_titles` can build that from + plain-titled data. Column groups are a preview feature; the client must be constructed with + ``preview=True`` to use them. :raises CacheNotConfiguredException: If the cache is not configured. :return: The new version of the block model with the added columns. """ - return await self._add_new_columns(bm_id, data, units, geometry_change=None, tags=tags) + return await self._add_new_columns( + bm_id, data, units, geometry_change=None, tags=tags, column_groups=column_groups + ) async def _update_columns( self, @@ -793,6 +894,7 @@ async def _update_columns( geometry_change: bool | None = None, fill_subblocks: bool | None = None, tags: dict[str, dict[str, Any]] | None = None, + column_groups: dict[str, str] | None = None, update_type: models.UpdateType = models.UpdateType.replace, ) -> Version: if self._cache is None: @@ -805,6 +907,8 @@ async def _update_columns( units = {} if tags is None: tags = {} + if column_groups is None: + column_groups = {} data_type_map = {name: data_type for name, data_type in zip(schema.names, schema.types)} if update_columns is None: @@ -813,10 +917,17 @@ async def _update_columns( if delete_columns is None: delete_columns = set() - # Check for any new or updated columns that are not in the data - missing = (set(new_columns) | update_columns) - data_type_map.keys() + # Every declared column is uploaded under its own title: new columns by their title in ``data``, + # existing data updates by the title the service currently stores them under (a qualified + # ``group▸title`` if grouped, a plain title if not). Data is never renamed, so validate the table + # directly against those titles. + expected_column_titles = set(new_columns) | update_columns + missing = expected_column_titles - data_type_map.keys() if missing: - raise MissingColumnInTable(f"Columns {missing} are not present in the provided table.") + raise MissingColumnInTable( + f"Columns {missing} are not present in the provided table. Key the data by each column's " + "exact title (qualified 'group▸title' for grouped columns, plain otherwise)." + ) unknown_unit_columns = set(units) - set(new_columns) if unknown_unit_columns: @@ -832,13 +943,24 @@ async def _update_columns( "To tag existing columns, use update_column_metadata." ) + # ``column_groups`` here only assigns *new* columns to a group as they are added. Moving or + # ungrouping an existing column is a metadata-only operation; use update_column_metadata. + unknown_group_columns = set(column_groups) - set(new_columns) + if unknown_group_columns: + raise MissingColumnInTable( + f"column_groups reference columns that are not in new_columns: {unknown_group_columns}. " + "column_groups only groups new columns; to move or ungroup an existing column use " + "update_column_metadata." + ) + columns = models.UpdateColumnsLite( new=[ models.ColumnLite( - title=new_column, + title=_title_from_column_title(new_column, column_groups.get(new_column)), data_type=convert_dtype(data_type_map[new_column]), unit_id=units.get(new_column), **({"tags": tags[new_column]} if new_column in tags else {}), + **({"group": column_groups[new_column]} if new_column in column_groups else {}), ) for new_column in new_columns ], @@ -872,6 +994,7 @@ async def update_block_model_columns( units: dict[str, str] | None = None, tags: dict[str, dict[str, Any]] | None = None, update_type: models.UpdateType = models.UpdateType.replace, + column_groups: dict[str, str] | None = None, ) -> Version: """Add, update, or delete regular block model columns. @@ -880,13 +1003,23 @@ async def update_block_model_columns( This method requires the `pyarrow` package to be installed, and the 'cache' parameter to be set in the constructor. :param bm_id: The ID of the block model to add columns to. - :param data: The data containing the new columns to add. - :param new_columns: A list of new column names to add to the block model. - :param update_columns: A set of column names to update in the block model. - :param delete_columns: A set of column names to delete from the block model. + :param data: The data containing the affected columns, keyed by each column's title + (a plain title for an ungrouped column, or the qualified ``group▸title`` for a grouped one). + :func:`~evo.blockmodels.data.qualify_column_titles` can build these titles from plain-titled data. + :param new_columns: A list of new columns to add, named by their title in `data` (qualified + ``group▸title`` for a grouped column, plain otherwise). + :param update_columns: A set of existing columns to re-upload, each identified by the title the service + currently stores it under: its qualified title (``group▸title``) if grouped, or its plain title if not. + :param delete_columns: A set of existing columns to delete, identified the same way as ``update_columns`` + (qualified title if grouped, plain otherwise). :param units: A dictionary mapping column names within `data` to units. :param tags: A dictionary mapping new column names to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param column_groups: A dictionary assigning **new** columns to groups: map a new column's qualified + title (its key in `data`, e.g. ``"Assays▸Cu"``) to the qualified title of the group it belongs to. + To move or ungroup an *existing* column, use :meth:`update_column_metadata` instead — a group change is + metadata-only and does not require re-uploading data. Column groups are a preview feature; the client + must be constructed with ``preview=True`` to use them. :param: update_type: Provide the type of update. Either 'replace' or 'merge' (default: replace) :raises CacheNotConfiguredException: If the cache is not configured. :return: The new version of the block model with the added columns. @@ -900,6 +1033,7 @@ async def update_block_model_columns( units, geometry_change=None, tags=tags, + column_groups=column_groups, update_type=update_type, ) @@ -915,6 +1049,7 @@ async def update_subblocked_columns( fill_subblocks: bool | None = None, tags: dict[str, dict[str, Any]] | None = None, update_type: models.UpdateType = models.UpdateType.replace, + column_groups: dict[str, str] | None = None, ) -> Version: """Add, update, or delete sub-blocked block model columns. @@ -928,10 +1063,15 @@ async def update_subblocked_columns( This method requires the `pyarrow` package to be installed, and the 'cache' parameter to be set in the constructor. :param bm_id: The ID of the block model to add columns to. - :param data: The data containing the new columns to add. - :param new_columns: A list of new column names to add to the block model. - :param update_columns: A set of column names to update in the block model. - :param delete_columns: A set of column names to delete from the block model. + :param data: The data containing the affected columns, keyed by each column's title + (a plain title for an ungrouped column, or the qualified ``group▸title`` for a grouped one). + :func:`~evo.blockmodels.data.qualify_column_titles` can build these titles from plain-titled data. + :param new_columns: A list of new columns to add, named by their title in `data` (qualified + ``group▸title`` for a grouped column, plain otherwise). + :param update_columns: A set of existing columns to re-upload, each identified by the title the service + currently stores it under: its qualified title (``group▸title``) if grouped, or its plain title if not. + :param delete_columns: A set of existing columns to delete, identified the same way as ``update_columns`` + (qualified title if grouped, plain otherwise). :param units: A dictionary mapping column names within `data` to units. :param geometry_change: Whether the geometry of the sub-blocked model changes. :param fill_subblocks: If ``True``, any missing sub-blocks will be filled with data from the parent block. @@ -939,6 +1079,11 @@ async def update_subblocked_columns( the block model's own ``fill_subblocks`` setting is used. :param tags: A dictionary mapping new column names to their tags object. Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. + :param column_groups: A dictionary assigning **new** columns to groups: map a new column's qualified + title (its key in `data`, e.g. ``"Assays▸Cu"``) to the qualified title of the group it belongs to. + To move or ungroup an *existing* column, use :meth:`update_column_metadata` instead — a group change is + metadata-only and does not require re-uploading data. Column groups are a preview feature; the client + must be constructed with ``preview=True`` to use them. :param: update_type: Provide the type of update. Either 'replace' or 'merge' (default: replace) """ return await self._update_columns( @@ -951,6 +1096,7 @@ async def update_subblocked_columns( geometry_change=geometry_change, fill_subblocks=fill_subblocks, tags=tags, + column_groups=column_groups, update_type=update_type, ) @@ -968,16 +1114,22 @@ async def update_column_metadata( - A ``str`` sets the column's unit ID. - ``None`` clears the column's unit ID. - - A :class:`ColumnMetadataUpdate` sets any combination of unit ID and/or tags. Only the + - A :class:`ColumnMetadataUpdate` sets any combination of unit ID, tags and/or group. Only the fields explicitly set on the object are sent; unset fields are left untouched. Set - ``tags={}`` to clear a column's tags, or ``unit_id=None`` to clear its unit. + ``tags={}`` to clear a column's tags, ``unit_id=None`` to clear its unit, or ``group=""`` to + move the column out of any group. + + A column's group is metadata, so it can be moved (or ungrouped) here without re-uploading its + data. Address the column by the title the service currently stores it under: its qualified title + (``group▸title``) if it is currently grouped, or its plain title if it is not. Set + ``ColumnMetadataUpdate(group=...)`` to the target group's qualified title (or ``""`` to ungroup). Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them. :param bm_id: The ID of the block model to update. :param column_updates: A dictionary mapping column titles to their metadata update. Example: {"Cu": "%[mass]", "Au": None, - "Ag": ColumnMetadataUpdate(tags={"source": "assay"})} + "Assays▸Ag": ColumnMetadataUpdate(group="Geology")} :param comment: An optional comment describing the metadata changes. This is max 250 characters. :return: The new version of the block model with updated metadata. """ @@ -1003,6 +1155,51 @@ def _to_values(value: str | None | ColumnMetadataUpdate) -> models.UpdateMetadat return await self._update_model_no_data(bm_id, columns, comment=comment) + async def update_groups( + self, + bm_id: UUID, + *, + new: list[GroupDefinition] | None = None, + update: dict[str, GroupMetadataUpdate] | None = None, + delete: list[str] | None = None, + comment: str | None = None, + ) -> Version: + """Create, update, and/or delete column groups on a block model. + + This method manages group definitions without requiring data upload or cache configuration. Any + combination of ``new``, ``update`` and ``delete`` can be supplied in a single call. + + Groups are addressed by their qualified title (a bare title for a top-level group, or segments + joined by ``▸`` for a nested group). To assign a *new* column to a group, use the ``column_groups`` + parameter on the column methods; to move or ungroup an *existing* column, use + :meth:`update_column_metadata`. To resolve a written group back to its server-assigned UUID and + resolved policy, use the helpers on the returned :class:`~evo.blockmodels.data.Version`, e.g. + :meth:`~evo.blockmodels.data.Version.group_by_qualified_title`. + + Column groups are a preview feature; the client must be constructed with ``preview=True`` to use them. + + :param bm_id: The ID of the block model to update. + :param new: Definitions of new groups to create. + :param update: A dictionary mapping the qualified title of an existing group to the metadata update + to apply to it. Use :class:`GroupMetadataUpdate` to rename, re-parent, change the missing-column + policy, replace tags, or toggle the hidden flag. + :param delete: Qualified titles of groups to delete. + :param comment: An optional comment describing the changes. This is max 250 characters. + :return: The new version of the block model with the updated groups. + """ + if not new and not update and not delete: + raise ValueError("At least one of 'new', 'update' or 'delete' must be provided.") + + columns = models.UpdateColumnsLite( + new=[], + update=[], + delete=[], + rename=[], + ) + groups = _build_update_groups_lite(new, update, delete) + + return await self._update_model_no_data(bm_id, columns, comment=comment, groups=groups) + async def rename_block_model_columns( self, bm_id: UUID, diff --git a/packages/evo-blockmodels/src/evo/blockmodels/data.py b/packages/evo-blockmodels/src/evo/blockmodels/data.py index c09e9473..ee7f66cf 100644 --- a/packages/evo-blockmodels/src/evo/blockmodels/data.py +++ b/packages/evo-blockmodels/src/evo/blockmodels/data.py @@ -14,36 +14,101 @@ from typing import Any from uuid import UUID +from pydantic import ConfigDict + from evo.common import ResourceMetadata from evo.workspaces import ServiceUser from ._model_config import CustomBaseModel +from ._types import Table from .endpoints.models import ( BBox, BBoxXYZ, Column, ListingColumn, ListingGroup, + MissingColumnPolicy, ResolvedGroup, RotationAxis, ) __all__ = [ + "QUALIFIED_TITLE_SEPARATOR", "BaseGridDefinition", "BlockModel", "Column", "ColumnMetadataUpdate", "FlexibleGridDefinition", "FullySubBlockedGridDefinition", + "GroupDefinition", + "GroupMetadataUpdate", "ListingColumn", "ListingGroup", "ListingVersion", + "MissingColumnPolicy", "OctreeGridDefinition", "RegularGridDefinition", "ResolvedGroup", "Version", + "get_qualified_title", + "qualify_column_titles", ] +QUALIFIED_TITLE_SEPARATOR = "\u25b8" +"""Default single-character separator (``▸``) used to build and parse qualified group titles.""" + + +def get_qualified_title(group: str | None, title: str, separator: str = QUALIFIED_TITLE_SEPARATOR) -> str: + """Build the column title the block model service expects for a column. + + A grouped column must be uploaded under its qualified title (``group▸…▸title``); an + ungrouped column (no group, or ``group == ""``) keeps its plain title. + + :param group: The qualified title (path) of the column's group, or ``None``/``""`` if ungrouped. + :param title: The column's title. + :param separator: Separator used to join the path segments. + :return: The qualified title (``group▸title``) if grouped, otherwise ``title``. + """ + if group: + return f"{group}{separator}{title}" + return title + + +def qualify_column_titles( + data: Table, groups: dict[str, str], separator: str = QUALIFIED_TITLE_SEPARATOR +) -> tuple[Table, dict[str, str]]: + """Rename a table's columns to the qualified titles the service expects. + + The column methods on :class:`~evo.blockmodels.client.BlockModelAPIClient` expect ``data`` to be + keyed by each column's title — a qualified ``group▸title`` for a grouped column, or the plain title + for an ungrouped one — and do not rename anything for you. This opt-in helper performs that shift: + pass a table keyed by plain titles plus a mapping of the columns you want grouped, and it returns the + renamed table alongside the ``column_groups`` mapping to hand back to the client. + + :param data: A table keyed by plain column titles. + :param groups: A mapping of a column's title to the qualified title of the group it should be + placed in. Columns absent from the mapping (or mapped to ``""``) are treated as ungrouped and + left with their plain title. + :param separator: Separator used to build qualified titles. + :return: A ``(table, column_groups)`` pair: the table with grouped columns renamed to their qualified + title, and a ``{qualified_title: group}`` mapping to pass as ``column_groups``. + :raises KeyError: If ``groups`` references a column that is not present in ``data``. + """ + existing = set(data.schema.names) + unknown = set(groups) - existing + if unknown: + raise KeyError(f"groups reference columns that are not present in the table: {unknown}") + + new_names: list[str] = [] + column_groups: dict[str, str] = {} + for name in data.schema.names: + group = groups.get(name) + column_title = get_qualified_title(group, name, separator) + new_names.append(column_title) + if group: + column_groups[column_title] = group + return data.rename_columns(new_names), column_groups + class ColumnMetadataUpdate(CustomBaseModel): """A per-column metadata update for :meth:`BlockModelAPIClient.update_column_metadata`. @@ -51,9 +116,18 @@ class ColumnMetadataUpdate(CustomBaseModel): Only the fields you explicitly set are sent to the service; unset fields are left untouched. """ + model_config = ConfigDict(extra="forbid", protected_namespaces=()) + unit_id: str | None = None """The new unit ID for the column.""" + group: str | None = None + """The qualified title of the new group for the column (a bare title for a top-level group, or + segments joined by ``▸`` for a nested group). Send ``""`` to move the column out of any group. + Omit this field to leave the column's current group unchanged. + + Column groups are a preview feature; the client must be constructed with ``preview=True`` to use them.""" + tags: dict[str, Any] | None = None """Replacement tags for the column. Send a populated object to replace the column's tags wholesale, or ``{}`` to clear them. Omit this field to leave the existing tags untouched. @@ -61,6 +135,63 @@ class ColumnMetadataUpdate(CustomBaseModel): Column tags are a preview feature; the client must be constructed with ``preview=True`` to use them.""" +class GroupDefinition(CustomBaseModel): + """Definition of a new column group to create via :meth:`BlockModelAPIClient.update_groups`. + + Column groups are a preview feature; the client must be constructed with ``preview=True`` to use them. + """ + + model_config = ConfigDict(extra="forbid", protected_namespaces=()) + + title: str + """Human-readable label for the group, unique across its siblings. Must not contain the qualified + title separator ``▸``.""" + + parent_group: str | None = None + """Qualified title of the parent group (a bare title for a top-level parent, or segments joined by + ``▸`` for a nested parent). ``None`` (the default) makes this a top-level group.""" + + missing_column_policy: MissingColumnPolicy | None = None + """Policy for columns in this group's zone that are absent from an update. ``None`` (the default) + lets the service apply its default of ``INHERIT``.""" + + tags: dict[str, Any] | None = None + """Publisher-supplied free-form metadata for the group.""" + + is_hidden: bool = False + """When ``True``, the group's direct member columns are excluded from wildcard queries unless + ``include_hidden`` is set.""" + + +class GroupMetadataUpdate(CustomBaseModel): + """A metadata update for an existing column group, used by :meth:`BlockModelAPIClient.update_groups`. + + Only the fields you explicitly set are sent to the service; unset fields are left unchanged. + + Column groups are a preview feature; the client must be constructed with ``preview=True`` to use them. + """ + + model_config = ConfigDict(extra="forbid", protected_namespaces=()) + + new_title: str | None = None + """Rename the group to this title. Omit to leave the title unchanged.""" + + parent_group: str | None = None + """Re-parent the group. Provide the qualified title of the new parent group, or an empty string + (``""``) to make it a top-level group. Omit to leave the parent unchanged.""" + + missing_column_policy: MissingColumnPolicy | None = None + """Set the group's missing-column policy. Use ``INHERIT`` to adopt the parent chain's policy zone. + Omit to leave the policy unchanged.""" + + tags: dict[str, Any] | None = None + """Replacement tags for the group. Send a populated object to replace the group's tags wholesale, + or ``{}`` to clear them. Omit to leave the existing tags untouched.""" + + is_hidden: bool | None = None + """Set the group's hidden flag. Omit to leave it unchanged.""" + + @dataclass(frozen=True, kw_only=True) class BaseGridDefinition: """Base class for grid definitions.""" @@ -279,6 +410,68 @@ def __repr__(self) -> str: f"columns={col_names})" ) + def group_by_uuid(self, group_uuid: UUID) -> "ResolvedGroup | ListingGroup | None": + """Return the group on this version with the given UUID, or ``None`` if there is no match.""" + for group in self.groups: + if group.group_uuid == group_uuid: + return group + return None + + def group_for_column(self, column: "Column | ListingColumn") -> "ResolvedGroup | ListingGroup | None": + """Resolve the group a column belongs to. + + This bridges the read path (columns reference their group by UUID) so a caller can get the + group's title, parent and resolved missing-column policy without hand-rolling a lookup. + + A column object is required rather than a title: titles are only unique within a group, so a + title can be ambiguous across groups. The column carries its group reference unambiguously. + + :param column: A column from this version. + :return: The column's group, or ``None`` if the column is ungrouped or its group is not on + this version. + """ + if column.group_uuid is None: + return None + return self.group_by_uuid(column.group_uuid) + + def qualified_group_title( + self, group: "ResolvedGroup | ListingGroup", separator: str = QUALIFIED_TITLE_SEPARATOR + ) -> str: + """Build the qualified (``▸``-joined) title of a group by walking its parent chain. + + :param group: A group from this version. + :param separator: Separator used to join the path segments. + :return: The fully-qualified group title. + """ + titles = [group.title] + seen = {group.group_uuid} + parent_uuid = group.parent_group_uuid + while parent_uuid is not None and parent_uuid not in seen: + parent = self.group_by_uuid(parent_uuid) + if parent is None: + break + titles.append(parent.title) + seen.add(parent.group_uuid) + parent_uuid = parent.parent_group_uuid + return separator.join(reversed(titles)) + + def group_by_qualified_title( + self, qualified_title: str, separator: str = QUALIFIED_TITLE_SEPARATOR + ) -> "ResolvedGroup | ListingGroup | None": + """Find a group by the qualified title it was created with. + + This bridges the write/read asymmetry: groups are written by title but read back by UUID, so a + caller who created a group by title can find it again on the returned version. + + :param qualified_title: A bare title for a top-level group, or ``▸``-joined segments for a nested group. + :param separator: Separator used to parse and rebuild qualified titles. + :return: The matching group, or ``None`` if there is no match. + """ + for group in self.groups: + if self.qualified_group_title(group, separator) == qualified_title: + return group + return None + @dataclass(frozen=True, kw_only=True, repr=False) class Version(_VersionBase): diff --git a/packages/evo-blockmodels/tests/test_group_helpers.py b/packages/evo-blockmodels/tests/test_group_helpers.py new file mode 100644 index 00000000..e404db1b --- /dev/null +++ b/packages/evo-blockmodels/tests/test_group_helpers.py @@ -0,0 +1,249 @@ +# Copyright © 2025 Bentley Systems, Incorporated +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +import uuid +from datetime import datetime, timezone + +from evo.blockmodels.data import QUALIFIED_TITLE_SEPARATOR, MissingColumnPolicy, Version +from evo.blockmodels.endpoints import models +from evo.common import ServiceUser + +DATE = datetime(2021, 1, 1, tzinfo=timezone.utc) +USER = ServiceUser.from_model(models.IMSUserInfo(email="test@test.com", name="Test User", id=uuid.uuid4())) + +PARENT_UUID = uuid.uuid4() +CHILD_UUID = uuid.uuid4() +COL_IN_CHILD = uuid.uuid4() + + +def _resolved_group(group_uuid, title, parent_group_uuid=None): + return models.ResolvedGroup( + group_uuid=group_uuid, + title=title, + parent_group_uuid=parent_group_uuid, + missing_column_policy=MissingColumnPolicy.SET_NULL, + resolved_missing_column_policy=MissingColumnPolicy.SET_NULL, + tags={"source": "assay"}, + ) + + +def _make_version() -> Version: + return Version( + bm_uuid=uuid.uuid4(), + version_id=2, + version_uuid=uuid.uuid4(), + parent_version_id=1, + base_version_id=1, + geoscience_version_id="3", + created_at=DATE, + created_by=USER, + comment="", + columns=[ + models.Column( + col_id=str(COL_IN_CHILD), title="Cu", data_type=models.DataType.Float64, group_uuid=CHILD_UUID + ), + models.Column(col_id=str(uuid.uuid4()), title="Au", data_type=models.DataType.Float64), + ], + groups=[ + _resolved_group(PARENT_UUID, "Assays"), + _resolved_group(CHILD_UUID, "Primary", parent_group_uuid=PARENT_UUID), + ], + ) + + +class TestVersionGroupHelpers(unittest.TestCase): + def setUp(self) -> None: + self.version = _make_version() + + def test_group_by_uuid(self) -> None: + group = self.version.group_by_uuid(CHILD_UUID) + self.assertIsNotNone(group) + self.assertEqual(group.title, "Primary") + self.assertIsNone(self.version.group_by_uuid(uuid.uuid4())) + + def test_qualified_group_title_nested(self) -> None: + child = self.version.group_by_uuid(CHILD_UUID) + self.assertEqual( + self.version.qualified_group_title(child), + f"Assays{QUALIFIED_TITLE_SEPARATOR}Primary", + ) + + def test_qualified_group_title_top_level(self) -> None: + parent = self.version.group_by_uuid(PARENT_UUID) + self.assertEqual(self.version.qualified_group_title(parent), "Assays") + + def test_qualified_group_title_custom_separator(self) -> None: + child = self.version.group_by_uuid(CHILD_UUID) + self.assertEqual(self.version.qualified_group_title(child, separator="/"), "Assays/Primary") + + def test_group_by_qualified_title_roundtrip(self) -> None: + # A caller who created "Assays▸Primary" by title finds it again by its qualified title. + group = self.version.group_by_qualified_title(f"Assays{QUALIFIED_TITLE_SEPARATOR}Primary") + self.assertIsNotNone(group) + self.assertEqual(group.group_uuid, CHILD_UUID) + self.assertIsNone(self.version.group_by_qualified_title("Missing")) + + def test_group_for_column_by_object(self) -> None: + column = self.version.columns[0] + group = self.version.group_for_column(column) + self.assertIsNotNone(group) + self.assertEqual(group.group_uuid, CHILD_UUID) + self.assertEqual(group.resolved_missing_column_policy, MissingColumnPolicy.SET_NULL) + + def test_group_for_ungrouped_column(self) -> None: + # columns[1] ("Au") has no group_uuid. + self.assertIsNone(self.version.group_for_column(self.version.columns[1])) + + def test_group_for_column_with_dangling_group(self) -> None: + # A column referencing a group that is not on this version resolves to None rather than guessing. + orphan_column = models.Column( + col_id=str(uuid.uuid4()), title="Cu", data_type=models.DataType.Float64, group_uuid=uuid.uuid4() + ) + self.assertIsNone(self.version.group_for_column(orphan_column)) + + def test_group_for_column_is_unambiguous_across_groups(self) -> None: + # Two columns share the title "Cu" but live in different groups; each resolves to its own group + # because the lookup uses the column's group_uuid, not its (non-unique) title. + other_uuid = uuid.uuid4() + version = Version( + bm_uuid=uuid.uuid4(), + version_id=2, + version_uuid=uuid.uuid4(), + parent_version_id=1, + base_version_id=1, + geoscience_version_id="3", + created_at=DATE, + created_by=USER, + comment="", + columns=[ + models.Column( + col_id=str(uuid.uuid4()), title="Cu", data_type=models.DataType.Float64, group_uuid=CHILD_UUID + ), + models.Column( + col_id=str(uuid.uuid4()), title="Cu", data_type=models.DataType.Float64, group_uuid=other_uuid + ), + ], + groups=[ + _resolved_group(CHILD_UUID, "Primary"), + _resolved_group(other_uuid, "Secondary"), + ], + ) + self.assertEqual(version.group_for_column(version.columns[0]).group_uuid, CHILD_UUID) + self.assertEqual(version.group_for_column(version.columns[1]).group_uuid, other_uuid) + + def test_qualified_group_title_tolerates_broken_parent_chain(self) -> None: + # A dangling parent reference must not loop or raise; it just stops walking. + orphan = models.ResolvedGroup( + group_uuid=uuid.uuid4(), + title="Orphan", + parent_group_uuid=uuid.uuid4(), + missing_column_policy=MissingColumnPolicy.INHERIT, + resolved_missing_column_policy=MissingColumnPolicy.USE_PREVIOUS, + ) + object.__setattr__(self.version, "groups", [*self.version.groups, orphan]) + self.assertEqual(self.version.qualified_group_title(orphan), "Orphan") + + +if __name__ == "__main__": + unittest.main() + + +class TestGroupInputModelsRejectExtras(unittest.TestCase): + def test_group_definition_forbids_unknown_fields(self) -> None: + from pydantic import ValidationError + + from evo.blockmodels.data import GroupDefinition + + with self.assertRaises(ValidationError): + GroupDefinition(title="Assays", parnet_group="typo") + + def test_group_metadata_update_forbids_title_field(self) -> None: + # The wire rename field is ``new_title``; passing ``title`` (the wire name) must be rejected + # so it can't silently collide with the new_title -> title remap. + from pydantic import ValidationError + + from evo.blockmodels.data import GroupMetadataUpdate + + with self.assertRaises(ValidationError): + GroupMetadataUpdate(title="X") + + def test_column_metadata_update_accepts_group_field(self) -> None: + # A column's group is metadata and can be changed without re-uploading data, so ``group`` is a + # supported field. Only fields explicitly set are forwarded onto the wire. + from evo.blockmodels.data import ColumnMetadataUpdate + + update = ColumnMetadataUpdate(group="Assays") + self.assertEqual(update.group, "Assays") + self.assertEqual(update.model_dump(exclude_unset=True), {"group": "Assays"}) + + # An empty string ungroups the column. + self.assertEqual(ColumnMetadataUpdate(group="").model_dump(exclude_unset=True), {"group": ""}) + + def test_column_metadata_update_forbids_unknown_field(self) -> None: + from pydantic import ValidationError + + from evo.blockmodels.data import ColumnMetadataUpdate + + with self.assertRaises(ValidationError): + ColumnMetadataUpdate(not_a_field="x") + + +class TestQualifyColumnTitles(unittest.TestCase): + def test_get_qualified_title_builds_qualified_and_bare_titles(self) -> None: + from evo.blockmodels.data import get_qualified_title + + self.assertEqual(get_qualified_title("Assays", "Cu"), f"Assays{QUALIFIED_TITLE_SEPARATOR}Cu") + self.assertEqual( + get_qualified_title("Assays\u25b8Primary", "Cu"), f"Assays\u25b8Primary{QUALIFIED_TITLE_SEPARATOR}Cu" + ) + # An ungrouped column keeps its bare title. + self.assertEqual(get_qualified_title(None, "Cu"), "Cu") + self.assertEqual(get_qualified_title("", "Cu"), "Cu") + + def test_qualify_column_titles_renames_and_builds_column_groups(self) -> None: + import pyarrow + + from evo.blockmodels.data import qualify_column_titles + + data = pyarrow.table({"i": [1], "Cu": [2.0], "Au": [3.0], "rock": ["x"]}) + renamed, column_groups = qualify_column_titles(data, {"Cu": "Assays", "Au": "Assays"}) + + self.assertEqual( + renamed.schema.names, + ["i", f"Assays{QUALIFIED_TITLE_SEPARATOR}Cu", f"Assays{QUALIFIED_TITLE_SEPARATOR}Au", "rock"], + ) + self.assertEqual( + column_groups, + {f"Assays{QUALIFIED_TITLE_SEPARATOR}Cu": "Assays", f"Assays{QUALIFIED_TITLE_SEPARATOR}Au": "Assays"}, + ) + # Untouched columns keep their plain title and are absent from column_groups. + self.assertEqual(renamed.column("i").to_pylist(), [1]) + + def test_qualify_column_titles_treats_empty_group_as_ungrouped(self) -> None: + import pyarrow + + from evo.blockmodels.data import qualify_column_titles + + data = pyarrow.table({"Cu": [2.0]}) + renamed, column_groups = qualify_column_titles(data, {"Cu": ""}) + + self.assertEqual(renamed.schema.names, ["Cu"]) + self.assertEqual(column_groups, {}) + + def test_qualify_column_titles_rejects_unknown_columns(self) -> None: + import pyarrow + + from evo.blockmodels.data import qualify_column_titles + + data = pyarrow.table({"Cu": [2.0]}) + with self.assertRaises(KeyError): + qualify_column_titles(data, {"Au": "Assays"}) diff --git a/packages/evo-blockmodels/tests/test_update.py b/packages/evo-blockmodels/tests/test_update.py index d7f42670..13697a44 100644 --- a/packages/evo-blockmodels/tests/test_update.py +++ b/packages/evo-blockmodels/tests/test_update.py @@ -16,10 +16,11 @@ from unittest import mock import pyarrow +import pyarrow.parquet from parameterized import parameterized from evo.blockmodels import BlockModelAPIClient -from evo.blockmodels.data import ColumnMetadataUpdate +from evo.blockmodels.data import ColumnMetadataUpdate, GroupDefinition, GroupMetadataUpdate, MissingColumnPolicy from evo.blockmodels.endpoints import models from evo.blockmodels.endpoints.models import JobResponse, JobStatus from evo.blockmodels.exceptions import CacheNotConfiguredException, JobFailedException, MissingColumnInTable @@ -141,6 +142,12 @@ def setUp(self) -> None: environment=self.environment, preview=True, ) + self.preview_client = BlockModelAPIClient( + connector=self.connector, + environment=self.environment, + cache=self.cache, + preview=True, + ) self.setup_universal_headers(get_header_metadata(BlockModelAPIClient.__module__)) @property @@ -891,6 +898,75 @@ async def test_update_column_metadata_with_preview_sends_header(self) -> None: ) self.assertEqual(version.version_id, 2) + async def test_update_column_metadata_moves_column_between_groups(self) -> None: + """An existing grouped column is moved to another group as a metadata-only update (no data).""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + await self.preview_client_without_cache.update_column_metadata( + BM_UUID, + # The column is addressed by its current qualified title; the group is set to the new path. + column_updates={"Assays\u25b8col1": ColumnMetadataUpdate(group="Geology")}, + ) + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[], + update=[], + rename=[], + delete=[], + update_metadata=[ + models.UpdateMetadataLite( + title="Assays\u25b8col1", values=models.UpdateMetadataValuesLite(group="Geology") + ), + ], + ), + comment=None, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS | {"API-Preview": "opt-in"}, + ) + + async def test_update_column_metadata_ungroups_column(self) -> None: + """An existing grouped column is ungrouped as a metadata-only update by sending ``group=""``.""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + await self.bms_client_without_cache.update_column_metadata( + BM_UUID, + column_updates={"Assays\u25b8col2": ColumnMetadataUpdate(group="")}, + ) + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[], + update=[], + rename=[], + delete=[], + update_metadata=[ + models.UpdateMetadataLite( + title="Assays\u25b8col2", values=models.UpdateMetadataValuesLite(group="") + ), + ], + ), + comment=None, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS, + ) + async def test_rename_block_model_columns(self) -> None: self.transport.set_request_handler( UpdateRequestHandler( @@ -1038,3 +1114,295 @@ async def test_delete_block_model_columns_job_failed(self) -> None: ) with self.assertRaises(JobFailedException): await self.bms_client_without_cache.delete_block_model_columns(BM_UUID, ["col1"]) + + async def test_add_new_columns_with_column_groups(self) -> None: + """New columns can be assigned to a group via ``column_groups``.""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + with ( + mock.patch("evo.common.io.upload.StorageDestination") as mock_destination, + mock.patch("pyarrow.parquet.write_table", wraps=pyarrow.parquet.write_table) as mock_write, + ): + mock_destination.upload_file = mock.AsyncMock() + # The caller keys the grouped column by its qualified title; the SDK never renames data. + data = pyarrow.table( + { + "i": [1, 2, 3], + "j": [4, 5, 6], + "k": [7, 8, 9], + "col1": ["A", "B", "B"], + "Assays\u25b8col2": [4.5, 5.3, 6.2], + } + ) + await self.preview_client.add_new_columns( + BM_UUID, + data, + column_groups={"Assays\u25b8col2": "Assays"}, + ) + mock_destination.upload_file.assert_called_once() + + # col2 is placed in the "Assays" group; col1 stays ungrouped so has no group field on the wire. + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[ + models.ColumnLite(title="col1", data_type=models.DataType.Utf8, unit_id=None), + models.ColumnLite( + title="col2", data_type=models.DataType.Float64, unit_id=None, group="Assays" + ), + ], + update=[], + rename=[], + delete=[], + ), + update_type=models.UpdateType.replace, + geometry_change=None, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS | {"API-Preview": "opt-in"}, + ) + + # The uploaded data is untouched: the caller already keyed the grouped column by its + # qualified title, and the ungrouped column keeps its plain title. + uploaded_table = mock_write.call_args.args[0] + self.assertEqual(uploaded_table.schema.names, ["i", "j", "k", "col1", "Assays\u25b8col2"]) + + async def test_add_new_columns_with_unknown_group_column(self) -> None: + with self.assertRaises(MissingColumnInTable): + await self.bms_client.add_new_columns( + BM_UUID, + REGULAR_DATA, + column_groups={"does_not_exist": "Assays"}, + ) + + async def test_update_block_model_columns_with_column_groups(self) -> None: + """New columns added through update_block_model_columns can be grouped.""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + with ( + mock.patch("evo.common.io.upload.StorageDestination") as mock_destination, + mock.patch("pyarrow.parquet.write_table", wraps=pyarrow.parquet.write_table) as mock_write, + ): + mock_destination.upload_file = mock.AsyncMock() + # New columns are named by their title; the grouped one is qualified. + data = pyarrow.table( + { + "i": [1, 2, 3], + "j": [4, 5, 6], + "k": [7, 8, 9], + "Assays\u25b8Primary\u25b8col1": ["A", "B", "B"], + "col2": [4.5, 5.3, 6.2], + } + ) + await self.bms_client.update_block_model_columns( + BM_UUID, + data, + new_columns=["Assays\u25b8Primary\u25b8col1", "col2"], + column_groups={"Assays\u25b8Primary\u25b8col1": "Assays\u25b8Primary"}, + ) + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[ + models.ColumnLite( + title="col1", + data_type=models.DataType.Utf8, + unit_id=None, + group="Assays\u25b8Primary", + ), + models.ColumnLite(title="col2", data_type=models.DataType.Float64, unit_id=None), + ], + update=[], + rename=[], + delete=[], + ), + update_type=models.UpdateType.replace, + geometry_change=None, + fill_subblocks=None, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS, + ) + + uploaded_table = mock_write.call_args.args[0] + # Data is uploaded exactly as provided (already keyed by the qualified title). + self.assertEqual(uploaded_table.schema.names, ["i", "j", "k", "Assays\u25b8Primary\u25b8col1", "col2"]) + + async def test_update_columns_with_unknown_group_column(self) -> None: + with self.assertRaises(MissingColumnInTable): + await self.bms_client.update_block_model_columns( + BM_UUID, + REGULAR_DATA, + new_columns=["col1"], + column_groups={"col2": "Assays"}, # col2 is not a new column + ) + + async def test_update_block_model_columns_data_only_update_of_grouped_column(self) -> None: + """A plain data update of an already-grouped column references it by its current qualified title.""" + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + with ( + mock.patch("evo.common.io.upload.StorageDestination") as mock_destination, + mock.patch("pyarrow.parquet.write_table", wraps=pyarrow.parquet.write_table) as mock_write, + ): + mock_destination.upload_file = mock.AsyncMock() + # A data-only update keeps the grouped column under its current qualified title. + data = pyarrow.table( + { + "i": [1, 2, 3], + "j": [4, 5, 6], + "k": [7, 8, 9], + "Assays\u25b8col1": ["A", "B", "B"], + "col2": [4.5, 5.3, 6.2], + } + ) + await self.bms_client.update_block_model_columns( + BM_UUID, + data, + new_columns=[], + update_columns={"Assays\u25b8col1"}, # current qualified title, no group change + ) + + # No group change, but the grouped column is still referenced by its current qualified title. + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite( + new=[], + update=["Assays\u25b8col1"], + rename=[], + delete=[], + ), + update_type=models.UpdateType.replace, + geometry_change=None, + fill_subblocks=None, + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS, + ) + + uploaded_table = mock_write.call_args.args[0] + # The grouped column keeps its current qualified title so its data binds correctly. + self.assertEqual(uploaded_table.schema.names, ["i", "j", "k", "Assays\u25b8col1", "col2"]) + + async def test_update_groups_create(self) -> None: + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + version = await self.preview_client_without_cache.update_groups( + BM_UUID, + new=[ + GroupDefinition(title="Assays", missing_column_policy=MissingColumnPolicy.SET_NULL), + GroupDefinition(title="Primary", parent_group="Assays", is_hidden=True), + ], + ) + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite(new=[], update=[], delete=[], rename=[]), + comment=None, + ) + expected_update_body.groups = models.UpdateGroupsLite( + new=[ + models.GroupLite(title="Assays", missing_column_policy=MissingColumnPolicy.SET_NULL), + models.GroupLite(title="Primary", parent_group="Assays", is_hidden=True), + ], + update_metadata=[], + delete=[], + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS | {"API-Preview": "opt-in"}, + ) + self.assertEqual(version.bm_uuid, BM_UUID) + self.assertEqual(version.version_id, 2) + + async def test_update_groups_update_metadata_and_delete(self) -> None: + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse(job_status=JobStatus.COMPLETE, payload=UPDATED_VERSION), + ) + ) + await self.preview_client_without_cache.update_groups( + BM_UUID, + update={ + "Assays": GroupMetadataUpdate( + new_title="Assay Results", missing_column_policy=MissingColumnPolicy.INHERIT + ), + "Waste": GroupMetadataUpdate(parent_group="", tags={}), + }, + delete=["Old Group"], + comment="Reorganise groups", + ) + + expected_update_body = models.UpdateDataLite1( + columns=models.UpdateColumnsLite(new=[], update=[], delete=[], rename=[]), + comment="Reorganise groups", + ) + expected_update_body.groups = models.UpdateGroupsLite( + new=[], + update_metadata=[ + models.GroupUpdateMetadataLite( + title="Assays", + values=models.GroupUpdateMetadataValuesLite( + title="Assay Results", missing_column_policy=MissingColumnPolicy.INHERIT + ), + ), + models.GroupUpdateMetadataLite( + title="Waste", + values=models.GroupUpdateMetadataValuesLite(parent_group="", tags={}), + ), + ], + delete=["Old Group"], + ) + self.assert_any_request_made( + method=RequestMethod.PATCH, + path=f"{self.base_path}/block-models/{BM_UUID}/blocks", + body=expected_update_body.model_dump(mode="json", exclude_unset=True), + headers=DEFAULT_EXPECTED_HEADERS | {"API-Preview": "opt-in"}, + ) + + async def test_update_groups_requires_an_operation(self) -> None: + with self.assertRaises(ValueError): + await self.preview_client_without_cache.update_groups(BM_UUID) + + async def test_update_groups_job_failed(self) -> None: + self.transport.set_request_handler( + UpdateRequestHandler( + update_result=UPDATE_RESULT, + job_response=JobResponse( + job_status=JobStatus.FAILED, + payload=models.JobErrorPayload( + detail="Group update failed", + status=500, + title="Group update failed", + type="https://seequent.com/error-codes/block-model-service/job/internal-error", + ), + ), + ) + ) + with self.assertRaises(JobFailedException): + await self.preview_client_without_cache.update_groups(BM_UUID, delete=["Assays"])