diff --git a/code-samples/geoscience-objects/simplified-object-interactions/create-downhole-collection.ipynb b/code-samples/geoscience-objects/simplified-object-interactions/create-downhole-collection.ipynb index b7a6d9b3..74092c65 100644 --- a/code-samples/geoscience-objects/simplified-object-interactions/create-downhole-collection.ipynb +++ b/code-samples/geoscience-objects/simplified-object-interactions/create-downhole-collection.ipynb @@ -152,7 +152,7 @@ " \"count\": [4],\n", " }\n", " ).astype({\"hole_index\": np.int32, \"offset\": np.uint64, \"count\": np.uint64}),\n", - " distance_table=pd.DataFrame(\n", + " table=pd.DataFrame(\n", " {\n", " \"distance\": [0.0, 10.0, 20.0, 30.0],\n", " \"attr_str\": [\"a\", \"b\", \"a\", \"c\"],\n", @@ -160,6 +160,7 @@ " \"attr_num\": [1.1, 2.2, 3.3, 4.4],\n", " }\n", " ),\n", + " unit=\"m\",\n", ")\n", "\n", "dhc_data = DownholeCollectionData(\n", diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index fe3a7fba..ab93e4d3 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -69,6 +69,42 @@ You can also get a list of all objects. Internally, this recursively calls the ` Check out the other methods on the `ObjectAPIClient` for more details on how to upload and download objects, or get object versions. +### Typed downhole collections + +`DownholeCollection` provides a DataFrame-based API for creating and reading downhole objects. A collection can contain +distance tables, interval tables, or both. `hole_index` is an integer key in the object's `hole_id` lookup table, not a +row position or a pandas categorical code. The creation helper uses dense zero-based keys in sorted hole-ID order unless +an explicit `hole_indices` mapping is supplied. Persisted objects may use arbitrary integer keys. Collection holes may +contain a subset or repeated chunks. + +```python +import pandas as pd + +from evo.objects.typed import DownholeCollection, DownholeCollectionData +from evo.objects.typed.downhole_collection import IntervalCollection +from evo.objects.utils.downhole import hole_chunks_from_ids + +intervals = pd.DataFrame({"from": [0.0], "to": [1.5], "lithology": ["sandstone"]}) +collections = [ + IntervalCollection( + name="geology", + holes=hole_chunks_from_ids(pd.Series(["DH-01"])), + table=intervals, + unit="m", # Collection coordinate units are set explicitly here. + ) +] + +# Build DownholeCollectionData with matching path, location-hole chunks, and collar properties, +# then create it with: await DownholeCollection.create(context, data) +``` + +Use `await dhc.location.to_dataframe()` and `await dhc.location.path_to_dataframe()` to read collars and paths. +Distance and interval tables provide `to_dataframe()` and `to_dataframe_by_hole()`. Before reading a large object, call +`await dhc.prefetch_collections("geology")` to warm only the requested collection data (and location data by default). +Attribute descriptions for ordinary attributes round-trip through `DataFrame.attrs["attribute_descriptions"]`. +Collection coordinate units come only from `DistanceCollection.unit` or `IntervalCollection.unit`; unit metadata on the +distance or `from` column is ignored. + ## Contributing For instructions on contributing to the development of this library, please refer to the [evo-python-sdk documentation](https://github.com/seequentevo/evo-python-sdk). diff --git a/packages/evo-objects/pyproject.toml b/packages/evo-objects/pyproject.toml index 7b2081c1..bbf0d214 100644 --- a/packages/evo-objects/pyproject.toml +++ b/packages/evo-objects/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "evo-objects" description = "Python SDK for using the Seequent Evo Geoscience Object API" -version = "0.6.1" +version = "0.6.2" requires-python = ">=3.10" license-files = ["LICENSE.md"] dynamic = ["readme"] diff --git a/packages/evo-objects/src/evo/objects/typed/__init__.py b/packages/evo-objects/src/evo/objects/typed/__init__.py index d13d270b..501cce39 100644 --- a/packages/evo-objects/src/evo/objects/typed/__init__.py +++ b/packages/evo-objects/src/evo/objects/typed/__init__.py @@ -12,17 +12,24 @@ from ._grid import BlockModelData, BlockModelGeometry from .attributes import ( Attribute, + AttributeDescription, Attributes, BlockModelAttribute, BlockModelAttributes, BlockModelPendingAttribute, PendingAttribute, + Unit, ) from .base import BaseObject, object_from_path, object_from_reference, object_from_uuid from .block_model_ref import ( BlockModel, ) -from .downhole_collection import DownholeCollection, DownholeCollectionData +from .downhole_collection import ( + DistanceCollection, + DownholeCollection, + DownholeCollectionData, + IntervalCollection, +) from .downhole_intervals import DownholeIntervals, DownholeIntervalsData from .pointset import ( Locations, @@ -70,6 +77,7 @@ __all__ = [ "Attribute", + "AttributeDescription", "Attributes", "BaseObject", "BaseSpatialObject", @@ -82,6 +90,7 @@ "BoundingBox", "CoordinateReferenceSystem", "CubicStructure", + "DistanceCollection", "DownholeCollection", "DownholeCollectionData", "DownholeIntervals", @@ -92,6 +101,7 @@ "ExponentialStructure", "GaussianStructure", "GeneralisedCauchyStructure", + "IntervalCollection", "LinearStructure", "Locations", "MaskedCells", @@ -110,6 +120,7 @@ "SpheroidalStructure", "Tensor3DGrid", "Tensor3DGridData", + "Unit", "Variogram", "VariogramCurveData", "VariogramData", diff --git a/packages/evo-objects/src/evo/objects/typed/_model.py b/packages/evo-objects/src/evo/objects/typed/_model.py index bc4476f4..85d0fb61 100644 --- a/packages/evo-objects/src/evo/objects/typed/_model.py +++ b/packages/evo-objects/src/evo/objects/typed/_model.py @@ -450,15 +450,30 @@ def __init_subclass__(cls, **kwargs: Any) -> None: if get_origin(base) is SchemaList: args = get_args(base) if args: - cls._item_type = args[0] + item_type = args[0] + if isinstance(item_type, type): + cls._item_type = item_type + elif "_resolve_item_type" not in cls.__dict__ or "_data_to_schema" not in cls.__dict__: + raise TypeError( + "SchemaList with a non-class item type must override _resolve_item_type() and _data_to_schema()" + ) break def __getitem__(self, index: int) -> _M: - return self._item_type(self._context, self._document[index]) + document = self._document[index] + return self._resolve_item_type(document)(self._context, document) def __iter__(self): for item in self._document: - yield self._item_type(self._context, item) + yield self._resolve_item_type(item)(self._context, item) + + @classmethod + def _resolve_item_type(cls, document: dict[str, Any]) -> type[_M]: + """Resolve the schema model used for a document item. + + Subclasses with heterogeneous document items can override this hook. + """ + return cls._item_type def __len__(self) -> int: return len(self._document) diff --git a/packages/evo-objects/src/evo/objects/typed/_prefetch.py b/packages/evo-objects/src/evo/objects/typed/_prefetch.py new file mode 100644 index 00000000..d0edc839 --- /dev/null +++ b/packages/evo-objects/src/evo/objects/typed/_prefetch.py @@ -0,0 +1,72 @@ +# Copyright © 2026 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. + +"""Helpers for warming typed-object data in the local cache.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Sequence +from typing import Any + +from evo.common import IFeedback +from evo.common.utils import NoFeedback +from evo.objects import DownloadedObject +from evo.objects.io import _CACHE_SCOPE + + +def collect_data_ids(documents: Any) -> list[str]: + """Collect de-duplicated ``data`` references from document values.""" + ids: list[str] = [] + + def visit(value: Any) -> None: + if isinstance(value, dict): + for key, nested in value.items(): + if key == "data" and isinstance(nested, str) and nested not in ids: + ids.append(nested) + else: + visit(nested) + elif isinstance(value, (list, tuple)): + for nested in value: + visit(nested) + + visit(documents) + return ids + + +async def prefetch_object_data( + obj: DownloadedObject, + *, + data_ids: Sequence[str] | None = None, + max_concurrent: int = 8, + fb: IFeedback = NoFeedback, +) -> None: + """Warm cache entries referenced by an object, downloading each ID at most once.""" + if max_concurrent < 1: + raise ValueError("max_concurrent must be at least 1") + cache = obj.get_cache() + if cache is None: + raise ValueError("prefetch requires an IContext with a cache") + identifiers = list(dict.fromkeys(data_ids if data_ids is not None else collect_data_ids(obj.as_dict()))) + if not identifiers: + return + cache_location = cache.get_location(environment=obj.get_environment(), scope=_CACHE_SCOPE) + identifiers = [identifier for identifier in identifiers if not (cache_location / identifier).exists()] + if not identifiers: + return + contexts = list(obj.prepare_data_download(identifiers)) + semaphore = asyncio.Semaphore(max_concurrent) + + async def download(context: Any) -> None: + async with semaphore: + await context.download_to_cache(cache, obj.get_connector().transport, fb=fb) + + await asyncio.gather(*(download(context) for context in contexts)) diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 8ea42595..3f35a4ff 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -14,7 +14,7 @@ import typing import uuid from dataclasses import dataclass -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, Protocol, runtime_checkable from uuid import UUID import pandas as pd @@ -41,11 +41,13 @@ __all__ = [ "Attribute", + "AttributeDescription", "Attributes", "BlockModelAttribute", "BlockModelAttributes", "BlockModelPendingAttribute", "PendingAttribute", + "Unit", ] @@ -53,6 +55,13 @@ class UnSupportedDataTypeError(Exception): """An unsupported data type was encountered while processing data.""" +@runtime_checkable +class Unit(Protocol): + """A schema unit enum represented by its string value.""" + + value: str + + def _infer_attribute_type_from_series(series: pd.Series) -> str: """Infer the attribute type from a Pandas Series. @@ -89,13 +98,20 @@ def _infer_attribute_type_from_series(series: pd.Series) -> str: class AttributeDescription: discipline: str = "" type: str = "" - unit: str | None = None + unit: str | Unit | None = None scale: str | None = None extensions: dict[str, typing.Any] | None = None tags: dict[str, str] | None = None - def to_schema(self): - result = { + def __post_init__(self) -> None: + if self.unit is None or isinstance(self.unit, str): + return + if not isinstance(self.unit, Unit) or not isinstance(self.unit.value, str): + raise TypeError("unit must be a str, a Unit with a string value, or None") + self.unit = self.unit.value + + def to_schema(self) -> dict[str, Any]: + result: dict[str, Any] = { "discipline": self.discipline, "type": self.type, } @@ -117,6 +133,7 @@ class Attribute(SchemaModel): _attribute_type: Annotated[str, SchemaLocation("attribute_type")] _key: Annotated[str | None, SchemaLocation("key")] _data: Annotated[str, SchemaLocation("values.data")] + _attribute_description: Annotated[dict[str, Any] | None, SchemaLocation("attribute_description")] @property def key(self) -> str: @@ -132,6 +149,12 @@ def attribute_type(self) -> str: """The type of this attribute.""" return self._attribute_type + @property + def attribute_description(self) -> AttributeDescription | None: + """Optional descriptive metadata associated with this attribute.""" + raw = self._attribute_description + return AttributeDescription(**raw) if raw else None + @property def exists(self) -> bool: """Whether this attribute exists on the object. @@ -241,7 +264,7 @@ class Attributes(SchemaList[Attribute]): attribute descriptions are attached to the DataFrame's `attrs` attribute. >>> df.attrs - {'attribute_description': {: }, ...} + {'attribute_descriptions': {: }, ...} """ _schema_path: str | None = None @@ -320,8 +343,8 @@ async def _upload_attributes_to_list( if attr_desc is not None: if not isinstance(attr_desc, AttributeDescription): raise TypeError("attribute description must be a AttributeDescription.") - if attr_desc.unit is not None: - attr_doc["attribute_description"] = attr_desc.to_schema() + if description := attr_desc.to_schema(): + attr_doc["attribute_description"] = description attributes_list.append(attr_doc) @@ -337,7 +360,15 @@ async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataF """ attributes = [self[key] for key in keys] if keys else list(self) parts = [await attribute.to_dataframe(fb=fb_part) for attribute, fb_part in iter_with_fb(attributes, fb)] - return pd.concat(parts, axis=1) if len(parts) > 0 else pd.DataFrame() + result = pd.concat(parts, axis=1) if len(parts) > 0 else pd.DataFrame() + descriptions = { + attribute.name: description + for attribute in attributes + if isinstance(attribute, Attribute) and (description := attribute.attribute_description) is not None + } + if descriptions: + result.attrs["attribute_descriptions"] = descriptions + return result async def append_attribute(self, df: pd.DataFrame, fb: IFeedback = NoFeedback): """Add a new attribute to the object. @@ -601,3 +632,14 @@ async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: if self._context.is_data_modified(self._data): raise DataLoaderError("Data was modified since the object was downloaded") return await self._obj.download_category_dataframe(self.as_dict(), fb=fb) + + async def to_indexed_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Load the persisted category lookup as ``[key, value]`` rows. + + Pandas categorical codes are dense positional values and must not be used + as schema lookup keys. This method is intended for joins involving an + index column such as ``hole_index``. + """ + if self._context.is_data_modified(self._data): + raise DataLoaderError("Data was modified since the object was downloaded") + return await self._obj.download_dataframe(self.as_dict()["table"], fb=fb) diff --git a/packages/evo-objects/src/evo/objects/typed/base.py b/packages/evo-objects/src/evo/objects/typed/base.py index c194cbd2..534365df 100644 --- a/packages/evo-objects/src/evo/objects/typed/base.py +++ b/packages/evo-objects/src/evo/objects/typed/base.py @@ -14,15 +14,18 @@ import copy import sys import weakref +from collections.abc import Sequence from dataclasses import dataclass from typing import Annotated, Any, ClassVar from uuid import UUID from evo import jmespath -from evo.common import IContext, StaticContext +from evo.common import IContext, IFeedback, StaticContext +from evo.common.utils import NoFeedback from evo.objects import DownloadedObject, ObjectMetadata, ObjectReference, ObjectSchema, SchemaVersion from ._model import ModelContext, SchemaLocation, SchemaModel +from ._prefetch import prefetch_object_data from ._utils import ( create_geoscience_object, replace_geoscience_object, @@ -416,6 +419,17 @@ class BaseObject(_BaseObject): tags: Annotated[dict[str, str], SchemaLocation("tags")] = {} extensions: Annotated[dict, SchemaLocation("extensions")] = {} + async def prefetch( + self, + *, + data_ids: Sequence[str] | None = None, + max_concurrent: int = 8, + fb: IFeedback = NoFeedback, + ) -> None: + """Warm cached data files referenced by this object.""" + + await prefetch_object_data(self._obj, data_ids=data_ids, max_concurrent=max_concurrent, fb=fb) + @classmethod def create( cls, diff --git a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py index 5fa5c870..d5d2c83e 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -11,6 +11,7 @@ from __future__ import annotations +import warnings from dataclasses import dataclass from typing import Annotated, Any, ClassVar, TypeAlias @@ -18,15 +19,16 @@ import pandas as pd from numpy._typing import NDArray +from evo.common import IFeedback from evo.common.interfaces import IContext +from evo.common.utils import NoFeedback from evo.objects import SchemaVersion +from evo.objects.typed import attributes as _attributes from evo.objects.typed._data import DataTable, DataTableAndAttributes -from evo.objects.typed._downhole import HoleIdCategory +from evo.objects.typed._downhole import DepthIntervalsTable, HoleIdCategory from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel -from evo.objects.typed.attributes import ( - AttributeDescription, - Attributes, -) +from evo.objects.typed._prefetch import collect_data_ids +from evo.objects.typed.attributes import Attributes from evo.objects.typed.exceptions import ObjectValidationError from evo.objects.typed.spatial import BaseSpatialObject, BaseSpatialObjectData from evo.objects.typed.types import BoundingBox @@ -34,12 +36,16 @@ DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, FLOAT_ARRAY_3, + INTEGER_ARRAY_1_INT32, + LOOKUP_TABLE_INT32, KnownTableFormat, ) __all__ = [ + "DistanceCollection", "DownholeCollection", "DownholeCollectionData", + "IntervalCollection", ] _X = "x" @@ -49,39 +55,60 @@ HolePath: TypeAlias = pd.DataFrame # [ distance | dip | azimuth | ] -HoleChunks: TypeAlias = pd.DataFrame # [ hole_id | offset | count ] +HoleChunks: TypeAlias = pd.DataFrame # [ hole_index | offset | count ] HoleProperties: TypeAlias = pd.DataFrame # [ hole_id | final | target | current | x | y | z ] HoleAttributes: TypeAlias = pd.DataFrame -# If `Depths` has unit descriptions in its `DataFrame.attrs` dictionary, then those units will be used when building -# the schema object. -# This is the expected structure: -# >>> depths_df.attrs -# {'attribute_description': {: }, ...} Depths: TypeAlias = pd.DataFrame # [ distance | ] +Intervals: TypeAlias = pd.DataFrame # [ from | to | ] @dataclass class DistanceCollection: name: str holes: HoleChunks - distance_table: Depths + table: Depths + unit: str | None collection_type: str = "distance" +@dataclass +class IntervalCollection: + name: str + holes: HoleChunks + table: Intervals + unit: str | None + collection_type: str = "interval" + + +DownholeCollectionEntry: TypeAlias = DistanceCollection | IntervalCollection + + +def _property_hole_indices(properties: HoleProperties) -> pd.Series: + hole_ids = properties["hole_id"].astype(str) + if "hole_index" in properties: + indices = properties["hole_index"].astype(int) + if indices.duplicated().any(): + raise ObjectValidationError("properties hole_index values must be unique") + return indices + dense = {hole_id: index for index, hole_id in enumerate(sorted(hole_ids.unique()))} + return hole_ids.map(dense).astype(int) + + @dataclass(kw_only=True, frozen=True) class DownholeCollectionData(BaseSpatialObjectData): - """Data class for creating a new DownholeCollection + """Data class representing a DownholeCollection. :param name: The name of the object. :param holes: A DataFrame describing which parts of `path` belong to which holes. - Columns: hole_id, offset, count - :param properties: DataFrame for the properties of the holes. The ith row corresponds to the ith element of `holes`. - Mandatory columns: hole_id, final, target, current, x, y, z - :param attributes: DataFrame for the attributes of the holes. The ith row corresponds to the ith element of `holes`. + Columns: hole_index, offset, count. ``hole_index`` is an integer lookup key, not a row position. + :param properties: DataFrame for the properties of the holes, with one row per unique hole ID. + Mandatory columns: hole_id, final, target, current, x, y, z. An optional hole_index column supplies explicit + lookup keys; otherwise dense zero-based keys are generated in sorted hole-ID order. + :param attributes: DataFrame for the attributes of the holes, in the same order as ``properties``. :param path: Dataframe of [ distance | dip | azimuth | ]. Distance/dip/azimuth describe the geometry as the step since the previous row. - :param collections: A list of `DistanceCollection` describing a table of distances with attributes. + :param collections: Distance and interval collection tables. :param distance_unit: The distance unit for the `path` table and the `properties` x/y/y. :param desurvey: The desurvey method appropriate for this collection. Must be one of: "minimum_curvature", "balanced_tangent", "trench". @@ -95,26 +122,96 @@ class DownholeCollectionData(BaseSpatialObjectData): holes: HoleChunks properties: HoleProperties attributes: HoleAttributes | None - collections: list[DistanceCollection] + collections: list[DownholeCollectionEntry] distance_unit: str | None desurvey: str | None def __post_init__(self): - if self.attributes is not None and len(self.holes) != len(self.attributes): - raise ObjectValidationError("The number of attributes rows must match the number or holes rows") + hole_ids = self.properties["hole_id"] + if hole_ids.isna().any(): + raise ObjectValidationError("properties hole_id values cannot be missing") + if hole_ids.astype(str).duplicated().any(): + raise ObjectValidationError("properties hole_id values must be unique") + + if self.attributes is not None and len(self.properties) != len(self.attributes): + raise ObjectValidationError("The number of attributes rows must match the number of properties rows") + + names = [collection.name for collection in self.collections] + if len(names) != len(set(names)): + raise ObjectValidationError("Collection names must be unique") + + valid_indices = set(_property_hole_indices(self.properties)) + self._validate_hole_chunks( + self.holes, + len(self.path), + valid_indices=valid_indices, + require_hole_coverage=True, + require_table_coverage=True, + ) + for collection in self.collections: + table = collection.table + self._validate_hole_chunks( + collection.holes, + len(table), + valid_indices=valid_indices, + require_hole_coverage=False, + require_table_coverage=False, + ) - assert self.attributes is None or len(self.holes) == len(self.attributes) + @staticmethod + def _validate_hole_chunks( + holes: HoleChunks, + table_length: int, + *, + valid_indices: set[int], + require_hole_coverage: bool, + require_table_coverage: bool, + ) -> None: + required = {"hole_index", "offset", "count"} + if missing := required - set(holes.columns): + raise ObjectValidationError(f"Hole chunks are missing columns: {sorted(missing)}") + indices = holes["hole_index"].astype(int) + if not set(indices).issubset(valid_indices): + raise ObjectValidationError("hole_index must reference a valid hole lookup key") + if require_hole_coverage and (indices.duplicated().any() or set(indices) != valid_indices): + raise ObjectValidationError("Location holes must contain each hole lookup key exactly once") + offsets = holes["offset"].astype(int) + counts = holes["count"].astype(int) + if (offsets < 0).any() or (counts < 0).any() or ((offsets + counts) > table_length).any(): + raise ObjectValidationError("Hole chunk offsets and counts must be within the associated table") + if not require_table_coverage: + return + + non_empty = sorted(zip(offsets[counts > 0], counts[counts > 0], strict=True)) + expected_offset = 0 + for offset, count in non_empty: + if offset != expected_offset: + raise ObjectValidationError("Hole chunk ranges must cover the associated table exactly once") + expected_offset = offset + count + if expected_offset != table_length: + raise ObjectValidationError("Hole chunk ranges must cover the associated table exactly once") + + def _collars_by_hole_index(self) -> dict[int, tuple[float, float, float]]: + """Map declared hole lookup keys to collar coordinates.""" + keys = _property_hole_indices(self.properties) + coordinates = self.properties[_COORDINATE_COLUMNS].astype(float) + return { + int(key): coordinate + for key, coordinate in zip(keys, coordinates.itertuples(index=False, name=None), strict=True) + } def compute_bounding_box(self) -> BoundingBox: + collars_by_hole_index = self._collars_by_hole_index() bboxes = [] - - for i in range(len(self.holes)): - offset = self.holes.iat[i, 1] - count = self.holes.iat[i, 2] - collar = tuple(self.properties.loc[i, _COORDINATE_COLUMNS]) - path_table = self.path[offset : offset + count] - bboxes.append(self._compute_hole_bounding_box(path_table, collar)) - + for chunk in self.holes.itertuples(index=False): + collar = collars_by_hole_index[int(chunk.hole_index)] + offset = int(chunk.offset) + count = int(chunk.count) + if count: + bbox = self._compute_hole_bounding_box(self.path.iloc[offset : offset + count], collar) + else: + bbox = BoundingBox.from_points(np.asarray([collar])) + bboxes.append(bbox) return BoundingBox.combine(bboxes) @staticmethod @@ -246,14 +343,49 @@ async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: return await super()._data_to_schema(distances_df, context) +class LocationHoleIdCategory(HoleIdCategory): + """Collar hole IDs persisted with explicit integer lookup keys.""" + + @classmethod + async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: + hole_ids = data["hole_id"].astype(str) + if hole_ids.duplicated().any(): + raise ObjectValidationError("properties hole_id values must be unique") + keys = _property_hole_indices(data).astype("int32") + data_client = _attributes.get_data_client(context) + return { + "values": await data_client.upload_dataframe( + pd.DataFrame({"hole_id": keys}), table_format=INTEGER_ARRAY_1_INT32 + ), + "table": await data_client.upload_dataframe( + pd.DataFrame({"key": keys, "value": hole_ids}), table_format=LOOKUP_TABLE_INT32 + ), + } + + class DownholeLocation(SchemaModel): - hole_id: Annotated[HoleIdCategory, SchemaLocation("hole_id"), DataLocation("properties")] + hole_id: Annotated[LocationHoleIdCategory, SchemaLocation("hole_id"), DataLocation("properties")] path: Annotated[DownholePath, SchemaLocation("path"), DataLocation("path")] holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] distances: Annotated[DistancesTable, SchemaLocation("distances"), DataLocation("properties")] coordinates: Annotated[CollarCoordinates, SchemaLocation("coordinates"), DataLocation("properties")] attributes: Annotated[Attributes, SchemaLocation("attributes"), DataLocation("attributes")] + async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return collars with resolved ``hole_id`` values and selected attributes.""" + parts = [ + await self.hole_id.to_dataframe(fb=fb), + await self.coordinates.to_dataframe(fb=fb), + await self.distances.to_dataframe(fb=fb), + ] + if len(self.attributes): + parts.append(await self.attributes.to_dataframe(*keys, fb=fb)) + return pd.concat(parts, axis=1) + + async def path_to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return the desurvey path and its attributes.""" + return await self.path.to_dataframe(*keys, fb=fb) + class _Distances(DataTable): table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_1 @@ -264,28 +396,163 @@ class DistanceTableDistances(DataTableAndAttributes): _table: Annotated[_Distances, SchemaLocation("values"), DataLocation("")] unit: Annotated[str | None, SchemaLocation("unit")] - @classmethod - async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: - result = await super()._data_to_schema(data, context) - attr_desc: AttributeDescription = data.attrs.get("attribute_descriptions", {}).get("distance") - if attr_desc is not None and attr_desc.unit is not None: - # "unit" can be missing, but it must not be `None` - result["unit"] = attr_desc.unit - return result - class DistanceTable(SchemaModel): name: Annotated[str, SchemaLocation("name"), DataLocation("name")] collection_type: Annotated[str, SchemaLocation("collection_type"), DataLocation("collection_type")] - distance: Annotated[DistanceTableDistances, SchemaLocation("distance"), DataLocation("distance_table")] + distance: Annotated[DistanceTableDistances, SchemaLocation("distance"), DataLocation("table")] + + async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return distance values and selected attributes.""" + return await self.distance.to_dataframe(*keys, fb=fb) + + +class _DownholeCollectionChild: + """Shared behavior for models that must be nested in a DownholeCollection.""" + + @property + def _downhole_collection(self) -> DownholeCollection: + """Return the containing DownholeCollection or raise if this model is used out of context.""" + root = self._context.root_model + if not isinstance(root, DownholeCollection): + raise ObjectValidationError( + f"{type(self).__name__} must be attached to a DownholeCollection, not {type(root).__name__}" + ) + return root + + async def _table_by_hole(self, data: pd.DataFrame, *, fb: IFeedback) -> dict[str, pd.DataFrame]: + """Group table rows by hole using the persisted hole-id lookup.""" + lookup = await self._downhole_collection.location.hole_id.to_indexed_dataframe(fb=fb) + hole_ids = dict(zip(lookup["key"].astype(int), lookup["value"].astype(str), strict=True)) + result: dict[str, list[tuple[int, pd.DataFrame]]] = {} + for chunk in (await self.holes.to_dataframe(fb=fb)).itertuples(index=False): + try: + hole_id = hole_ids[int(chunk.hole_index)] + except KeyError as exc: + raise ObjectValidationError(f"Unknown hole_index in collection chunks: {chunk.hole_index}") from exc + offset = int(chunk.offset) + result.setdefault(hole_id, []).append( + (offset, data.iloc[offset : offset + int(chunk.count)].reset_index(drop=True)) + ) + return { + hole_id: pd.concat([chunk for _, chunk in sorted(chunks, key=lambda item: item[0])], ignore_index=True) + for hole_id, chunks in result.items() + } + + +class DownholeDistanceTable(_DownholeCollectionChild, DistanceTable): + holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] + + async def to_dataframe_by_hole(self, *keys: str, fb: IFeedback = NoFeedback) -> dict[str, pd.DataFrame]: + """Return per-hole collection values and selected attributes.""" + return await self._table_by_hole(await self.to_dataframe(*keys, fb=fb), fb=fb) + + +class IntervalTableFromTo(DataTableAndAttributes): + _table: Annotated[DepthIntervalsTable, SchemaLocation("intervals.start_and_end"), DataLocation("")] + unit: Annotated[str | None, SchemaLocation("unit")] -class DownholeDistanceTable(DistanceTable): +class DownholeIntervalTable(_DownholeCollectionChild, SchemaModel): + name: Annotated[str, SchemaLocation("name"), DataLocation("name")] + collection_type: Annotated[str, SchemaLocation("collection_type"), DataLocation("collection_type")] + from_to: Annotated[IntervalTableFromTo, SchemaLocation("from_to"), DataLocation("table")] holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] + async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return collection intervals and selected attributes.""" + return await self.from_to.to_dataframe(*keys, fb=fb) -class DownholeCollectionTables(SchemaList[DownholeDistanceTable]): - pass + async def to_dataframe_by_hole(self, *keys: str, fb: IFeedback = NoFeedback) -> dict[str, pd.DataFrame]: + """Return per-hole collection intervals and selected attributes.""" + return await self._table_by_hole(await self.to_dataframe(*keys, fb=fb), fb=fb) + + +class DownholeCollectionTables(_DownholeCollectionChild, SchemaList[DownholeDistanceTable | DownholeIntervalTable]): + @classmethod + def _resolve_item_type(cls, document: dict[str, Any]) -> type[DownholeDistanceTable | DownholeIntervalTable]: + return DownholeIntervalTable if "from_to" in document else DownholeDistanceTable + + @classmethod + async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: + if data is None: + return [] + result = [] + for collection in data: + model = DownholeIntervalTable if isinstance(collection, IntervalCollection) else DownholeDistanceTable + schema = await model._data_to_schema(collection, context) + if collection.unit is not None: + if isinstance(collection, IntervalCollection): + schema["from_to"]["unit"] = collection.unit + else: + schema["distance"]["unit"] = collection.unit + result.append(schema) + return result + + def _indices_for_name(self, name: str) -> list[int]: + return [index for index, document in enumerate(self._document) if document.get("name") == name] + + def get(self, name: str) -> DownholeDistanceTable | DownholeIntervalTable | None: + """Return the first collection named ``name``. + + Collection names are the only schema-provided identifier. Legacy documents may contain duplicate names; in + that case the first collection in document order is returned and a warning is emitted. + """ + indices = self._indices_for_name(name) + if len(indices) > 1: + warnings.warn( + f"Multiple collections named '{name}' were found; returning the first one", + UserWarning, + ) + return self[indices[0]] if indices else None + + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and bool(self._indices_for_name(name)) + + def names(self) -> list[str]: + """Return collection names in document order.""" + return [collection.name for collection in self] + + async def add(self, collection: DownholeCollectionEntry, *, replace: bool = False) -> None: + """Add ``collection`` or replace an existing collection with the same name. + + Set ``replace`` to ``True`` to replace a single existing collection in place. Collection hole-index keys must + reference holes in the containing downhole collection. + """ + existing_indices = self._indices_for_name(collection.name) + if len(existing_indices) > 1: + raise ObjectValidationError( + f"Multiple collections named '{collection.name}' already exist; remove duplicates before adding a replacement" + ) + if existing_indices and not replace: + raise ObjectValidationError(f"Collection '{collection.name}' already exists") + lookup = await self._downhole_collection.location.hole_id.to_indexed_dataframe() + table = collection.table + valid_indices = set(lookup["key"].astype(int)) + DownholeCollectionData._validate_hole_chunks( + collection.holes, + len(table), + valid_indices=valid_indices, + require_hole_coverage=False, + require_table_coverage=False, + ) + schema = await self._data_to_schema([collection], self._obj) + if existing_indices: + self._document[existing_indices[0]] = schema[0] + else: + self._document.append(schema[0]) + + def remove(self, *names: str) -> int: + """Remove collections whose names are supplied and return the number of documents removed. + + The count includes every matching legacy document when duplicate collection names are present. Unknown names do + not raise an error and contribute zero to the result. To require that a name exists, check membership before + calling this method. + """ + requested = set(names) + previous = len(self._document) + self._document[:] = [item for item in self._document if item.get("name") not in requested] + return previous - len(self._document) class DownholeCollection(BaseSpatialObject): @@ -301,3 +568,22 @@ class DownholeCollection(BaseSpatialObject): desurvey: Annotated[str | None, SchemaLocation("desurvey")] type: ClassVar[Annotated[str, SchemaLocation("type")]] = "downhole" + + async def prefetch_collections( + self, + *names: str, + include_location: bool = True, + max_concurrent: int = 8, + fb: IFeedback = NoFeedback, + ) -> None: + """Prefetch data referenced by named collections and optionally location data.""" + + documents = [] + if include_location: + documents.append(self.location.as_dict()) + for name in names: + collection = self.collections.get(name) + if collection is None: + raise KeyError(f"Unknown collection '{name}'") + documents.append(collection.as_dict()) + await self.prefetch(data_ids=collect_data_ids(documents), max_concurrent=max_concurrent, fb=fb) diff --git a/packages/evo-objects/src/evo/objects/utils/__init__.py b/packages/evo-objects/src/evo/objects/utils/__init__.py index dde35e2c..af087953 100644 --- a/packages/evo-objects/src/evo/objects/utils/__init__.py +++ b/packages/evo-objects/src/evo/objects/utils/__init__.py @@ -23,6 +23,7 @@ DataFrame = None # type: ignore from .data import ObjectDataClient +from .downhole import expand_hole_index, hole_chunks_from_ids from .table_formats import all_known_formats, get_known_format from .tables import ArrowTableFormat, BaseTableFormat, KnownTableFormat from .types import ArrayTableInfo, AttributeInfo, CategoryInfo, LookupTableInfo, TableInfo @@ -43,5 +44,7 @@ "ObjectDataClient", "TableInfo", "all_known_formats", + "expand_hole_index", "get_known_format", + "hole_chunks_from_ids", ] diff --git a/packages/evo-objects/src/evo/objects/utils/data.py b/packages/evo-objects/src/evo/objects/utils/data.py index 045b4572..d22970bb 100644 --- a/packages/evo-objects/src/evo/objects/utils/data.py +++ b/packages/evo-objects/src/evo/objects/utils/data.py @@ -17,7 +17,7 @@ import numpy as np from evo import logging -from evo.common import APIConnector, Environment, ICache, IFeedback +from evo.common import APIConnector, Environment, ICache, IFeedback, StaticContext from evo.common.exceptions import StorageFileNotFoundError from evo.common.io.exceptions import DataExistsError from evo.common.utils import NoFeedback, PartialFeedback, split_feedback @@ -117,6 +117,15 @@ def clear_cache(self) -> None: """Clear the cache used by this client.""" self._cache.clear_cache(environment=self._environment, scope=_CACHE_SCOPE) + def get_static_context(self) -> StaticContext: + """Build a typed-object context from this client's environment and infrastructure.""" + return StaticContext( + connector=self._connector, + cache=self._cache, + org_id=self._environment.org_id, + workspace_id=self._environment.workspace_id, + ) + async def upload_referenced_data(self, object_model: dict, fb: IFeedback = NoFeedback) -> None: """Upload all data referenced by a geoscience object. diff --git a/packages/evo-objects/src/evo/objects/utils/downhole.py b/packages/evo-objects/src/evo/objects/utils/downhole.py new file mode 100644 index 00000000..81337ef2 --- /dev/null +++ b/packages/evo-objects/src/evo/objects/utils/downhole.py @@ -0,0 +1,76 @@ +# Copyright © 2026 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. + +"""Utilities for the indexed tables used by downhole collections.""" + +from __future__ import annotations + +from collections.abc import Mapping + +import numpy as np +import pandas as pd + +__all__ = ["expand_hole_index", "hole_chunks_from_ids"] + + +def hole_chunks_from_ids(hole_ids: pd.Series, *, hole_indices: Mapping[str, int] | None = None) -> pd.DataFrame: + """Run-length encode contiguous hole IDs into ``[hole_index, offset, count]`` chunks. + + ``hole_indices`` maps hole IDs to their lookup-table keys. When omitted, this + creation helper assigns dense zero-based keys in sorted ID order. Persisted + objects are not required to use those keys. Only IDs present in ``hole_ids`` + produce chunks; an empty collection therefore produces no chunks. + """ + if hole_ids.isna().any(): + raise ValueError("hole_ids cannot contain missing values") + ids = hole_ids.astype(str) + if hole_indices is None: + indices = {hole_id: index for index, hole_id in enumerate(sorted(set(ids)))} + else: + indices = dict(hole_indices) + unknown = sorted(set(ids) - set(indices)) + if unknown: + raise ValueError(f"hole_ids contains values absent from hole_indices: {unknown}") + if len(set(indices.values())) != len(indices): + raise ValueError("hole_indices must map each hole_id to a unique hole_index") + codes = ids.map(indices).to_numpy(dtype=np.int32) + chunks: list[tuple[int, int, int]] = [] + start = 0 + while start < len(codes): + code = int(codes[start]) + end = start + 1 + while end < len(codes) and codes[end] == code: + end += 1 + if any(chunk_code == code for chunk_code, _, _ in chunks): + raise ValueError("Rows for each hole_id must be contiguous") + chunks.append((code, start, end - start)) + start = end + return pd.DataFrame( + { + "hole_index": np.array([code for code, _, _ in chunks], dtype=np.int32), + "offset": np.array([offset for _, offset, _ in chunks], dtype=np.uint64), + "count": np.array([count for _, _, count in chunks], dtype=np.uint64), + } + ) + + +def expand_hole_index(holes: pd.DataFrame, num_rows: int) -> pd.Series: + """Expand ``hole_index`` chunks into a per-row nullable integer Series.""" + required = {"hole_index", "offset", "count"} + if missing := required - set(holes.columns): + raise ValueError(f"holes is missing columns: {sorted(missing)}") + result = pd.Series(pd.array([pd.NA] * num_rows, dtype="Int32")) + for chunk in holes.itertuples(index=False): + offset, count = int(chunk.offset), int(chunk.count) + if offset < 0 or count < 0 or offset + count > num_rows: + raise ValueError("Hole chunk offsets and counts must be within num_rows") + result.iloc[offset : offset + count] = int(chunk.hole_index) + return result diff --git a/packages/evo-objects/tests/helpers.py b/packages/evo-objects/tests/helpers.py index 459aa83e..f3e604ef 100644 --- a/packages/evo-objects/tests/helpers.py +++ b/packages/evo-objects/tests/helpers.py @@ -104,7 +104,7 @@ def _generate_timestamp_data(n_samples: int) -> Iterator[datetime]: min_ = datetime(1970, 1, 1, tzinfo=timezone.utc).timestamp() max_ = datetime(2038, 12, 31, 23, 59, 59, 999999, tzinfo=timezone.utc).timestamp() for _ in range(n_samples): - yield datetime.utcfromtimestamp(random.uniform(min_, max_)) + yield datetime.fromtimestamp(random.uniform(min_, max_), timezone.utc) def _generate_data(format_id: str, n_samples: int) -> Iterator: diff --git a/packages/evo-objects/tests/test_downhole_utils.py b/packages/evo-objects/tests/test_downhole_utils.py new file mode 100644 index 00000000..a4a8ab9d --- /dev/null +++ b/packages/evo-objects/tests/test_downhole_utils.py @@ -0,0 +1,125 @@ +# Copyright © 2026 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 pandas as pd + +from evo.objects.utils.downhole import expand_hole_index, hole_chunks_from_ids + + +class TestDownholeUtils(unittest.TestCase): + def _chunks(self, values, hole_indices=None): + return hole_chunks_from_ids(pd.Series(values), hole_indices=hole_indices) + + def test_single_hole_uses_zero_based_code(self): + chunks = self._chunks(["A"]) + self.assertListEqual(chunks.to_dict("records"), [{"hole_index": 0, "offset": 0, "count": 1}]) + + def test_multiple_rows_for_one_hole_are_one_chunk(self): + chunks = self._chunks(["A", "A", "A"]) + self.assertListEqual(chunks.to_dict("records"), [{"hole_index": 0, "offset": 0, "count": 3}]) + + def test_contiguous_holes_have_consecutive_offsets(self): + chunks = self._chunks(["A", "A", "B", "B"]) + self.assertListEqual( + chunks.to_dict("records"), + [{"hole_index": 0, "offset": 0, "count": 2}, {"hole_index": 1, "offset": 2, "count": 2}], + ) + + def test_varying_hole_counts_are_preserved(self): + chunks = self._chunks(["A", "B", "B", "B", "C", "C"]) + self.assertListEqual( + chunks.to_dict("records"), + [ + {"hole_index": 0, "offset": 0, "count": 1}, + {"hole_index": 1, "offset": 1, "count": 3}, + {"hole_index": 2, "offset": 4, "count": 2}, + ], + ) + + def test_mapping_ids_absent_from_data_do_not_emit_chunks(self): + chunks = self._chunks(["A", "A"], {"A": 1, "B": 42}) + self.assertListEqual(chunks.to_dict("records"), [{"hole_index": 1, "offset": 0, "count": 2}]) + + def test_id_absent_from_hole_indices_raises(self): + with self.assertRaises(ValueError): + self._chunks(["A", "C"], {"A": 1, "B": 42}) + + def test_empty_input_emits_no_chunks(self): + chunks = self._chunks([], {"A": 1, "B": 42}) + self.assertListEqual(chunks.to_dict("records"), []) + + def test_hole_index_dtype_is_int32(self): + self.assertEqual(str(self._chunks(["A"])["hole_index"].dtype), "int32") + + def test_offset_and_count_dtypes_are_uint64(self): + chunks = self._chunks(["A"]) + self.assertEqual(str(chunks["offset"].dtype), "uint64") + self.assertEqual(str(chunks["count"].dtype), "uint64") + + def test_fifty_holes_with_two_hundred_rows_each_round_trip(self): + hole_ids = [f"H{index:02d}" for index in range(50)] + values = [hole_id for hole_id in hole_ids for _ in range(200)] + chunks = self._chunks(values) + self.assertEqual(len(chunks), 50) + self.assertListEqual(chunks["offset"].tolist(), list(range(0, 10_000, 200))) + self.assertListEqual(chunks["count"].tolist(), [200] * 50) + self.assertListEqual( + expand_hole_index(chunks, len(values)).tolist(), [index for index in range(50) for _ in range(200)] + ) + + def test_creation_keys_are_sorted_by_hole_id(self): + chunks = self._chunks(["M", "M", "Z"]) + self.assertListEqual( + chunks.to_dict("records"), + [ + {"hole_index": 0, "offset": 0, "count": 2}, + {"hole_index": 1, "offset": 2, "count": 1}, + ], + ) + + def test_chunks_use_dense_creation_keys_and_round_trip(self): + values = pd.Series(["M", "M", "Z"]) + chunks = hole_chunks_from_ids(values) + self.assertListEqual(chunks["hole_index"].tolist(), [0, 1]) + self.assertListEqual(chunks["count"].tolist(), [2, 1]) + self.assertListEqual(expand_hole_index(chunks, len(values)).tolist(), [0, 0, 1]) + + def test_chunks_preserve_explicit_sparse_lookup_keys(self): + chunks = self._chunks(["M", "M", "Z"], {"Z": 1, "M": 42}) + self.assertListEqual( + chunks.to_dict("records"), + [{"hole_index": 42, "offset": 0, "count": 2}, {"hole_index": 1, "offset": 2, "count": 1}], + ) + + def test_non_contiguous_and_unknown_ids_raise(self): + with self.assertRaises(ValueError): + hole_chunks_from_ids(pd.Series(["A", "B", "A"]), hole_indices={"A": 1, "B": 42}) + with self.assertRaises(ValueError): + hole_chunks_from_ids(pd.Series(["C"]), hole_indices={"A": 1, "B": 42}) + + def test_empty_input_uses_schema_dtypes(self): + chunks = hole_chunks_from_ids(pd.Series([], dtype="string"), hole_indices={"A": 1, "B": 42}) + self.assertListEqual(chunks["hole_index"].tolist(), []) + self.assertListEqual(chunks["offset"].tolist(), []) + self.assertListEqual(chunks["count"].tolist(), []) + self.assertEqual(str(chunks["hole_index"].dtype), "int32") + self.assertEqual(str(chunks["offset"].dtype), "uint64") + self.assertEqual(str(chunks["count"].dtype), "uint64") + + def test_expand_rejects_out_of_bounds_chunks(self): + with self.assertRaises(ValueError): + expand_hole_index(pd.DataFrame({"hole_index": [0], "offset": [1], "count": [2]}), 2) + + def test_expand_preserves_hole_indices(self): + holes = pd.DataFrame({"hole_index": [1, 42], "offset": [0, 1], "count": [1, 2]}) + self.assertListEqual(expand_hole_index(holes, 3).tolist(), [1, 42, 42]) diff --git a/packages/evo-objects/tests/test_tables.py b/packages/evo-objects/tests/test_tables.py index c0fd9955..416f5c22 100644 --- a/packages/evo-objects/tests/test_tables.py +++ b/packages/evo-objects/tests/test_tables.py @@ -27,6 +27,7 @@ from evo.common.test_tools import BASE_URL from evo.common.utils import Cache from evo.objects.exceptions import SchemaValidationError, TableFormatError +from evo.objects.parquet import ParquetLoader from evo.objects.utils import ( ArrowTableFormat, BaseTableFormat, @@ -82,6 +83,13 @@ def _get_table_digest(table: pa.Table) -> str: return _get_buffer_digest(buffer=buffer) +def _load_table(table_info: dict, source: Path) -> pa.Table: + with pa.OSFile(str(source / str(table_info["data"])), mode="r") as parquet_file: + with ParquetLoader(parquet_file) as loader: + loader.validate_with_table_info(table_info) + return loader.load_as_table() + + def _test_name_from_known_format(cls: type, num: int, params_dict: dict) -> str: """Create parameterized test name from DataFormat""" data_format = params_dict["data_format"] @@ -185,7 +193,7 @@ def test_load_table(self) -> None: self.assertTrue(self.parquet_file.is_file()) inferred_format = KnownTableFormat.from_table_info(table_info) - actual_table = KnownTableFormat.load_table(table_info, self.data_dir) + actual_table = _load_table(table_info, self.data_dir) self.assertEqual(inferred_format.width, actual_table.num_columns) self.assertEqual(table_info["length"], actual_table.num_rows) @@ -196,34 +204,34 @@ def test_load_table_with_too_many_columns(self) -> None: table_info = self._save_parquet_file(add_column=True) self.assertTrue(self.parquet_file.is_file()) - with self.assertRaises(TableFormatError): - KnownTableFormat.load_table(table_info, self.data_dir) + with self.assertRaises(SchemaValidationError): + _load_table(table_info, self.data_dir) def test_load_table_with_too_many_rows(self) -> None: table_info = self._save_parquet_file(add_row=True) self.assertTrue(self.parquet_file.is_file()) with self.assertRaises(SchemaValidationError): - KnownTableFormat.load_table(table_info, self.data_dir) + _load_table(table_info, self.data_dir) def test_load_table_with_wrong_data_types(self) -> None: table_info = self._save_parquet_file(change_type=True) self.assertTrue(self.parquet_file.is_file()) - with self.assertRaises(TableFormatError): - KnownTableFormat.load_table(table_info, self.data_dir) + with self.assertRaises(SchemaValidationError): + _load_table(table_info, self.data_dir) - def test_load_table_from_uuid(self) -> None: + def test_load_table_with_uuid_identifier(self) -> None: table_info = self._save_parquet_file() self.assertTrue(self.parquet_file.is_file()) - table_info["data"] = uuid.uuid4() - self.parquet_file = self.parquet_file.rename(self.parquet_file.parent / str(table_info["data"])) + table_info["data"] = str(uuid.uuid4()) + self.parquet_file = self.parquet_file.rename(self.parquet_file.parent / table_info["data"]) self.assertTrue(self.parquet_file.is_file()) inferred_format = KnownTableFormat.from_table_info(table_info) - actual_table = KnownTableFormat.load_table(table_info, self.data_dir) + actual_table = _load_table(table_info, self.data_dir) self.assertEqual(inferred_format.width, actual_table.num_columns) self.assertEqual(table_info["length"], actual_table.num_rows) diff --git a/packages/evo-objects/tests/typed/helpers.py b/packages/evo-objects/tests/typed/helpers.py index 630766a3..004f4a3b 100644 --- a/packages/evo-objects/tests/typed/helpers.py +++ b/packages/evo-objects/tests/typed/helpers.py @@ -68,7 +68,13 @@ async def download_attribute_dataframe(self, data: dict, fb) -> pd.DataFrame: return self.mock_client.get_dataframe(data["values"]) async def download_category_dataframe(self, category_info: dict, fb) -> pd.DataFrame: - return self.mock_client.get_dataframe(category_info["values"]) + values = self.mock_client.get_dataframe(category_info["values"]) + if not pd.api.types.is_integer_dtype(values.iloc[:, 0]): + # Legacy mock uploads store category values directly rather than as lookup keys. + return values + lookup = self.mock_client.get_dataframe(category_info["table"]) + mapping = dict(zip(lookup["key"], lookup["value"])) + return values.apply(lambda column: column.map(mapping)) async def download_array(self, jmespath_expr: str, fb=None): """Download an array from the object using a JMESPath expression.""" @@ -84,6 +90,8 @@ async def download_array(self, jmespath_expr: str, fb=None): async def update(self, object_dict): new_version_id = str(int(self.metadata.version_id) + 1) + persisted = copy.deepcopy(object_dict) + self.mock_client.objects[persisted["uuid"]] = persisted return MockDownloadedObject(self.mock_client, object_dict, new_version_id) @@ -110,8 +118,13 @@ async def upload_table(self, table, *args, **kwargs) -> dict: return {"data": data_id, "length": len(table)} async def upload_category_dataframe(self, df: pd.DataFrame, *args, **kwargs) -> dict: + series = df.iloc[:, 0].astype("category") + categories = series.cat.categories return { - "values": await self.upload_dataframe(df), + "values": await self.upload_dataframe(pd.DataFrame({df.columns[0]: series})), + "table": await self.upload_dataframe( + pd.DataFrame({"key": range(len(categories)), "value": categories.astype(str)}) + ), "category_data": True, } diff --git a/packages/evo-objects/tests/typed/test_attributes.py b/packages/evo-objects/tests/typed/test_attributes.py index c8928cec..1a33c32c 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -11,12 +11,18 @@ from __future__ import annotations +from typing import Any, cast from unittest import TestCase import pandas as pd from parameterized import parameterized -from evo.objects.typed.attributes import PendingAttribute, UnSupportedDataTypeError, _infer_attribute_type_from_series +from evo.objects.typed.attributes import ( + AttributeDescription, + PendingAttribute, + UnSupportedDataTypeError, + _infer_attribute_type_from_series, +) class TestAttributeTypeInference(TestCase): @@ -65,3 +71,27 @@ def test_pending_attribute_repr(self): """Test that PendingAttribute has a useful repr.""" pending = PendingAttribute(None, "test_attr") self.assertEqual(repr(pending), "PendingAttribute(name='test_attr', exists=False)") + + +class TestAttributeDescription(TestCase): + def test_default_description_preserves_required_fields(self): + self.assertEqual(AttributeDescription().to_schema(), {"discipline": "", "type": ""}) + + def test_unitless_description_omits_unit(self): + description = AttributeDescription(discipline="Geology", type="Azimuth") + self.assertEqual(description.to_schema(), {"discipline": "Geology", "type": "Azimuth"}) + + def test_description_preserves_string_units(self): + description = AttributeDescription(discipline="geology", type="length", unit="m") + self.assertEqual(description.to_schema(), {"discipline": "geology", "type": "length", "unit": "m"}) + + def test_description_normalizes_value_units(self): + class Unit: + value = "m" + + description = AttributeDescription(discipline="geology", type="length", unit=Unit()) + self.assertEqual(description.to_schema(), {"discipline": "geology", "type": "length", "unit": "m"}) + + def test_description_rejects_invalid_units(self): + with self.assertRaisesRegex(TypeError, "unit must be"): + AttributeDescription(unit=cast(Any, object())) diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 563b08fb..e31edc47 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -12,11 +12,14 @@ from __future__ import annotations import contextlib +import copy import dataclasses +import inspect import math import uuid +import warnings from datetime import date -from unittest.mock import patch +from unittest.mock import AsyncMock, patch import numpy as np import numpy.testing as npt @@ -25,13 +28,16 @@ from evo.common import Environment, StaticContext from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector +from evo.common.utils import NoFeedback from evo.objects import ObjectReference from evo.objects.typed import BoundingBox +from evo.objects.typed.attributes import AttributeDescription from evo.objects.typed.base import BaseObject from evo.objects.typed.downhole_collection import ( DistanceCollection, DownholeCollection, DownholeCollectionData, + IntervalCollection, ) from evo.objects.typed.exceptions import ObjectValidationError @@ -65,7 +71,7 @@ def _make_example_data( "count": [4], } ), - distance_table=pd.DataFrame( + table=pd.DataFrame( { "distance": [0.0, 10.0, 20.0, 30.0], "attr_str": ["a", "b", "a", "c"], @@ -73,6 +79,7 @@ def _make_example_data( "attr_num": [1.1, 2.2, 3.3, 4.4], } ), + unit=None, ) holes = pd.DataFrame( @@ -113,6 +120,17 @@ def _make_example_data( class TestDownholeCollection(TestWithConnector): + def test_nullable_domain_arguments_are_required(self): + for data_class, arguments in ( + (DistanceCollection, ("unit",)), + (IntervalCollection, ("unit",)), + (DownholeCollectionData, ("distance_unit", "desurvey")), + ): + parameters = inspect.signature(data_class).parameters + for argument in arguments: + with self.subTest(data_class=data_class.__name__, argument=argument): + self.assertIs(parameters[argument].default, inspect.Parameter.empty) + def setUp(self) -> None: TestWithConnector.setUp(self) self.environment = Environment(hub_url=BASE_URL, org_id=ORG.id, workspace_id=WORKSPACE_ID) @@ -152,6 +170,7 @@ async def _check_locations(self, expected: DownholeCollectionData, result: Downh npt.assert_array_equal(expected.properties[xyz], await loc.coordinates.to_dataframe()) npt.assert_array_equal(expected.properties[distances], await loc.distances.to_dataframe()) npt.assert_array_equal(expected.properties[["hole_id"]], await loc.hole_id.to_dataframe()) + npt.assert_array_equal(expected.properties[["hole_id"]], (await loc.to_dataframe())[["hole_id"]]) npt.assert_array_equal(expected.holes, await loc.holes.to_dataframe()) if expected.attributes: npt.assert_array_equal(expected.attributes, await result.location.hole_id.to_dataframe()) @@ -167,12 +186,10 @@ async def _check_path(self, expected: DownholeCollectionData, result: DownholeCo async def _check_collections(self, expected: DownholeCollectionData, result: DownholeCollection): for expected_collection, result_collection in zip(expected.collections, result.collections, strict=True): - expected_distance_unit = expected_collection.distance_table.attrs.get("attribute_descriptions", {}).get( - "distance" - ) + expected_distance_unit = expected_collection.table.attrs.get("attribute_descriptions", {}).get("distance") self.assertEqual(expected_distance_unit, result_collection.distance.unit) - expected_table = expected_collection.distance_table + expected_table = expected_collection.table result_table = await result_collection.distance.to_dataframe() for col in result_table.columns: @@ -208,6 +225,190 @@ async def test_create_with_empty_collections(self): result = await DownholeCollection.create(context=self.context, data=data) self.assertIsInstance(result, DownholeCollection) + async def test_mixed_collections_round_trip_and_mutation(self): + interval = IntervalCollection( + name="intervals", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), + table=pd.DataFrame({"from": [0.0, 1.0], "to": [1.0, 2.0], "lithology": ["a", "b"]}), + unit="m", + ) + data = _make_example_data(collections=[_make_example_data().collections[0], interval]) + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create(context=self.context, data=data) + self.assertEqual(result.collections.names(), ["collection1", "intervals"]) + self.assertEqual(result.collections.get("intervals").from_to.unit, "m") + self.assertListEqual( + (await result.collections.get("intervals").to_dataframe()).columns.tolist(), ["from", "to", "lithology"] + ) + self.assertEqual(result.collections.remove("missing", "collection1"), 1) + await result.collections.add(data.collections[0]) + self.assertEqual(result.collections.names(), ["intervals", "collection1"]) + object_json = mock_client.objects[str(result.metadata.url.object_id)] + self.assertIn("start_and_end", object_json["collections"][1]["from_to"]["intervals"]) + + async def test_interval_collection_ignores_coordinate_attribute_metadata(self): + table = pd.DataFrame({"from": [0.0, 1.0], "to": [1.0, 2.0], "grade": [1.0, 2.0]}) + table.attrs["attribute_descriptions"] = {"from": AttributeDescription(unit="ft")} + interval = IntervalCollection( + name="intervals", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), + table=table, + unit=None, + ) + data = _make_example_data(collections=[interval]) + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create(context=self.context, data=data) + collection = result.collections.get("intervals") + self.assertIsNotNone(collection) + self.assertIsNone(collection.from_to.unit) + self.assertListEqual((await collection.to_dataframe()).columns.tolist(), ["from", "to", "grade"]) + document = mock_client.objects[str(result.metadata.url.object_id)] + self.assertNotIn("unit", document["collections"][0]["from_to"]) + + @parameterized.expand( + [ + ("explicit", "m", "m"), + ("omitted", None, None), + ] + ) + async def test_interval_collection_uses_explicit_unit_only(self, _name, explicit_unit, expected_unit): + table = pd.DataFrame({"from": [0.0], "to": [1.0]}) + collection = IntervalCollection( + name="intervals", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=table, + unit=explicit_unit, + ) + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create( + context=self.context, data=_make_example_data(collections=[collection]) + ) + document = mock_client.objects[str(result.metadata.url.object_id)] + from_to = document["collections"][0]["from_to"] + self.assertEqual(result.collections.get("intervals").from_to.unit, expected_unit) + if expected_unit is None: + self.assertNotIn("unit", from_to) + else: + self.assertEqual(from_to["unit"], expected_unit) + + async def test_interval_collection_round_trips_all_attribute_types(self): + table = pd.DataFrame( + { + "from": [0.0, 1.0], + "to": [1.0, 2.0], + "scalar": pd.Series([1.5, 2.5], dtype="float64"), + "integer": pd.Series([1, 2], dtype="int64"), + "boolean": pd.Series([True, False], dtype="bool"), + "string": pd.Series(["a", "b"], dtype="string"), + "category": pd.Series(pd.Categorical(["ore", "waste"])), + "date_time": pd.to_datetime(["2026-01-01", "2026-01-02"]), + } + ) + collection = IntervalCollection( + name="intervals", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), + table=table, + unit=None, + ) + expected_types = { + "scalar": "scalar", + "integer": "integer", + "boolean": "bool", + "string": "string", + "category": "category", + "date_time": "date_time", + } + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create( + context=self.context, data=_make_example_data(collections=[collection]) + ) + document = mock_client.objects[str(result.metadata.url.object_id)] + interval = result.collections.get("intervals") + self.assertIsNotNone(interval) + self.assertEqual( + { + attribute["name"]: attribute["attribute_type"] + for attribute in document["collections"][0]["from_to"]["attributes"] + }, + expected_types, + ) + self.assertListEqual((await interval.to_dataframe()).columns.tolist(), list(table.columns)) + + @parameterized.expand([("missing_from", {"to": [1.0]}), ("missing_to", {"from": [0.0]})]) + async def test_interval_collection_requires_from_and_to_columns(self, _name, table_data): + collection = IntervalCollection( + name="invalid", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=pd.DataFrame(table_data), + unit=None, + ) + with self._mock_geoscience_objects(), self.assertRaises(ObjectValidationError): + await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[collection])) + + @parameterized.expand( + [ + ("explicit", "m", "m"), + ("metadata_only", None, None), + ] + ) + async def test_distance_collection_uses_explicit_unit_only(self, _name, collection_unit, expected_unit): + table = pd.DataFrame({"distance": [0.0], "grade": [1.0]}) + table.attrs["attribute_descriptions"] = {"distance": AttributeDescription(unit="ft")} + collection = DistanceCollection( + name="distances", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=table, + unit=collection_unit, + ) + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create( + context=self.context, data=_make_example_data(collections=[collection]) + ) + self.assertEqual(result.collections.get("distances").distance.unit, expected_unit) + document = mock_client.objects[str(result.metadata.url.object_id)] + distance = document["collections"][0]["distance"] + if expected_unit is None: + self.assertNotIn("unit", distance) + else: + self.assertEqual(distance["unit"], expected_unit) + + async def test_attribute_descriptions_round_trip_to_dataframe_metadata(self): + table = pd.DataFrame({"distance": [0.0], "grade": [1.0]}) + description = AttributeDescription( + discipline="geology", + type="grade", + unit="ppm", + scale="linear", + tags={"source": "assay"}, + ) + table.attrs["attribute_descriptions"] = {"grade": description} + collection = DistanceCollection( + name="grades", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=table, + unit=None, + ) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create( + context=self.context, data=_make_example_data(collections=[collection]) + ) + round_tripped = await result.collections.get("grades").to_dataframe() + self.assertEqual(round_tripped.attrs["attribute_descriptions"]["grade"], description) + + async def test_attribute_without_description_reads_as_none(self): + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=_make_example_data()) + self.assertIsNone(result.collections.get("collection1").distance.attributes["attr_str"].attribute_description) + + async def test_none_optional_fields_are_omitted(self): + data = _make_example_data() + data = dataclasses.replace(data, distance_unit=None, desurvey=None) + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create(context=self.context, data=data) + document = mock_client.objects[str(result.metadata.url.object_id)] + self.assertNotIn("distance_unit", document) + self.assertNotIn("desurvey", document) + @parameterized.expand([BaseObject, DownholeCollection]) async def test_replace(self, class_to_call): data = _make_example_data() @@ -260,6 +461,18 @@ def test_bounding_box(self): bbox = data.compute_bounding_box() self._assert_bounding_box_equal(bbox, 100.0, 200.0, 150.0, 300.0, -30.0, 50.0) + def test_bounding_box_uses_explicit_hole_index_mapping(self): + data = _make_example_data() + expected = data.compute_bounding_box() + properties = data.properties.iloc[[1, 0]].copy() + properties.index = [10, 20] + properties["hole_index"] = [1, 0] + data = dataclasses.replace(data, properties=properties) + bbox = data.compute_bounding_box() + self._assert_bounding_box_equal( + bbox, expected.min_x, expected.max_x, expected.min_y, expected.max_y, expected.min_z, expected.max_z + ) + def test_bounding_box_from_spiral(self): # First hole spirals, second hole zig-zags path = pd.DataFrame( @@ -327,7 +540,7 @@ async def test_description_and_tags(self): self.assertEqual(result.tags, {"site": "alpha", "status": "active"}) def test_attributes_length_raises(self): - """attributes length must match holes length.""" + """attributes length must match properties length.""" path = pd.DataFrame({"distance": [0.0, 10.0], "azimuth": [0.0, 0.0], "dip": [90.0, 90.0]}) holes = pd.DataFrame({"hole_index": [0, 1], "offset": [0, 1], "count": [1, 1]}) properties = pd.DataFrame( @@ -355,6 +568,280 @@ def test_attributes_length_raises(self): desurvey=None, ) + def test_location_chunks_require_each_lookup_key_and_exact_path_coverage(self): + base = _make_example_data(collections=[]) + for holes in ( + pd.DataFrame({"hole_index": [0], "offset": [0], "count": [7]}), + pd.DataFrame({"hole_index": [0, 0], "offset": [0, 4], "count": [4, 3]}), + pd.DataFrame({"hole_index": [1, 2], "offset": [0, 4], "count": [4, 3]}), + pd.DataFrame({"hole_index": [0, 1], "offset": [0, 5], "count": [4, 2]}), + pd.DataFrame({"hole_index": [0, 1], "offset": [0, 3], "count": [4, 3]}), + ): + with self.subTest(holes=holes.to_dict("records")), self.assertRaises(ObjectValidationError): + dataclasses.replace(base, holes=holes) + + def test_collection_chunk_ranges_are_not_required_to_partition_the_table(self): + base = _make_example_data(collections=[]) + table = pd.DataFrame({"distance": [0.0, 1.0, 2.0]}) + collection = DistanceCollection( + name="repeated", + holes=pd.DataFrame({"hole_index": [0, 0], "offset": [0, 2], "count": [1, 1]}), + table=table, + unit=None, + ) + dataclasses.replace(base, collections=[collection]) + + invalid_collection = DistanceCollection( + name="invalid-index", + holes=pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), + table=table, + unit=None, + ) + with self.assertRaisesRegex(ObjectValidationError, "hole_index must reference a valid hole lookup key"): + dataclasses.replace(base, collections=[invalid_collection]) + + async def test_location_holes_may_be_out_of_hole_index_order(self): + data = _make_example_data(collections=[]) + holes = pd.DataFrame({"hole_index": [1, 0], "offset": [4, 0], "count": [3, 4]}) + collection = DistanceCollection( + name="measurements", + holes=pd.DataFrame({"hole_index": [1], "offset": [0], "count": [1]}), + table=pd.DataFrame({"distance": [1.0]}), + unit=None, + ) + data = dataclasses.replace(data, holes=holes, collections=[collection]) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=data) + self.assertListEqual(list((await result.collections.get("measurements").to_dataframe_by_hole())), ["H002"]) + self.assertEqual(data.compute_bounding_box(), _make_example_data(collections=[]).compute_bounding_box()) + + async def test_data_and_reads_support_one_based_sparse_lookup_keys(self): + data = _make_example_data(collections=[]) + properties = data.properties.assign(hole_index=[1, 42]) + holes = pd.DataFrame({"hole_index": [1, 42], "offset": [0, 4], "count": [4, 3]}) + collection = DistanceCollection( + name="measurements", + holes=pd.DataFrame({"hole_index": [42], "offset": [0], "count": [1]}), + table=pd.DataFrame({"distance": [1.0]}), + unit=None, + ) + data = dataclasses.replace(data, properties=properties, holes=holes, collections=[collection]) + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create(context=self.context, data=data) + lookup_info = result.location.hole_id.as_dict()["table"] + self.assertListEqual(mock_client.data[lookup_info["data"]]["key"].tolist(), [1, 42]) + self.assertListEqual((await result.location.to_dataframe())["hole_id"].tolist(), ["H001", "H002"]) + self.assertListEqual(list((await result.collections.get("measurements").to_dataframe_by_hole())), ["H002"]) + self.assertEqual(data.compute_bounding_box(), _make_example_data(collections=[]).compute_bounding_box()) + + def test_undeclared_lookup_keys_are_rejected(self): + base = _make_example_data(collections=[]) + with self.assertRaisesRegex(ObjectValidationError, "hole_index must reference a valid hole lookup key"): + dataclasses.replace( + base, + holes=pd.DataFrame({"hole_index": [7, 9], "offset": [0, 4], "count": [4, 3]}), + ) + + def test_bounding_box_includes_collar_for_hole_without_path_rows(self): + base = _make_example_data(collections=[]) + data = dataclasses.replace( + base, + path=base.path.iloc[:4], + holes=pd.DataFrame({"hole_index": [0, 1], "offset": [0, 4], "count": [4, 0]}), + ) + bbox = data.compute_bounding_box() + self._assert_bounding_box_equal(bbox, 100.0, 200.0, 150.0, 300.0, -30.0, 50.0) + + async def test_collection_allows_zero_row_hole(self): + base = _make_example_data(collections=[]) + collection = DistanceCollection( + name="sparse", + holes=pd.DataFrame({"hole_index": [0, 1], "offset": [0, 0], "count": [2, 0]}), + table=pd.DataFrame({"distance": [0.0, 1.0]}), + unit=None, + ) + data = dataclasses.replace(base, collections=[collection]) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=data) + by_hole = await result.collections.get("sparse").to_dataframe_by_hole() + self.assertEqual({hole_id: len(table) for hole_id, table in by_hole.items()}, {"H001": 2, "H002": 0}) + + async def test_collection_add_allows_repeated_hole_chunks_and_replaces_in_place(self): + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[])) + collection = DistanceCollection( + name="measurements", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=pd.DataFrame({"distance": [0.0]}), + unit=None, + ) + await result.collections.add(collection) + replacement = dataclasses.replace(collection, table=pd.DataFrame({"distance": [1.0]})) + with self.assertRaises(ObjectValidationError): + await result.collections.add(replacement) + await result.collections.add(replacement, replace=True) + self.assertEqual(result.collections.names(), ["measurements"]) + self.assertEqual((await result.collections.get("measurements").to_dataframe()).iloc[0, 0], 1.0) + repeated = dataclasses.replace( + collection, + holes=pd.DataFrame({"hole_index": [0, 0], "offset": [0, 1], "count": [1, 0]}), + ) + await result.collections.add(repeated, replace=True) + + def test_collection_creation_rejects_duplicate_names(self): + collection = _make_example_data().collections[0] + with self.assertRaisesRegex(ObjectValidationError, "Collection names must be unique"): + _make_example_data(collections=[collection, dataclasses.replace(collection)]) + + async def test_legacy_duplicate_collection_names_warn_and_return_first(self): + data = _make_example_data() + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create(context=self.context, data=data) + document = mock_client.objects[str(result.metadata.url.object_id)] + document["collections"].append(copy.deepcopy(document["collections"][0])) + legacy = await DownholeCollection.from_reference(context=self.context, reference=result.metadata.url) + + self.assertEqual(legacy.collections.names(), ["collection1", "collection1"]) + with warnings.catch_warnings(): + warnings.simplefilter("error") + self.assertIn("collection1", legacy.collections) + with self.assertWarnsRegex(UserWarning, "Multiple collections named 'collection1'"): + collection = legacy.collections.get("collection1") + self.assertIsNotNone(collection) + with self.assertRaisesRegex(ObjectValidationError, "Multiple collections named 'collection1' already exist"): + await legacy.collections.add(data.collections[0], replace=True) + + async def test_collection_replacement_preserves_position_and_persists_after_update(self): + first = DistanceCollection( + name="first", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=pd.DataFrame({"distance": [0.0]}), + unit=None, + ) + middle = DistanceCollection( + name="middle", + holes=pd.DataFrame({"hole_index": [1], "offset": [0], "count": [1]}), + table=pd.DataFrame({"distance": [1.0]}), + unit=None, + ) + last = DistanceCollection( + name="last", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=pd.DataFrame({"distance": [2.0]}), + unit=None, + ) + replacement = dataclasses.replace(middle, table=pd.DataFrame({"distance": [10.0]})) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[])) + await result.collections.add(first) + await result.collections.add(middle) + await result.collections.add(last) + await result.collections.add(replacement, replace=True) + self.assertEqual(result.collections.names(), ["first", "middle", "last"]) + await result.update() + persisted = await DownholeCollection.from_reference(context=self.context, reference=result.metadata.url) + self.assertEqual(persisted.collections.names(), ["first", "middle", "last"]) + self.assertEqual((await persisted.collections.get("middle").to_dataframe()).iloc[0, 0], 10.0) + + async def test_location_path_and_collection_reads(self): + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=_make_example_data()) + collars = await result.location.to_dataframe() + self.assertListEqual(collars.columns.tolist(), ["hole_id", "x", "y", "z", "final", "target", "current"]) + self.assertListEqual( + (await result.location.path_to_dataframe()).columns.tolist(), ["distance", "azimuth", "dip"] + ) + by_hole = await result.collections.get("collection1").to_dataframe_by_hole() + self.assertListEqual(list(by_hole), ["H001"]) + self.assertEqual(len(by_hole["H001"]), 4) + + async def test_prefetch_collections_exposes_and_forwards_explicit_options(self): + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=_make_example_data()) + + parameters = inspect.signature(DownholeCollection.prefetch_collections).parameters + self.assertNotIn("kwargs", parameters) + self.assertEqual(parameters["max_concurrent"].default, 8) + self.assertIs(parameters["fb"].default, NoFeedback) + + base_parameters = inspect.signature(BaseObject.prefetch).parameters + self.assertEqual(base_parameters["max_concurrent"].default, 8) + + with patch.object(DownholeCollection, "prefetch", new_callable=AsyncMock) as prefetch: + await result.prefetch_collections("collection1", include_location=False, max_concurrent=2, fb=NoFeedback) + + self.assertEqual(prefetch.await_count, 1) + self.assertEqual(prefetch.await_args.kwargs["max_concurrent"], 2) + self.assertIs(prefetch.await_args.kwargs["fb"], NoFeedback) + self.assertTrue(prefetch.await_args.kwargs["data_ids"]) + + async def test_dataframe_read_wrappers_select_attributes(self): + interval = IntervalCollection( + name="intervals", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + table=pd.DataFrame({"from": [0.0], "to": [1.0], "lithology": ["ore"]}), + unit=None, + ) + data = _make_example_data( + attributes=pd.DataFrame({"collar_attribute": ["first", "second"]}), + collections=[_make_example_data().collections[0], interval], + ) + data = dataclasses.replace(data, path=data.path.assign(path_attribute=[1] * len(data.path))) + + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=data) + + self.assertListEqual( + (await result.location.to_dataframe("collar_attribute")).columns.tolist(), + ["hole_id", "x", "y", "z", "final", "target", "current", "collar_attribute"], + ) + self.assertListEqual( + (await result.location.path_to_dataframe("path_attribute")).columns.tolist(), + ["distance", "azimuth", "dip", "path_attribute"], + ) + distance = result.collections.get("collection1") + self.assertIsNotNone(distance) + self.assertListEqual((await distance.to_dataframe("attr_num")).columns.tolist(), ["distance", "attr_num"]) + self.assertListEqual( + (await distance.to_dataframe_by_hole("attr_num"))["H001"].columns.tolist(), ["distance", "attr_num"] + ) + intervals = result.collections.get("intervals") + self.assertIsNotNone(intervals) + self.assertListEqual((await intervals.to_dataframe("lithology")).columns.tolist(), ["from", "to", "lithology"]) + self.assertListEqual( + (await intervals.to_dataframe_by_hole("lithology"))["H001"].columns.tolist(), ["from", "to", "lithology"] + ) + + async def test_collection_read_uses_persisted_non_contiguous_lookup_keys(self): + with self._mock_geoscience_objects() as mock_client: + result = await DownholeCollection.create(context=self.context, data=_make_example_data()) + lookup_info = result.location.hole_id.as_dict()["table"] + mock_client.data[lookup_info["data"]]["key"] = [10, 20] + collection = result.collections.get("collection1") + holes_info = collection.holes.as_dict() + mock_client.data[holes_info["data"]]["hole_index"] = [10] + + by_hole = await collection.to_dataframe_by_hole() + + self.assertListEqual(list(by_hole), ["H001"]) + self.assertEqual(len(by_hole["H001"]), 4) + + async def test_repeated_collection_chunks_are_concatenated_in_table_order(self): + collection = DistanceCollection( + name="repeated", + holes=pd.DataFrame({"hole_index": [0, 0], "offset": [2, 0], "count": [1, 1]}), + table=pd.DataFrame({"distance": [0.0, 1.0, 2.0]}), + unit=None, + ) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create( + context=self.context, + data=_make_example_data(collections=[collection]), + ) + by_hole = await result.collections.get("repeated").to_dataframe_by_hole() + + self.assertListEqual(by_hole["H001"]["distance"].tolist(), [0.0, 2.0]) + async def test_update_dataframe_after_creation(self): """Test updating the path DataFrame after downhole collection creation.""" with self._mock_geoscience_objects(): diff --git a/packages/evo-objects/tests/typed/test_model.py b/packages/evo-objects/tests/typed/test_model.py index 1243508c..a432590b 100644 --- a/packages/evo-objects/tests/typed/test_model.py +++ b/packages/evo-objects/tests/typed/test_model.py @@ -23,7 +23,7 @@ from evo.common import Environment, StaticContext from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector from evo.objects.typed._data import DataTable, DataTableAndAttributes -from evo.objects.typed._model import SchemaBuilder, SchemaLocation, SchemaModel +from evo.objects.typed._model import SchemaBuilder, SchemaList, SchemaLocation, SchemaModel from evo.objects.utils.table_formats import FLOAT_ARRAY_3, KnownTableFormat from .helpers import MockClient @@ -59,14 +59,26 @@ class FakeData: self.assertEqual(result["format_version"], "1.0.0") self.assertEqual(result["name"], "test") + def test_union_schema_list_requires_polymorphic_overrides(self): + class First(SchemaModel): + pass -class TestTable(DataTable): + class Second(SchemaModel): + pass + + with self.assertRaisesRegex(TypeError, r"must override _resolve_item_type\(\) and _data_to_schema\(\)"): + + class InvalidUnionList(SchemaList[First | Second]): + pass + + +class CoordinateTable(DataTable): table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 data_columns: ClassVar[list[str]] = ["x", "y", "z"] class ExtendedLocations(DataTableAndAttributes): - _table: Annotated[TestTable, SchemaLocation("coordinates")] + _table: Annotated[CoordinateTable, SchemaLocation("coordinates")] point_count: Annotated[int, SchemaLocation("point_count")] diff --git a/packages/evo-objects/tests/typed/test_prefetch.py b/packages/evo-objects/tests/typed/test_prefetch.py new file mode 100644 index 00000000..262f758d --- /dev/null +++ b/packages/evo-objects/tests/typed/test_prefetch.py @@ -0,0 +1,103 @@ +# Copyright © 2026 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. + +from __future__ import annotations + +import asyncio +import tempfile +from pathlib import Path +from types import SimpleNamespace +from unittest import IsolatedAsyncioTestCase + +from evo.objects.typed._prefetch import collect_data_ids, prefetch_object_data + + +class _Cache: + def __init__(self, path: Path) -> None: + self.path = path + + def get_location(self, **_kwargs) -> Path: + return self.path + + +class _Download: + def __init__(self, identifier: str, owner: _Object) -> None: + self.identifier = identifier + self.owner = owner + + async def download_to_cache(self, cache: _Cache, _transport, fb=None) -> None: + self.owner.active += 1 + self.owner.max_active = max(self.owner.max_active, self.owner.active) + try: + await asyncio.sleep(0.001) + (cache.path / self.identifier).touch() + self.owner.downloaded.append(self.identifier) + finally: + self.owner.active -= 1 + + +class _Object: + def __init__(self, document: dict, cache: _Cache) -> None: + self.document = document + self.cache = cache + self.downloaded: list[str] = [] + self.requested: list[str] = [] + self.active = 0 + self.max_active = 0 + + def as_dict(self) -> dict: + return self.document + + def get_cache(self) -> _Cache: + return self.cache + + def get_environment(self): + return object() + + def get_connector(self): + return SimpleNamespace(transport=object()) + + def prepare_data_download(self, identifiers): + self.requested.extend(identifiers) + return (_Download(identifier, self) for identifier in identifiers) + + +class TestPrefetch(IsolatedAsyncioTestCase): + async def test_prefetch_deduplicates_honours_concurrency_and_warms_cache(self): + with tempfile.TemporaryDirectory() as directory: + cache = _Cache(Path(directory)) + obj = _Object({"first": {"data": "a"}, "second": [{"data": "b"}, {"data": "a"}]}, cache) + await prefetch_object_data(obj, max_concurrent=1) + self.assertListEqual(obj.requested, ["a", "b"]) + self.assertListEqual(obj.downloaded, ["a", "b"]) + self.assertEqual(obj.max_active, 1) + + await prefetch_object_data(obj) + self.assertListEqual(obj.requested, ["a", "b"]) + + async def test_prefetch_collects_category_values_and_lookup_and_accepts_subset(self): + document = { + "category": {"values": {"data": "values"}, "lookup": {"data": "lookup"}}, + "other": {"data": "other"}, + } + self.assertListEqual(collect_data_ids(document), ["values", "lookup", "other"]) + with tempfile.TemporaryDirectory() as directory: + obj = _Object(document, _Cache(Path(directory))) + await prefetch_object_data(obj, data_ids=["lookup"]) + self.assertListEqual(obj.downloaded, ["lookup"]) + + async def test_prefetch_empty_document_is_a_no_op_and_rejects_invalid_concurrency(self): + with tempfile.TemporaryDirectory() as directory: + obj = _Object({}, _Cache(Path(directory))) + await prefetch_object_data(obj) + self.assertListEqual(obj.requested, []) + with self.assertRaises(ValueError): + await prefetch_object_data(obj, max_concurrent=0) diff --git a/packages/evo-objects/tests/typed/test_regular_grid.py b/packages/evo-objects/tests/typed/test_regular_grid.py index 506cb66e..20ac9411 100644 --- a/packages/evo-objects/tests/typed/test_regular_grid.py +++ b/packages/evo-objects/tests/typed/test_regular_grid.py @@ -59,7 +59,7 @@ def _mock_geoscience_objects(self): cell_data=pd.DataFrame( { "value": np.random.rand(10 * 10 * 5), - "cat": pd.Categorical(np.random.choice(range(4), size=10 * 10 * 5), ["a", "b", "c", "d"]), + "cat": pd.Categorical(np.random.choice(["a", "b", "c", "d"], size=10 * 10 * 5)), } ), vertex_data=pd.DataFrame( diff --git a/packages/evo-objects/tests/typed/test_regular_masked_grid.py b/packages/evo-objects/tests/typed/test_regular_masked_grid.py index 8ef6fb57..05e249d1 100644 --- a/packages/evo-objects/tests/typed/test_regular_masked_grid.py +++ b/packages/evo-objects/tests/typed/test_regular_masked_grid.py @@ -61,7 +61,7 @@ def _mock_geoscience_objects(self): { "value": np.random.rand(np.sum(example_mask)), "cat": pd.Categorical( - np.random.choice(range(4), size=np.sum(example_mask)), categories=["a", "b", "c", "d"] + np.random.choice(["a", "b", "c", "d"], size=np.sum(example_mask)), categories=["a", "b", "c", "d"] ), } ), diff --git a/packages/evo-objects/tests/typed/test_tensor_grid.py b/packages/evo-objects/tests/typed/test_tensor_grid.py index 93eca2b3..780eea7f 100644 --- a/packages/evo-objects/tests/typed/test_tensor_grid.py +++ b/packages/evo-objects/tests/typed/test_tensor_grid.py @@ -59,7 +59,7 @@ def _mock_geoscience_objects(self): cell_data=pd.DataFrame( { "value": np.random.rand(10 * 10 * 5), - "cat": pd.Categorical(np.random.choice(range(4), size=10 * 10 * 5), ["a", "b", "c", "d"]), + "cat": pd.Categorical(np.random.choice(["a", "b", "c", "d"], size=10 * 10 * 5)), } ), vertex_data=pd.DataFrame( diff --git a/uv.lock b/uv.lock index 80f8210f..6461c28d 100644 --- a/uv.lock +++ b/uv.lock @@ -1107,7 +1107,7 @@ test = [ [[package]] name = "evo-objects" -version = "0.6.1" +version = "0.6.2" source = { editable = "packages/evo-objects" } dependencies = [ { name = "evo-sdk-common", extra = ["jmespath"] },