From 532621ff4a202b637b08326b1342a06cf98b619c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Fri, 31 Jul 2026 09:25:08 +0100 Subject: [PATCH 01/30] feat(objects): add interval and indexed downhole collections --- .../src/evo/objects/typed/__init__.py | 11 +- .../src/evo/objects/typed/_model.py | 13 +- .../evo/objects/typed/downhole_collection.py | 210 +++++++++++++++++- .../src/evo/objects/utils/__init__.py | 3 + .../src/evo/objects/utils/downhole.py | 67 ++++++ .../evo-objects/tests/test_downhole_utils.py | 33 +++ .../tests/typed/test_downhole_collection.py | 42 ++++ 7 files changed, 364 insertions(+), 15 deletions(-) create mode 100644 packages/evo-objects/src/evo/objects/utils/downhole.py create mode 100644 packages/evo-objects/tests/test_downhole_utils.py diff --git a/packages/evo-objects/src/evo/objects/typed/__init__.py b/packages/evo-objects/src/evo/objects/typed/__init__.py index d13d270b..6f591b3b 100644 --- a/packages/evo-objects/src/evo/objects/typed/__init__.py +++ b/packages/evo-objects/src/evo/objects/typed/__init__.py @@ -12,6 +12,7 @@ from ._grid import BlockModelData, BlockModelGeometry from .attributes import ( Attribute, + AttributeDescription, Attributes, BlockModelAttribute, BlockModelAttributes, @@ -22,7 +23,12 @@ 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 +76,7 @@ __all__ = [ "Attribute", + "AttributeDescription", "Attributes", "BaseObject", "BaseSpatialObject", @@ -82,6 +89,7 @@ "BoundingBox", "CoordinateReferenceSystem", "CubicStructure", + "DistanceCollection", "DownholeCollection", "DownholeCollectionData", "DownholeIntervals", @@ -92,6 +100,7 @@ "ExponentialStructure", "GaussianStructure", "GeneralisedCauchyStructure", + "IntervalCollection", "LinearStructure", "Locations", "MaskedCells", diff --git a/packages/evo-objects/src/evo/objects/typed/_model.py b/packages/evo-objects/src/evo/objects/typed/_model.py index bc4476f4..e8a1f0b8 100644 --- a/packages/evo-objects/src/evo/objects/typed/_model.py +++ b/packages/evo-objects/src/evo/objects/typed/_model.py @@ -454,11 +454,20 @@ def __init_subclass__(cls, **kwargs: Any) -> None: 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/downhole_collection.py b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py index 8162da6a..2ce9c893 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -11,7 +11,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Annotated, Any, ClassVar, TypeAlias import numpy as np @@ -33,13 +33,16 @@ from evo.objects.utils.table_formats import ( DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, + FLOAT_ARRAY_2, FLOAT_ARRAY_3, KnownTableFormat, ) __all__ = [ + "DistanceCollection", "DownholeCollection", "DownholeCollectionData", + "IntervalCollection", ] _X = "x" @@ -59,6 +62,7 @@ # >>> depths_df.attrs # {'attribute_description': {: }, ...} Depths: TypeAlias = pd.DataFrame # [ distance | ] +Intervals: TypeAlias = pd.DataFrame # [ from | to | ] @dataclass @@ -67,6 +71,19 @@ class DistanceCollection: holes: HoleChunks distance_table: Depths collection_type: str = "distance" + unit: str | None = None + + +@dataclass +class IntervalCollection: + name: str + holes: HoleChunks + interval_table: Intervals + collection_type: str = "interval" + unit: str | None = None + + +DownholeCollectionEntry: TypeAlias = DistanceCollection | IntervalCollection @dataclass(kw_only=True, frozen=True) @@ -95,9 +112,17 @@ class DownholeCollectionData(BaseSpatialObjectData): holes: HoleChunks properties: HoleProperties attributes: HoleAttributes | None - collections: list[DistanceCollection] - distance_unit: str | None - desurvey: str | None + collections: list[DownholeCollectionEntry] + distance_unit: str | None = None + desurvey: str | None = None + + @property + def hole_id_dtype(self) -> pd.CategoricalDtype: + """The categorical dtype whose codes are used by every ``hole_index``.""" + hole_ids = self.properties["hole_id"] + if isinstance(hole_ids.dtype, pd.CategoricalDtype): + return hole_ids.dtype + return pd.CategoricalDtype(categories=sorted(hole_ids.dropna().unique())) def __post_init__(self): if self.attributes is not None and len(self.holes) != len(self.attributes): @@ -105,13 +130,40 @@ def __post_init__(self): assert self.attributes is None or len(self.holes) == len(self.attributes) + self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) + for collection in self.collections: + table = ( + collection.distance_table if isinstance(collection, DistanceCollection) else collection.interval_table + ) + self._validate_hole_chunks(collection.holes, len(table), require_coverage=False) + + def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require_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) + valid = set(range(len(self.hole_id_dtype.categories))) + if not set(indices).issubset(valid): + raise ObjectValidationError("hole_index must be a code in properties['hole_id'] categorical dtype") + if indices.duplicated().any(): + raise ObjectValidationError("Each hole_index may occur only once in a holes table") + if require_coverage and set(indices) != valid: + raise ObjectValidationError("Location holes must cover every hole_id categorical code 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") + def compute_bounding_box(self) -> BoundingBox: 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]) + collars = self.properties.copy() + collars["_hole_index"] = collars["hole_id"].astype(self.hole_id_dtype).cat.codes + collars_by_index = collars.set_index("_hole_index") + for chunk in self.holes.itertuples(index=False): + offset = int(chunk.offset) + count = int(chunk.count) + collar = tuple(collars_by_index.loc[int(chunk.hole_index), _COORDINATE_COLUMNS]) path_table = self.path[offset : offset + count] bboxes.append(self._compute_hole_bounding_box(path_table, collar)) @@ -254,6 +306,21 @@ class DownholeLocation(SchemaModel): coordinates: Annotated[CollarCoordinates, SchemaLocation("coordinates"), DataLocation("properties")] attributes: Annotated[Attributes, SchemaLocation("attributes"), DataLocation("attributes")] + async def to_dataframe(self) -> pd.DataFrame: + """Return collars with a categorical ``hole_id`` aligned to hole-index codes.""" + parts = [ + await self.hole_id.to_dataframe(), + await self.coordinates.to_dataframe(), + await self.distances.to_dataframe(), + ] + if len(self.attributes): + parts.append(await self.attributes.to_dataframe()) + return pd.concat(parts, axis=1) + + async def path_to_dataframe(self) -> pd.DataFrame: + """Return the desurvey path and its attributes.""" + return await self.path.to_dataframe() + class _Distances(DataTable): table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_1 @@ -267,10 +334,13 @@ class DistanceTableDistances(DataTableAndAttributes): @classmethod async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: result = await super()._data_to_schema(data, context) + unit = data.attrs.get("unit") attr_desc: AttributeDescription = data.attrs.get("attribute_descriptions", {}).get("distance") - if attr_desc is not None and attr_desc.unit is not None: + if unit is None and attr_desc is not None: + unit = attr_desc.unit + if unit is not None: # "unit" can be missing, but it must not be `None` - result["unit"] = attr_desc.unit + result["unit"] = unit return result @@ -283,9 +353,98 @@ class DistanceTable(SchemaModel): class DownholeDistanceTable(DistanceTable): holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] + async def to_dataframe(self) -> pd.DataFrame: + return await self.distance.to_dataframe() -class DownholeCollectionTables(SchemaList[DownholeDistanceTable]): - pass + async def to_dataframe_by_hole(self) -> dict[str, pd.DataFrame]: + return await _table_by_hole(self, await self.to_dataframe()) + + +class _Intervals(DataTable): + table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_2 + data_columns: ClassVar[list[str]] = ["from", "to"] + + +class IntervalTableFromTo(DataTableAndAttributes): + _table: Annotated[_Intervals, SchemaLocation("intervals.start_and_end"), 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) + unit = data.attrs.get("unit") + description = data.attrs.get("attribute_descriptions", {}).get("from") + if unit is None and description is not None: + unit = description.unit + if unit is not None: + result["unit"] = unit + return result + + +class DownholeIntervalTable(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("interval_table")] + holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] + + async def to_dataframe(self) -> pd.DataFrame: + return await self.from_to.to_dataframe() + + async def to_dataframe_by_hole(self) -> dict[str, pd.DataFrame]: + return await _table_by_hole(self, await self.to_dataframe()) + + +class DownholeCollectionTables(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 + table = ( + collection.interval_table if isinstance(collection, IntervalCollection) else collection.distance_table + ) + table = table.copy() + table.attrs = dict(table.attrs) + if collection.unit is not None: + table.attrs["unit"] = collection.unit + field = "interval_table" if isinstance(collection, IntervalCollection) else "distance_table" + result.append(await model._data_to_schema(replace(collection, **{field: table}), context)) + return result + + def get(self, name: str) -> DownholeDistanceTable | DownholeIntervalTable | None: + return next((collection for collection in self if collection.name == name), None) + + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and self.get(name) is not None + + def names(self) -> list[str]: + return [collection.name for collection in self] + + async def add(self, collection: DownholeCollectionEntry, *, replace: bool = False) -> None: + existing = self.names() + if collection.name in existing and not replace: + raise ValueError(f"Collection '{collection.name}' already exists") + location_holes = await self._context.root_model.location.holes.to_dataframe() + valid_indices = set(location_holes["hole_index"].astype(int)) + if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): + raise ObjectValidationError("Collection hole_index is not present in the location holes table") + schema = await self._data_to_schema([collection], self._obj) + if collection.name in existing: + self._document[existing.index(collection.name)] = schema[0] + else: + self._document.append(schema[0]) + + def remove(self, *names: str) -> int: + 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 +460,30 @@ 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, **kwargs: Any) -> None: + """Prefetch data referenced by named collections and optionally location data.""" + from evo.objects.typed._prefetch import collect_data_ids + + 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), **kwargs) + + +async def _table_by_hole( + table: DownholeDistanceTable | DownholeIntervalTable, data: pd.DataFrame +) -> dict[str, pd.DataFrame]: + root = table._context.root_model + collar_ids = await root.location.hole_id.to_dataframe() + categories = collar_ids.iloc[:, 0] + result: dict[str, pd.DataFrame] = {} + for chunk in (await table.holes.to_dataframe()).itertuples(index=False): + hole_id = str(categories.cat.categories[int(chunk.hole_index)]) + result[hole_id] = data.iloc[int(chunk.offset) : int(chunk.offset) + int(chunk.count)].reset_index(drop=True) + return result 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/downhole.py b/packages/evo-objects/src/evo/objects/utils/downhole.py new file mode 100644 index 00000000..0dedea8f --- /dev/null +++ b/packages/evo-objects/src/evo/objects/utils/downhole.py @@ -0,0 +1,67 @@ +# 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 + +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, *, dtype: pd.CategoricalDtype) -> pd.DataFrame: + """Run-length encode contiguous hole IDs using their categorical codes. + + A zero-count entry is emitted for categories absent from ``hole_ids``. An ID + outside ``dtype`` or a repeated non-contiguous run is rejected. + """ + unknown_mask = hole_ids.notna() & ~hole_ids.isin(dtype.categories) + if unknown_mask.any(): + unknown = hole_ids[unknown_mask].unique().tolist() + raise ValueError(f"hole_ids contains values absent from dtype: {unknown}") + categorical = hole_ids.astype(dtype) + codes = categorical.cat.codes.to_numpy(dtype=np.int32) + chunks: dict[int, tuple[int, int]] = {} + start = 0 + while start < len(codes): + code = int(codes[start]) + if code < 0: + raise ValueError("hole_ids cannot contain missing values") + end = start + 1 + while end < len(codes) and codes[end] == code: + end += 1 + if code in chunks: + raise ValueError("Rows for each hole_id must be contiguous") + chunks[code] = (start, end - start) + start = end + return pd.DataFrame( + { + "hole_index": np.arange(len(dtype.categories), dtype=np.int32), + "offset": np.array([chunks.get(code, (0, 0))[0] for code in range(len(dtype.categories))], dtype=np.uint64), + "count": np.array([chunks.get(code, (0, 0))[1] for code in range(len(dtype.categories))], 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/test_downhole_utils.py b/packages/evo-objects/tests/test_downhole_utils.py new file mode 100644 index 00000000..6449e0be --- /dev/null +++ b/packages/evo-objects/tests/test_downhole_utils.py @@ -0,0 +1,33 @@ +# 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 test_chunks_use_dtype_codes_and_round_trip(self): + dtype = pd.CategoricalDtype(categories=["Z", "A", "M"]) + values = pd.Series(["M", "M", "Z"]) + chunks = hole_chunks_from_ids(values, dtype=dtype) + self.assertListEqual(chunks["hole_index"].tolist(), [0, 1, 2]) + self.assertListEqual(chunks["count"].tolist(), [1, 0, 2]) + self.assertListEqual(expand_hole_index(chunks, len(values)).tolist(), [2, 2, 0]) + + def test_non_contiguous_and_unknown_ids_raise(self): + dtype = pd.CategoricalDtype(categories=["A", "B"]) + with self.assertRaises(ValueError): + hole_chunks_from_ids(pd.Series(["A", "B", "A"]), dtype=dtype) + with self.assertRaises(ValueError): + hole_chunks_from_ids(pd.Series(["C"]), dtype=dtype) diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 4c6f1fd8..492f8269 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -32,6 +32,7 @@ DistanceCollection, DownholeCollection, DownholeCollectionData, + IntervalCollection, ) from evo.objects.typed.exceptions import ObjectValidationError @@ -208,6 +209,36 @@ 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]}), + interval_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_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 +291,17 @@ 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_hole_index_not_property_position(self): + data = _make_example_data() + expected = data.compute_bounding_box() + properties = data.properties.iloc[[1, 0]].copy() + properties.index = [10, 20] + 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( From 9808cd96a57fd6727665240087df0ac557451cfa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Fri, 31 Jul 2026 09:25:34 +0100 Subject: [PATCH 02/30] feat(objects): add typed object prefetch support --- .../src/evo/objects/typed/_prefetch.py | 72 +++++++++++++++++++ .../evo-objects/src/evo/objects/typed/base.py | 16 ++++- 2 files changed, 87 insertions(+), 1 deletion(-) create mode 100644 packages/evo-objects/src/evo/objects/typed/_prefetch.py 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..bb8b0c83 --- /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 = 100, + 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/base.py b/packages/evo-objects/src/evo/objects/typed/base.py index c194cbd2..55675fb3 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 = 100, + 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, From aec5e9e5aeddae0c1857b7f2b1b635b7560b0cb0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Fri, 31 Jul 2026 09:25:48 +0100 Subject: [PATCH 03/30] feat(objects): add typed attribute metadata --- .../src/evo/objects/typed/attributes.py | 43 ++++++++++++++----- .../tests/typed/test_attributes.py | 19 +++++++- 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 8ea42595..d5b20abc 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -41,6 +41,7 @@ __all__ = [ "Attribute", + "AttributeDescription", "Attributes", "BlockModelAttribute", "BlockModelAttributes", @@ -87,18 +88,23 @@ def _infer_attribute_type_from_series(series: pd.Series) -> str: @dataclass class AttributeDescription: - discipline: str = "" - type: str = "" - unit: str | None = None + discipline: str | None = None + type: str | None = None + unit: Any | None = None scale: str | None = None extensions: dict[str, typing.Any] | None = None tags: dict[str, str] | None = None - def to_schema(self): - result = { - "discipline": self.discipline, - "type": self.type, - } + def __post_init__(self) -> None: + if self.unit is not None and not isinstance(self.unit, str): + self.unit = str(self.unit.value) + + def to_schema(self) -> dict[str, Any]: + result: dict[str, Any] = {} + if self.discipline: + result["discipline"] = self.discipline + if self.type: + result["type"] = self.type if self.unit: result["unit"] = self.unit if self.scale: @@ -117,6 +123,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 +139,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. @@ -320,8 +333,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 +350,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. diff --git a/packages/evo-objects/tests/typed/test_attributes.py b/packages/evo-objects/tests/typed/test_attributes.py index c8928cec..ad2004ff 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -16,7 +16,12 @@ 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 +70,15 @@ 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_empty_description_is_omitted(self): + self.assertEqual(AttributeDescription().to_schema(), {}) + + 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"}) From 1053f66c93323ca31916d7a7b689cc0117814995 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Fri, 31 Jul 2026 09:26:00 +0100 Subject: [PATCH 04/30] feat(objects): expose data client context (bump evo-objects to 0.6.2) --- packages/evo-objects/pyproject.toml | 2 +- packages/evo-objects/src/evo/objects/utils/data.py | 11 ++++++++++- uv.lock | 2 +- 3 files changed, 12 insertions(+), 3 deletions(-) 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/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/uv.lock b/uv.lock index fcb05bef..d33e1ff0 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"] }, From 798b0fc1fd75480574cf7672fdae0dcd5e2258f2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Wed, 5 Aug 2026 10:39:15 +0100 Subject: [PATCH 05/30] feat(objects): refine typed model and downhole collection; add prefetch test --- packages/evo-objects/README.md | 34 + .../src/evo/objects/typed/_model.py | 8 +- .../src/evo/objects/typed/attributes.py | 2 +- .../evo/objects/typed/downhole_collection.py | 1009 ++++++++-------- .../evo-objects/tests/test_downhole_utils.py | 13 + .../tests/typed/test_downhole_collection.py | 1010 +++++++++-------- .../evo-objects/tests/typed/test_model.py | 14 +- .../evo-objects/tests/typed/test_prefetch.py | 96 ++ 8 files changed, 1240 insertions(+), 946 deletions(-) create mode 100644 packages/evo-objects/tests/typed/test_prefetch.py diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index fe3a7fba..d3721e3b 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -69,6 +69,40 @@ 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 chunk tables always use a zero-based `hole_index`: it is the code in the +shared categorical `properties["hole_id"]` dtype, not the row position of a collar. + +```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 + +hole_dtype = pd.CategoricalDtype(categories=["DH-01"]) +intervals = pd.DataFrame({"from": [0.0], "to": [1.5], "lithology": ["sandstone"]}) +collections = [ + IntervalCollection( + name="geology", + holes=hole_chunks_from_ids(pd.Series(["DH-01"]), dtype=hole_dtype), + interval_table=intervals, + unit="m", # Explicit collection units override DataFrame metadata. + ) +] + +# 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 round-trip through `DataFrame.attrs["attribute_descriptions"]`; an explicit collection `unit` +takes precedence over unit metadata on the distance or `from` column. + ## 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/src/evo/objects/typed/_model.py b/packages/evo-objects/src/evo/objects/typed/_model.py index e8a1f0b8..85d0fb61 100644 --- a/packages/evo-objects/src/evo/objects/typed/_model.py +++ b/packages/evo-objects/src/evo/objects/typed/_model.py @@ -450,7 +450,13 @@ 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: diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index d5b20abc..6e1defbb 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -254,7 +254,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 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 2ce9c893..41447558 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -1,489 +1,520 @@ -# 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 - -from dataclasses import dataclass, replace -from typing import Annotated, Any, ClassVar, TypeAlias - -import numpy as np -import pandas as pd -from numpy._typing import NDArray - -from evo.common.interfaces import IContext -from evo.objects import SchemaVersion -from evo.objects.typed._data import DataTable, DataTableAndAttributes -from evo.objects.typed._downhole import HoleIdCategory -from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel -from evo.objects.typed.attributes import ( - AttributeDescription, - Attributes, -) -from evo.objects.typed.exceptions import ObjectValidationError -from evo.objects.typed.spatial import BaseSpatialObject, BaseSpatialObjectData -from evo.objects.typed.types import BoundingBox -from evo.objects.utils.table_formats import ( - DOWNHOLE_COLLECTION_LOCATION_HOLES, - FLOAT_ARRAY_1, - FLOAT_ARRAY_2, - FLOAT_ARRAY_3, - KnownTableFormat, -) - -__all__ = [ - "DistanceCollection", - "DownholeCollection", - "DownholeCollectionData", - "IntervalCollection", -] - -_X = "x" -_Y = "y" -_Z = "z" -_COORDINATE_COLUMNS = [_X, _Y, _Z] - - -HolePath: TypeAlias = pd.DataFrame # [ distance | dip | azimuth | ] -HoleChunks: TypeAlias = pd.DataFrame # [ hole_id | 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 - collection_type: str = "distance" - unit: str | None = None - - -@dataclass -class IntervalCollection: - name: str - holes: HoleChunks - interval_table: Intervals - collection_type: str = "interval" - unit: str | None = None - - -DownholeCollectionEntry: TypeAlias = DistanceCollection | IntervalCollection - - -@dataclass(kw_only=True, frozen=True) -class DownholeCollectionData(BaseSpatialObjectData): - """Data class for creating a new 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`. - :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 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". - :param coordinate_reference_system: Optional EPSG code or WKT string for the coordinate reference system. - :param description: Optional description of the object. - :param tags: Optional dictionary of tags for the object. - :param extensions: Optional dictionary of extensions for the object. - """ - - path: HolePath - holes: HoleChunks - properties: HoleProperties - attributes: HoleAttributes | None - collections: list[DownholeCollectionEntry] - distance_unit: str | None = None - desurvey: str | None = None - - @property - def hole_id_dtype(self) -> pd.CategoricalDtype: - """The categorical dtype whose codes are used by every ``hole_index``.""" - hole_ids = self.properties["hole_id"] - if isinstance(hole_ids.dtype, pd.CategoricalDtype): - return hole_ids.dtype - return pd.CategoricalDtype(categories=sorted(hole_ids.dropna().unique())) - - 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") - - assert self.attributes is None or len(self.holes) == len(self.attributes) - - self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) - for collection in self.collections: - table = ( - collection.distance_table if isinstance(collection, DistanceCollection) else collection.interval_table - ) - self._validate_hole_chunks(collection.holes, len(table), require_coverage=False) - - def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require_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) - valid = set(range(len(self.hole_id_dtype.categories))) - if not set(indices).issubset(valid): - raise ObjectValidationError("hole_index must be a code in properties['hole_id'] categorical dtype") - if indices.duplicated().any(): - raise ObjectValidationError("Each hole_index may occur only once in a holes table") - if require_coverage and set(indices) != valid: - raise ObjectValidationError("Location holes must cover every hole_id categorical code 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") - - def compute_bounding_box(self) -> BoundingBox: - bboxes = [] - - collars = self.properties.copy() - collars["_hole_index"] = collars["hole_id"].astype(self.hole_id_dtype).cat.codes - collars_by_index = collars.set_index("_hole_index") - for chunk in self.holes.itertuples(index=False): - offset = int(chunk.offset) - count = int(chunk.count) - collar = tuple(collars_by_index.loc[int(chunk.hole_index), _COORDINATE_COLUMNS]) - path_table = self.path[offset : offset + count] - bboxes.append(self._compute_hole_bounding_box(path_table, collar)) - - return BoundingBox.combine(bboxes) - - @staticmethod - def _compute_bounding_box_np( - depths: NDArray[np.float64], - dips: NDArray[np.float64], - azimuths: NDArray[np.float64], - offset: tuple[float, float, float] = (0.0, 0.0, 0.0), - ) -> BoundingBox: - if not np.all(depths[:-1] <= depths[1:]): - raise ObjectValidationError("depths must be sorted") - - if len(depths) != len(dips) or len(depths) != len(azimuths): - raise ObjectValidationError("depths, dips, and azimuths must have same length") - - # Process NaNs - # `depths`, `dips`, and `azimuths` could be read-only views, so take copies instead of mutating - depths = depths[~np.isnan(depths)] - dips = np.where(np.isnan(dips), 90.0, dips) - azimuths = np.where(np.isnan(azimuths), 0.0, azimuths) - - dips_rad = np.deg2rad(dips) - azimuths_rad = np.deg2rad(azimuths) - - # Prepend 0 so `step` has the same shape as `dips` and `azimuths`, and so the first depth gets treated as the - # first step. The depth column might already start with 0, in which case the first step will be length 0, which - # is a no-op as far as the following calculation is concerned. - step = np.diff(depths, prepend=0.0) - - dz_down = step * np.sin(dips_rad) - horiz = step * np.cos(dips_rad) - - # Horizontal into N/E (0° = North, 90° = East) - dN = horiz * np.cos(azimuths_rad) - dE = horiz * np.sin(azimuths_rad) - - # Convert to XYZ increments (Z up) - dX = dE - dY = dN - dZ = -dz_down - - x = np.cumsum(dX) - y = np.cumsum(dY) - z = np.cumsum(dZ) - - def ensure_zero(a, b): - return min(a, 0), max(b, 0) - - x0, x1 = ensure_zero(x.min(), x.max()) - y0, y1 = ensure_zero(y.min(), y.max()) - z0, z1 = ensure_zero(z.min(), z.max()) - - return BoundingBox( - min_x=x0 + offset[0], - max_x=x1 + offset[0], - min_y=y0 + offset[1], - max_y=y1 + offset[1], - min_z=z0 + offset[2], - max_z=z1 + offset[2], - ) - - @staticmethod - def _compute_hole_bounding_box( - depths_dips_azimuths_table: pd.DataFrame, - collar: tuple[float, float, float], - ) -> BoundingBox: - """ - Compute 3D bounding box for a deviated hole given collar XYZ and - depth / dip / azimuth data. - - Conventions - ----------- - - depths: measured depth along the hole (m), positive downward. - - dips: inclination FROM VERTICAL (degrees). - 90° = vertical down, 0° = horizontal. - - azimuths: degrees clockwise from North. - - Coordinates: X = Easting, Y = Northing, Z = elevation (up). - """ - df = depths_dips_azimuths_table.dropna(subset=["distance"]) - box = DownholeCollectionData._compute_bounding_box_np( - df["distance"].astype(float).to_numpy(), - df["dip"].astype(float).to_numpy(), - df["azimuth"].astype(float).to_numpy(), - offset=collar, - ) - - return box - - -class HoleChunksTable(DataTable): - table_format: ClassVar[KnownTableFormat] = DOWNHOLE_COLLECTION_LOCATION_HOLES - data_columns: ClassVar[list[str]] = ["hole_index", "offset", "count"] - - -class PathTable(DataTable): - table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 - data_columns: ClassVar[list[str]] = ["distance", "azimuth", "dip"] - - -class DownholePath(DataTableAndAttributes): - _table: Annotated[PathTable, SchemaLocation(""), DataLocation("")] - - -class DistancesTable(DataTable): - table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 - data_columns: ClassVar[list[str]] = ["final", "target", "current"] - - @classmethod - def _extract_distances(cls, data: HoleAttributes) -> pd.DataFrame: - return data[["final", "target", "current"]].astype(np.float64) - - @classmethod - async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: - distances_df = cls._extract_distances(data) - return await super()._data_to_schema(distances_df, context) - - -class CollarCoordinates(DataTable): - table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 - data_columns: ClassVar[list[str]] = _COORDINATE_COLUMNS - - @classmethod - def _extract_coordinates(cls, data: HoleAttributes): - return data[["x", "y", "z"]].astype(np.float64) - - @classmethod - async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: - distances_df = cls._extract_coordinates(data) - return await super()._data_to_schema(distances_df, context) - - -class DownholeLocation(SchemaModel): - hole_id: Annotated[HoleIdCategory, 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) -> pd.DataFrame: - """Return collars with a categorical ``hole_id`` aligned to hole-index codes.""" - parts = [ - await self.hole_id.to_dataframe(), - await self.coordinates.to_dataframe(), - await self.distances.to_dataframe(), - ] - if len(self.attributes): - parts.append(await self.attributes.to_dataframe()) - return pd.concat(parts, axis=1) - - async def path_to_dataframe(self) -> pd.DataFrame: - """Return the desurvey path and its attributes.""" - return await self.path.to_dataframe() - - -class _Distances(DataTable): - table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_1 - data_columns: ClassVar[list[str]] = ["distance"] - - -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) - unit = data.attrs.get("unit") - attr_desc: AttributeDescription = data.attrs.get("attribute_descriptions", {}).get("distance") - if unit is None and attr_desc is not None: - unit = attr_desc.unit - if unit is not None: - # "unit" can be missing, but it must not be `None` - result["unit"] = 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")] - - -class DownholeDistanceTable(DistanceTable): - holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] - - async def to_dataframe(self) -> pd.DataFrame: - return await self.distance.to_dataframe() - - async def to_dataframe_by_hole(self) -> dict[str, pd.DataFrame]: - return await _table_by_hole(self, await self.to_dataframe()) - - -class _Intervals(DataTable): - table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_2 - data_columns: ClassVar[list[str]] = ["from", "to"] - - -class IntervalTableFromTo(DataTableAndAttributes): - _table: Annotated[_Intervals, SchemaLocation("intervals.start_and_end"), 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) - unit = data.attrs.get("unit") - description = data.attrs.get("attribute_descriptions", {}).get("from") - if unit is None and description is not None: - unit = description.unit - if unit is not None: - result["unit"] = unit - return result - - -class DownholeIntervalTable(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("interval_table")] - holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] - - async def to_dataframe(self) -> pd.DataFrame: - return await self.from_to.to_dataframe() - - async def to_dataframe_by_hole(self) -> dict[str, pd.DataFrame]: - return await _table_by_hole(self, await self.to_dataframe()) - - -class DownholeCollectionTables(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 - table = ( - collection.interval_table if isinstance(collection, IntervalCollection) else collection.distance_table - ) - table = table.copy() - table.attrs = dict(table.attrs) - if collection.unit is not None: - table.attrs["unit"] = collection.unit - field = "interval_table" if isinstance(collection, IntervalCollection) else "distance_table" - result.append(await model._data_to_schema(replace(collection, **{field: table}), context)) - return result - - def get(self, name: str) -> DownholeDistanceTable | DownholeIntervalTable | None: - return next((collection for collection in self if collection.name == name), None) - - def __contains__(self, name: object) -> bool: - return isinstance(name, str) and self.get(name) is not None - - def names(self) -> list[str]: - return [collection.name for collection in self] - - async def add(self, collection: DownholeCollectionEntry, *, replace: bool = False) -> None: - existing = self.names() - if collection.name in existing and not replace: - raise ValueError(f"Collection '{collection.name}' already exists") - location_holes = await self._context.root_model.location.holes.to_dataframe() - valid_indices = set(location_holes["hole_index"].astype(int)) - if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): - raise ObjectValidationError("Collection hole_index is not present in the location holes table") - schema = await self._data_to_schema([collection], self._obj) - if collection.name in existing: - self._document[existing.index(collection.name)] = schema[0] - else: - self._document.append(schema[0]) - - def remove(self, *names: str) -> int: - 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): - """A GeoscienceObject representing a collection of downholes.""" - - _data_class = DownholeCollectionData - sub_classification = "downhole-collection" - creation_schema_version = SchemaVersion(major=1, minor=3, patch=1) - - location: Annotated[DownholeLocation, SchemaLocation("location"), DataLocation("")] - collections: Annotated[DownholeCollectionTables, SchemaLocation("collections"), DataLocation("collections")] - distance_unit: Annotated[str | None, SchemaLocation("distance_unit")] - desurvey: Annotated[str | None, SchemaLocation("desurvey")] - - type: ClassVar[Annotated[str, SchemaLocation("type")]] = "downhole" - - async def prefetch_collections(self, *names: str, include_location: bool = True, **kwargs: Any) -> None: - """Prefetch data referenced by named collections and optionally location data.""" - from evo.objects.typed._prefetch import collect_data_ids - - 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), **kwargs) - - -async def _table_by_hole( - table: DownholeDistanceTable | DownholeIntervalTable, data: pd.DataFrame -) -> dict[str, pd.DataFrame]: - root = table._context.root_model - collar_ids = await root.location.hole_id.to_dataframe() - categories = collar_ids.iloc[:, 0] - result: dict[str, pd.DataFrame] = {} - for chunk in (await table.holes.to_dataframe()).itertuples(index=False): - hole_id = str(categories.cat.categories[int(chunk.hole_index)]) - result[hole_id] = data.iloc[int(chunk.offset) : int(chunk.offset) + int(chunk.count)].reset_index(drop=True) - return result +# 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 + +from dataclasses import dataclass, replace +from typing import Annotated, Any, ClassVar, TypeAlias + +import numpy as np +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._data import DataTable, DataTableAndAttributes +from evo.objects.typed._downhole import HoleIdCategory +from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel +from evo.objects.typed.attributes import ( + AttributeDescription, + Attributes, +) +from evo.objects.typed.exceptions import ObjectValidationError +from evo.objects.typed.spatial import BaseSpatialObject, BaseSpatialObjectData +from evo.objects.typed.types import BoundingBox +from evo.objects.utils.table_formats import ( + DOWNHOLE_COLLECTION_LOCATION_HOLES, + FLOAT_ARRAY_1, + FLOAT_ARRAY_2, + FLOAT_ARRAY_3, + KnownTableFormat, +) + +__all__ = [ + "DistanceCollection", + "DownholeCollection", + "DownholeCollectionData", + "IntervalCollection", +] + +_X = "x" +_Y = "y" +_Z = "z" +_COORDINATE_COLUMNS = [_X, _Y, _Z] + + +HolePath: TypeAlias = pd.DataFrame # [ distance | dip | azimuth | ] +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_descriptions': {: }, ...} +Depths: TypeAlias = pd.DataFrame # [ distance | ] +Intervals: TypeAlias = pd.DataFrame # [ from | to | ] + + +@dataclass +class DistanceCollection: + name: str + holes: HoleChunks + distance_table: Depths + collection_type: str = "distance" + unit: str | None = None + + +@dataclass +class IntervalCollection: + name: str + holes: HoleChunks + interval_table: Intervals + collection_type: str = "interval" + unit: str | None = None + + +DownholeCollectionEntry: TypeAlias = DistanceCollection | IntervalCollection + + +@dataclass(kw_only=True, frozen=True) +class DownholeCollectionData(BaseSpatialObjectData): + """Data class for creating a new DownholeCollection + + :param name: The name of the object. + :param holes: A DataFrame describing which parts of `path` belong to which holes. + Columns: hole_index, offset, count. ``hole_index`` is a code in ``properties.hole_id``. + :param properties: DataFrame for the properties of the holes, joined to ``holes`` by the categorical hole-id code. + Mandatory columns: hole_id, final, target, current, x, y, z + :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: 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". + :param coordinate_reference_system: Optional EPSG code or WKT string for the coordinate reference system. + :param description: Optional description of the object. + :param tags: Optional dictionary of tags for the object. + :param extensions: Optional dictionary of extensions for the object. + """ + + path: HolePath + holes: HoleChunks + properties: HoleProperties + attributes: HoleAttributes | None + collections: list[DownholeCollectionEntry] + distance_unit: str | None = None + desurvey: str | None = None + + @property + def hole_id_dtype(self) -> pd.CategoricalDtype: + """The categorical dtype whose codes are used by every ``hole_index``.""" + hole_ids = self.properties["hole_id"] + if isinstance(hole_ids.dtype, pd.CategoricalDtype): + return hole_ids.dtype + return pd.CategoricalDtype(categories=sorted(hole_ids.dropna().unique())) + + 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") + + assert self.attributes is None or len(self.holes) == len(self.attributes) + + self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) + for collection in self.collections: + table = ( + collection.distance_table if isinstance(collection, DistanceCollection) else collection.interval_table + ) + self._validate_hole_chunks(collection.holes, len(table), require_coverage=False) + + def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require_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) + valid = set(range(len(self.hole_id_dtype.categories))) + if not set(indices).issubset(valid): + raise ObjectValidationError("hole_index must be a code in properties['hole_id'] categorical dtype") + if indices.duplicated().any(): + raise ObjectValidationError("Each hole_index may occur only once in a holes table") + if require_coverage and set(indices) != valid: + raise ObjectValidationError("Location holes must cover every hole_id categorical code 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") + 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 compute_bounding_box(self) -> BoundingBox: + bboxes = [] + + collars = self.properties.copy() + collars["_hole_index"] = collars["hole_id"].astype(self.hole_id_dtype).cat.codes + collars_by_index = collars.set_index("_hole_index") + for chunk in self.holes.itertuples(index=False): + offset = int(chunk.offset) + count = int(chunk.count) + collar = tuple(collars_by_index.loc[int(chunk.hole_index), _COORDINATE_COLUMNS]) + path_table = self.path[offset : offset + count] + bboxes.append(self._compute_hole_bounding_box(path_table, collar)) + + return BoundingBox.combine(bboxes) + + @staticmethod + def _compute_bounding_box_np( + depths: NDArray[np.float64], + dips: NDArray[np.float64], + azimuths: NDArray[np.float64], + offset: tuple[float, float, float] = (0.0, 0.0, 0.0), + ) -> BoundingBox: + if not np.all(depths[:-1] <= depths[1:]): + raise ObjectValidationError("depths must be sorted") + + if len(depths) != len(dips) or len(depths) != len(azimuths): + raise ObjectValidationError("depths, dips, and azimuths must have same length") + + # Process NaNs + # `depths`, `dips`, and `azimuths` could be read-only views, so take copies instead of mutating + depths = depths[~np.isnan(depths)] + dips = np.where(np.isnan(dips), 90.0, dips) + azimuths = np.where(np.isnan(azimuths), 0.0, azimuths) + + dips_rad = np.deg2rad(dips) + azimuths_rad = np.deg2rad(azimuths) + + # Prepend 0 so `step` has the same shape as `dips` and `azimuths`, and so the first depth gets treated as the + # first step. The depth column might already start with 0, in which case the first step will be length 0, which + # is a no-op as far as the following calculation is concerned. + step = np.diff(depths, prepend=0.0) + + dz_down = step * np.sin(dips_rad) + horiz = step * np.cos(dips_rad) + + # Horizontal into N/E (0° = North, 90° = East) + dN = horiz * np.cos(azimuths_rad) + dE = horiz * np.sin(azimuths_rad) + + # Convert to XYZ increments (Z up) + dX = dE + dY = dN + dZ = -dz_down + + x = np.cumsum(dX) + y = np.cumsum(dY) + z = np.cumsum(dZ) + + def ensure_zero(a, b): + return min(a, 0), max(b, 0) + + x0, x1 = ensure_zero(x.min(), x.max()) + y0, y1 = ensure_zero(y.min(), y.max()) + z0, z1 = ensure_zero(z.min(), z.max()) + + return BoundingBox( + min_x=x0 + offset[0], + max_x=x1 + offset[0], + min_y=y0 + offset[1], + max_y=y1 + offset[1], + min_z=z0 + offset[2], + max_z=z1 + offset[2], + ) + + @staticmethod + def _compute_hole_bounding_box( + depths_dips_azimuths_table: pd.DataFrame, + collar: tuple[float, float, float], + ) -> BoundingBox: + """ + Compute 3D bounding box for a deviated hole given collar XYZ and + depth / dip / azimuth data. + + Conventions + ----------- + - depths: measured depth along the hole (m), positive downward. + - dips: inclination FROM VERTICAL (degrees). + 90° = vertical down, 0° = horizontal. + - azimuths: degrees clockwise from North. + - Coordinates: X = Easting, Y = Northing, Z = elevation (up). + """ + df = depths_dips_azimuths_table.dropna(subset=["distance"]) + box = DownholeCollectionData._compute_bounding_box_np( + df["distance"].astype(float).to_numpy(), + df["dip"].astype(float).to_numpy(), + df["azimuth"].astype(float).to_numpy(), + offset=collar, + ) + + return box + + +class HoleChunksTable(DataTable): + table_format: ClassVar[KnownTableFormat] = DOWNHOLE_COLLECTION_LOCATION_HOLES + data_columns: ClassVar[list[str]] = ["hole_index", "offset", "count"] + + +class PathTable(DataTable): + table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 + data_columns: ClassVar[list[str]] = ["distance", "azimuth", "dip"] + + +class DownholePath(DataTableAndAttributes): + _table: Annotated[PathTable, SchemaLocation(""), DataLocation("")] + + +class DistancesTable(DataTable): + table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 + data_columns: ClassVar[list[str]] = ["final", "target", "current"] + + @classmethod + def _extract_distances(cls, data: HoleAttributes) -> pd.DataFrame: + return data[["final", "target", "current"]].astype(np.float64) + + @classmethod + async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: + distances_df = cls._extract_distances(data) + return await super()._data_to_schema(distances_df, context) + + +class CollarCoordinates(DataTable): + table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 + data_columns: ClassVar[list[str]] = _COORDINATE_COLUMNS + + @classmethod + def _extract_coordinates(cls, data: HoleAttributes): + return data[["x", "y", "z"]].astype(np.float64) + + @classmethod + async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: + distances_df = cls._extract_coordinates(data) + return await super()._data_to_schema(distances_df, context) + + +class DownholeLocation(SchemaModel): + hole_id: Annotated[HoleIdCategory, 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, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return collars with a categorical ``hole_id`` aligned to hole-index codes.""" + 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(fb=fb)) + return pd.concat(parts, axis=1) + + async def path_to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return the desurvey path and its attributes.""" + return await self.path.to_dataframe(fb=fb) + + +class _Distances(DataTable): + table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_1 + data_columns: ClassVar[list[str]] = ["distance"] + + +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) + unit = data.attrs.get("unit") + attr_desc: AttributeDescription = data.attrs.get("attribute_descriptions", {}).get("distance") + if unit is None and attr_desc is not None: + unit = attr_desc.unit + if unit is not None: + # "unit" can be missing, but it must not be `None` + result["unit"] = 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")] + + +class DownholeDistanceTable(DistanceTable): + holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] + + async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: + return await self.distance.to_dataframe(fb=fb) + + async def to_dataframe_by_hole(self, fb: IFeedback = NoFeedback) -> dict[str, pd.DataFrame]: + return await _table_by_hole(self, await self.to_dataframe(fb=fb), fb=fb) + + +class _Intervals(DataTable): + table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_2 + data_columns: ClassVar[list[str]] = ["from", "to"] + + +class IntervalTableFromTo(DataTableAndAttributes): + _table: Annotated[_Intervals, SchemaLocation("intervals.start_and_end"), 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) + unit = data.attrs.get("unit") + description = data.attrs.get("attribute_descriptions", {}).get("from") + if unit is None and description is not None: + unit = description.unit + if unit is not None: + result["unit"] = unit + return result + + +class DownholeIntervalTable(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("interval_table")] + holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] + + async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: + return await self.from_to.to_dataframe(fb=fb) + + async def to_dataframe_by_hole(self, fb: IFeedback = NoFeedback) -> dict[str, pd.DataFrame]: + return await _table_by_hole(self, await self.to_dataframe(fb=fb), fb=fb) + + +class DownholeCollectionTables(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 + table = ( + collection.interval_table if isinstance(collection, IntervalCollection) else collection.distance_table + ) + table = table.copy() + table.attrs = dict(table.attrs) + if collection.unit is not None: + table.attrs["unit"] = collection.unit + field = "interval_table" if isinstance(collection, IntervalCollection) else "distance_table" + result.append(await model._data_to_schema(replace(collection, **{field: table}), context)) + return result + + def get(self, name: str) -> DownholeDistanceTable | DownholeIntervalTable | None: + return next((collection for collection in self if collection.name == name), None) + + def __contains__(self, name: object) -> bool: + return isinstance(name, str) and self.get(name) is not None + + def names(self) -> list[str]: + return [collection.name for collection in self] + + async def add(self, collection: DownholeCollectionEntry, *, replace: bool = False) -> None: + existing = self.names() + if collection.name in existing and not replace: + raise ValueError(f"Collection '{collection.name}' already exists") + location_holes = await self._context.root_model.location.holes.to_dataframe() + valid_indices = set(location_holes["hole_index"].astype(int)) + if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): + raise ObjectValidationError("Collection hole_index is not present in the location holes table") + table = collection.distance_table if isinstance(collection, DistanceCollection) else collection.interval_table + _validate_chunk_ranges(collection.holes, len(table)) + schema = await self._data_to_schema([collection], self._obj) + if collection.name in existing: + self._document[existing.index(collection.name)] = schema[0] + else: + self._document.append(schema[0]) + + def remove(self, *names: str) -> int: + 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): + """A GeoscienceObject representing a collection of downholes.""" + + _data_class = DownholeCollectionData + sub_classification = "downhole-collection" + creation_schema_version = SchemaVersion(major=1, minor=3, patch=1) + + location: Annotated[DownholeLocation, SchemaLocation("location"), DataLocation("")] + collections: Annotated[DownholeCollectionTables, SchemaLocation("collections"), DataLocation("collections")] + distance_unit: Annotated[str | None, SchemaLocation("distance_unit")] + desurvey: Annotated[str | None, SchemaLocation("desurvey")] + + type: ClassVar[Annotated[str, SchemaLocation("type")]] = "downhole" + + async def prefetch_collections(self, *names: str, include_location: bool = True, **kwargs: Any) -> None: + """Prefetch data referenced by named collections and optionally location data.""" + from evo.objects.typed._prefetch import collect_data_ids + + 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), **kwargs) + + +def _validate_chunk_ranges(holes: HoleChunks, table_length: int) -> None: + required = {"hole_index", "offset", "count"} + if missing := required - set(holes.columns): + raise ObjectValidationError(f"Hole chunks are missing columns: {sorted(missing)}") + if holes["hole_index"].astype(int).duplicated().any(): + raise ObjectValidationError("Each hole_index may occur only once in a holes table") + 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") + expected_offset = 0 + for offset, count in sorted(zip(offsets[counts > 0], counts[counts > 0], strict=True)): + 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") + + +async def _table_by_hole( + table: DownholeDistanceTable | DownholeIntervalTable, data: pd.DataFrame, *, fb: IFeedback +) -> dict[str, pd.DataFrame]: + root = table._context.root_model + collar_ids = await root.location.hole_id.to_dataframe(fb=fb) + categories = collar_ids.iloc[:, 0] + result: dict[str, pd.DataFrame] = {} + for chunk in (await table.holes.to_dataframe(fb=fb)).itertuples(index=False): + hole_id = str(categories.cat.categories[int(chunk.hole_index)]) + result[hole_id] = data.iloc[int(chunk.offset) : int(chunk.offset) + int(chunk.count)].reset_index(drop=True) + return result diff --git a/packages/evo-objects/tests/test_downhole_utils.py b/packages/evo-objects/tests/test_downhole_utils.py index 6449e0be..08593121 100644 --- a/packages/evo-objects/tests/test_downhole_utils.py +++ b/packages/evo-objects/tests/test_downhole_utils.py @@ -31,3 +31,16 @@ def test_non_contiguous_and_unknown_ids_raise(self): hole_chunks_from_ids(pd.Series(["A", "B", "A"]), dtype=dtype) with self.assertRaises(ValueError): hole_chunks_from_ids(pd.Series(["C"]), dtype=dtype) + + def test_empty_input_emits_zero_count_entries_with_schema_dtypes(self): + chunks = hole_chunks_from_ids(pd.Series([], dtype="string"), dtype=pd.CategoricalDtype(categories=["A", "B"])) + self.assertListEqual(chunks["hole_index"].tolist(), [0, 1]) + self.assertListEqual(chunks["offset"].tolist(), [0, 0]) + self.assertListEqual(chunks["count"].tolist(), [0, 0]) + 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) diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 492f8269..1b575c32 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -1,454 +1,556 @@ -# 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 contextlib -import dataclasses -import math -import uuid -from datetime import date -from unittest.mock import patch - -import numpy as np -import numpy.testing as npt -import pandas as pd -from parameterized import parameterized - -from evo.common import Environment, StaticContext -from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector -from evo.objects import ObjectReference -from evo.objects.typed import BoundingBox -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 - -from .helpers import MockClient - - -def _make_example_data( - name: str = "Test DHC", - description: str | None = None, - tags: dict[str, str] | None = None, - attributes: pd.DataFrame | None = None, - collections: list[DistanceCollection] | None = None, -) -> DownholeCollectionData: - """Helper to build a simple two-hole DownholeCollectionData.""" - # Concatenated path table: hole 0 has 4 rows, hole 1 has 3 rows - path = pd.DataFrame( - { - "distance": [0.0, 10.0, 20.0, 30.0, 0.0, 15.0, 30.0], - "azimuth": [0.0, 0.0, 0.0, 0.0, 90.0, 90.0, 90.0], - "dip": [90.0, 90.0, 90.0, 90.0, 45.0, 45.0, 45.0], - } - ) - - collection1 = DistanceCollection( - name="collection1", - collection_type="distance", - holes=pd.DataFrame( - { - "hole_index": [0], - "offset": [0], - "count": [4], - } - ), - distance_table=pd.DataFrame( - { - "distance": [0.0, 10.0, 20.0, 30.0], - "attr_str": ["a", "b", "a", "c"], - "attr_dt": [date(2000, 1, 1), date(2000, 1, 2), date(2000, 1, 3), date(2000, 1, 4)], - "attr_num": [1.1, 2.2, 3.3, 4.4], - } - ), - ) - - holes = pd.DataFrame( - { - "hole_index": [0, 1], - "offset": [0, 4], - "count": [4, 3], - } - ) - - properties = pd.DataFrame( - { - "hole_id": ["H001", "H002"], - "x": [100.0, 200.0], - "y": [150.0, 300.0], - "z": [0.0, 50.0], - "final": [30.0, 30.0], - "target": [25.0, 25.0], - "current": [30.0, 28.0], - } - ) - - if collections is None: - collections = [collection1] - - return DownholeCollectionData( - name=name, - path=path, - holes=holes, - properties=properties, - attributes=attributes, - collections=collections, - distance_unit="m", - desurvey="trench", - description=description, - tags=tags, - ) - - -class TestDownholeCollection(TestWithConnector): - def setUp(self) -> None: - TestWithConnector.setUp(self) - self.environment = Environment(hub_url=BASE_URL, org_id=ORG.id, workspace_id=WORKSPACE_ID) - self.context = StaticContext.from_environment( - environment=self.environment, - connector=self.connector, - ) - - @contextlib.contextmanager - def _mock_geoscience_objects(self): - mock_client = MockClient(self.environment) - with ( - patch("evo.objects.typed.attributes.get_data_client", lambda _: mock_client), - patch("evo.objects.typed._data.get_data_client", lambda _: mock_client), - patch("evo.objects.typed._utils.get_data_client", lambda _: mock_client), - patch("evo.objects.typed.base.create_geoscience_object", mock_client.create_geoscience_object), - patch("evo.objects.typed.base.replace_geoscience_object", mock_client.replace_geoscience_object), - patch("evo.objects.DownloadedObject.from_context", mock_client.from_reference), - ): - yield mock_client - - def _assert_bounding_box_equal( - self, bbox: BoundingBox, min_x: float, max_x: float, min_y: float, max_y: float, min_z: float, max_z: float - ): - self.assertAlmostEqual(bbox.min_x, min_x, places=3) - self.assertAlmostEqual(bbox.max_x, max_x, places=3) - self.assertAlmostEqual(bbox.min_y, min_y, places=3) - self.assertAlmostEqual(bbox.max_y, max_y, places=3) - self.assertAlmostEqual(bbox.min_z, min_z, places=3) - self.assertAlmostEqual(bbox.max_z, max_z, places=3) - - async def _check_locations(self, expected: DownholeCollectionData, result: DownholeCollection): - loc = result.location - xyz = ["x", "y", "z"] - distances = ["final", "target", "current"] - - 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.holes, await loc.holes.to_dataframe()) - if expected.attributes: - npt.assert_array_equal(expected.attributes, await result.location.hole_id.to_dataframe()) - - async def _check_path(self, expected: DownholeCollectionData, result: DownholeCollection): - loc = result.location - path_columns = ["distance", "azimuth", "dip"] - attr_columns = [col for col in expected.path.columns if col not in path_columns] - - npt.assert_array_equal(expected.path[path_columns], await loc.path.to_dataframe()) - if attr_columns: - npt.assert_array_equal(expected.path[attr_columns], await loc.attributes.to_dataframe()) - - 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" - ) - self.assertEqual(expected_distance_unit, result_collection.distance.unit) - - expected_table = expected_collection.distance_table - result_table = await result_collection.distance.to_dataframe() - - for col in result_table.columns: - if pd.api.types.is_datetime64_any_dtype(result_table[col]): - for x, y in zip(expected_table[col], result_table[col]): - self.assertEqual(x.year, y.year) - self.assertEqual(x.month, y.month) - self.assertEqual(x.day, y.day) - else: - npt.assert_array_equal(expected_table[col], result_table[col]) - - async def _check_dhc(self, expected: DownholeCollectionData, result: DownholeCollection): - self.assertIsInstance(result, DownholeCollection) - self.assertEqual(expected.name, result.name) - self.assertEqual(expected.distance_unit, result.distance_unit) - self.assertEqual(expected.desurvey, result.desurvey) - - await self._check_locations(expected, result) - await self._check_path(expected, result) - await self._check_collections(expected, result) - - @parameterized.expand([BaseObject, DownholeCollection]) - async def test_create(self, class_to_call): - """Includes collections and attributes""" - data = _make_example_data() - with self._mock_geoscience_objects(): - result = await class_to_call.create(context=self.context, data=data) - await self._check_dhc(data, result) - - async def test_create_with_empty_collections(self): - data = _make_example_data(collections=[]) - with self._mock_geoscience_objects(): - 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]}), - interval_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_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() - with self._mock_geoscience_objects(): - result = await class_to_call.replace( - context=self.context, - reference=ObjectReference.new( - environment=self.context.get_environment(), - object_id=uuid.uuid4(), - ), - data=data, - ) - await self._check_dhc(data, result) - - @parameterized.expand([BaseObject, DownholeCollection]) - async def test_create_or_replace(self, class_to_call): - data = _make_example_data() - with self._mock_geoscience_objects(): - result = await class_to_call.create_or_replace( - context=self.context, - reference=ObjectReference.new( - environment=self.context.get_environment(), - object_id=uuid.uuid4(), - ), - data=data, - ) - await self._check_dhc(data, result) - - @parameterized.expand([BaseObject, DownholeCollection]) - async def test_from_reference(self, class_to_call): - data = _make_example_data() - with self._mock_geoscience_objects(): - original = await DownholeCollection.create(context=self.context, data=data) - result = await class_to_call.from_reference(context=self.context, reference=original.metadata.url) - await self._check_dhc(data, result) - - def test_bounding_box(self): - """Two vertical holes (dip=90deg) go straight down: bbox should reflect collar + depth. Azimuth doesn't matter""" - - path = pd.DataFrame( - { - "distance": [0.0, 10.0, 20.0, 30.0, 0.0, 15.0, 30.0], - "azimuth": [0.0, 45.0, 20.0, 0.0, 10.0, 90.0, 90.0], - "dip": [90.0, 90.0, 90.0, 90.0, 90.0, 90.0, 90.0], - } - ) - - data = _make_example_data() - data = dataclasses.replace(data, path=path) - 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_hole_index_not_property_position(self): - data = _make_example_data() - expected = data.compute_bounding_box() - properties = data.properties.iloc[[1, 0]].copy() - properties.index = [10, 20] - 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( - { - "distance": [0.0, 10.0, 20.0, 50.0, 0.0, 20.0, 40.0], - "azimuth": [0.0, 90.0, 180.0, 270.0, 0.0, 315.0, 90.0], - "dip": [60.0, 60.0, 60.0, 60.0, 60.0, 60.0, 60.0], - } - ) - - data = _make_example_data() - data = dataclasses.replace(data, path=path) - bbox = data.compute_bounding_box() - - # Expected geometry, based on having spiraled and zig-zagged with 30/60/90 and 45/45/90 dips/azimuths - xmin = 100.0 - 10 - xmax = 200.0 - 10 / math.sqrt(2) + 10 - ymin = 150.0 - 5 - ymax = 300.0 + 10 / math.sqrt(2) - zmin = (-50.0 / 2) * math.sqrt(3) - zmax = 50.0 - - self._assert_bounding_box_equal(bbox, xmin, xmax, ymin, ymax, zmin, zmax) - - def test_bounding_box_with_nans(self): - """Azimuth nans -> 0.0, dip nans -> 90.0""" - path = pd.DataFrame( - { - "distance": [0.0, 10.0, 20.0, 50.0, 0.0, 20.0, 40.0], - "azimuth": [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], - "dip": [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], - } - ) - - data = _make_example_data() - data = dataclasses.replace(data, path=path) - bbox = data.compute_bounding_box() - - self._assert_bounding_box_equal(bbox, 100.0, 200.0, 150.0, 300.0, -50.0, 50.0) - - def test_compute_bounding_box_np_unsorted_depths_raises(self): - with self.assertRaises(ObjectValidationError): - DownholeCollectionData._compute_bounding_box_np( - depths=np.array([10.0, 5.0, 20.0]), - dips=np.array([90.0, 90.0, 90.0]), - azimuths=np.array([0.0, 0.0, 0.0]), - ) - - def test_compute_bounding_box_np_length_mismatch_raises(self): - with self.assertRaises(ObjectValidationError): - DownholeCollectionData._compute_bounding_box_np( - depths=np.array([0.0, 10.0]), - dips=np.array([90.0]), - azimuths=np.array([0.0, 0.0]), - ) - - async def test_description_and_tags(self): - data = _make_example_data( - description="A test downhole collection", - tags={"site": "alpha", "status": "active"}, - ) - with self._mock_geoscience_objects(): - result = await DownholeCollection.create(context=self.context, data=data) - self.assertEqual(result.description, "A test downhole collection") - self.assertEqual(result.tags, {"site": "alpha", "status": "active"}) - - def test_attributes_length_raises(self): - """attributes length must match holes 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( - { - "hole_id": ["H1", "H2"], - "x": [0.0, 1.0], - "y": [0.0, 1.0], - "z": [0.0, 0.0], - "final": [10.0, 10.0], - "target": [10.0, 10.0], - "current": [10.0, 10.0], - } - ) - # attributes has 3 rows, but holes has 2 - should assert - bad_attributes = pd.DataFrame({"a": [1, 2, 3]}) - with self.assertRaises(ObjectValidationError): - DownholeCollectionData( - name="Bad", - path=path, - holes=holes, - properties=properties, - attributes=bad_attributes, - collections=[], - distance_unit=None, - desurvey=None, - ) - - async def test_update_dataframe_after_creation(self): - """Test updating the path DataFrame after downhole collection creation.""" - with self._mock_geoscience_objects(): - data = _make_example_data() - obj = await DownholeCollection.create(context=self.context, data=data) - - new_path = pd.DataFrame( - { - "distance": [0.0, 10.0, 20.0, 50.0, 0.0, 20.0, 40.0], - "azimuth": [0.0, 90.0, 180.0, 270.0, 0.0, 315.0, 90.0], - "dip": [60.0, 60.0, 60.0, 60.0, 60.0, 60.0, 60.0], - } - ) - await obj.location.path.from_dataframe(new_path) - - # Verify the data was updated - await obj.update() - expected = dataclasses.replace(data, path=new_path) - await self._check_dhc(expected, obj) - - async def test_json(self): - data = _make_example_data() - with self._mock_geoscience_objects() as mock_client: - obj = await DownholeCollection.create(context=self.context, data=data) - object_json = mock_client.objects[str(obj.metadata.url.object_id)] - - # Verify schema - self.assertIn("/objects/downhole-collection/", object_json["schema"]) - - # Verify base properties - self.assertEqual(object_json["name"], "Test DHC") - self.assertIn("uuid", object_json) - self.assertIn("bounding_box", object_json) - self.assertEqual(object_json["coordinate_reference_system"], "unspecified") - - # Verify DHC top level properties - self.assertEqual(object_json["type"], "downhole") - self.assertIn("distance_unit", object_json) - self.assertIn("desurvey", object_json) - - # Verify location structure - self.assertIn("location", object_json) - location = object_json["location"] - self.assertIn("path", location) - self.assertIn("holes", location) - self.assertIn("coordinates", location) - self.assertIn("distances", location) - self.assertIn("hole_id", location) - self.assertIn("collections", object_json) - collection = object_json["collections"][0] - self.assertIn("name", collection) - self.assertIn("collection_type", collection) - self.assertEqual(collection["collection_type"], "distance") - self.assertIn("holes", collection) - self.assertIn("distance", collection) +# 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 contextlib +import dataclasses +import math +import uuid +from datetime import date +from unittest.mock import patch + +import numpy as np +import numpy.testing as npt +import pandas as pd +from parameterized import parameterized + +from evo.common import Environment, StaticContext +from evo.common.test_tools import BASE_URL, ORG, WORKSPACE_ID, TestWithConnector +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 + +from .helpers import MockClient + + +def _make_example_data( + name: str = "Test DHC", + description: str | None = None, + tags: dict[str, str] | None = None, + attributes: pd.DataFrame | None = None, + collections: list[DistanceCollection] | None = None, +) -> DownholeCollectionData: + """Helper to build a simple two-hole DownholeCollectionData.""" + # Concatenated path table: hole 0 has 4 rows, hole 1 has 3 rows + path = pd.DataFrame( + { + "distance": [0.0, 10.0, 20.0, 30.0, 0.0, 15.0, 30.0], + "azimuth": [0.0, 0.0, 0.0, 0.0, 90.0, 90.0, 90.0], + "dip": [90.0, 90.0, 90.0, 90.0, 45.0, 45.0, 45.0], + } + ) + + collection1 = DistanceCollection( + name="collection1", + collection_type="distance", + holes=pd.DataFrame( + { + "hole_index": [0], + "offset": [0], + "count": [4], + } + ), + distance_table=pd.DataFrame( + { + "distance": [0.0, 10.0, 20.0, 30.0], + "attr_str": ["a", "b", "a", "c"], + "attr_dt": [date(2000, 1, 1), date(2000, 1, 2), date(2000, 1, 3), date(2000, 1, 4)], + "attr_num": [1.1, 2.2, 3.3, 4.4], + } + ), + ) + + holes = pd.DataFrame( + { + "hole_index": [0, 1], + "offset": [0, 4], + "count": [4, 3], + } + ) + + properties = pd.DataFrame( + { + "hole_id": ["H001", "H002"], + "x": [100.0, 200.0], + "y": [150.0, 300.0], + "z": [0.0, 50.0], + "final": [30.0, 30.0], + "target": [25.0, 25.0], + "current": [30.0, 28.0], + } + ) + + if collections is None: + collections = [collection1] + + return DownholeCollectionData( + name=name, + path=path, + holes=holes, + properties=properties, + attributes=attributes, + collections=collections, + distance_unit="m", + desurvey="trench", + description=description, + tags=tags, + ) + + +class TestDownholeCollection(TestWithConnector): + def setUp(self) -> None: + TestWithConnector.setUp(self) + self.environment = Environment(hub_url=BASE_URL, org_id=ORG.id, workspace_id=WORKSPACE_ID) + self.context = StaticContext.from_environment( + environment=self.environment, + connector=self.connector, + ) + + @contextlib.contextmanager + def _mock_geoscience_objects(self): + mock_client = MockClient(self.environment) + with ( + patch("evo.objects.typed.attributes.get_data_client", lambda _: mock_client), + patch("evo.objects.typed._data.get_data_client", lambda _: mock_client), + patch("evo.objects.typed._utils.get_data_client", lambda _: mock_client), + patch("evo.objects.typed.base.create_geoscience_object", mock_client.create_geoscience_object), + patch("evo.objects.typed.base.replace_geoscience_object", mock_client.replace_geoscience_object), + patch("evo.objects.DownloadedObject.from_context", mock_client.from_reference), + ): + yield mock_client + + def _assert_bounding_box_equal( + self, bbox: BoundingBox, min_x: float, max_x: float, min_y: float, max_y: float, min_z: float, max_z: float + ): + self.assertAlmostEqual(bbox.min_x, min_x, places=3) + self.assertAlmostEqual(bbox.max_x, max_x, places=3) + self.assertAlmostEqual(bbox.min_y, min_y, places=3) + self.assertAlmostEqual(bbox.max_y, max_y, places=3) + self.assertAlmostEqual(bbox.min_z, min_z, places=3) + self.assertAlmostEqual(bbox.max_z, max_z, places=3) + + async def _check_locations(self, expected: DownholeCollectionData, result: DownholeCollection): + loc = result.location + xyz = ["x", "y", "z"] + distances = ["final", "target", "current"] + + 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.holes, await loc.holes.to_dataframe()) + if expected.attributes: + npt.assert_array_equal(expected.attributes, await result.location.hole_id.to_dataframe()) + + async def _check_path(self, expected: DownholeCollectionData, result: DownholeCollection): + loc = result.location + path_columns = ["distance", "azimuth", "dip"] + attr_columns = [col for col in expected.path.columns if col not in path_columns] + + npt.assert_array_equal(expected.path[path_columns], await loc.path.to_dataframe()) + if attr_columns: + npt.assert_array_equal(expected.path[attr_columns], await loc.attributes.to_dataframe()) + + 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" + ) + self.assertEqual(expected_distance_unit, result_collection.distance.unit) + + expected_table = expected_collection.distance_table + result_table = await result_collection.distance.to_dataframe() + + for col in result_table.columns: + if pd.api.types.is_datetime64_any_dtype(result_table[col]): + for x, y in zip(expected_table[col], result_table[col]): + self.assertEqual(x.year, y.year) + self.assertEqual(x.month, y.month) + self.assertEqual(x.day, y.day) + else: + npt.assert_array_equal(expected_table[col], result_table[col]) + + async def _check_dhc(self, expected: DownholeCollectionData, result: DownholeCollection): + self.assertIsInstance(result, DownholeCollection) + self.assertEqual(expected.name, result.name) + self.assertEqual(expected.distance_unit, result.distance_unit) + self.assertEqual(expected.desurvey, result.desurvey) + + await self._check_locations(expected, result) + await self._check_path(expected, result) + await self._check_collections(expected, result) + + @parameterized.expand([BaseObject, DownholeCollection]) + async def test_create(self, class_to_call): + """Includes collections and attributes""" + data = _make_example_data() + with self._mock_geoscience_objects(): + result = await class_to_call.create(context=self.context, data=data) + await self._check_dhc(data, result) + + async def test_create_with_empty_collections(self): + data = _make_example_data(collections=[]) + with self._mock_geoscience_objects(): + 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]}), + interval_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_only_collection_uses_attribute_unit_when_not_explicit(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]}), + interval_table=table, + ) + data = _make_example_data(collections=[interval]) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=data) + collection = result.collections.get("intervals") + self.assertIsNotNone(collection) + self.assertEqual(collection.from_to.unit, "ft") + self.assertListEqual((await collection.to_dataframe()).columns.tolist(), ["from", "to", "grade"]) + + async def test_explicit_collection_unit_overrides_dataframe_metadata(self): + 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]}), + distance_table=table, + unit="m", + ) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create( + context=self.context, data=_make_example_data(collections=[collection]) + ) + self.assertEqual(result.collections.get("distances").distance.unit, "m") + + 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]}), + distance_table=table, + ) + 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_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() + with self._mock_geoscience_objects(): + result = await class_to_call.replace( + context=self.context, + reference=ObjectReference.new( + environment=self.context.get_environment(), + object_id=uuid.uuid4(), + ), + data=data, + ) + await self._check_dhc(data, result) + + @parameterized.expand([BaseObject, DownholeCollection]) + async def test_create_or_replace(self, class_to_call): + data = _make_example_data() + with self._mock_geoscience_objects(): + result = await class_to_call.create_or_replace( + context=self.context, + reference=ObjectReference.new( + environment=self.context.get_environment(), + object_id=uuid.uuid4(), + ), + data=data, + ) + await self._check_dhc(data, result) + + @parameterized.expand([BaseObject, DownholeCollection]) + async def test_from_reference(self, class_to_call): + data = _make_example_data() + with self._mock_geoscience_objects(): + original = await DownholeCollection.create(context=self.context, data=data) + result = await class_to_call.from_reference(context=self.context, reference=original.metadata.url) + await self._check_dhc(data, result) + + def test_bounding_box(self): + """Two vertical holes (dip=90deg) go straight down: bbox should reflect collar + depth. Azimuth doesn't matter""" + + path = pd.DataFrame( + { + "distance": [0.0, 10.0, 20.0, 30.0, 0.0, 15.0, 30.0], + "azimuth": [0.0, 45.0, 20.0, 0.0, 10.0, 90.0, 90.0], + "dip": [90.0, 90.0, 90.0, 90.0, 90.0, 90.0, 90.0], + } + ) + + data = _make_example_data() + data = dataclasses.replace(data, path=path) + 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_hole_index_not_property_position(self): + data = _make_example_data() + expected = data.compute_bounding_box() + properties = data.properties.iloc[[1, 0]].copy() + properties.index = [10, 20] + 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( + { + "distance": [0.0, 10.0, 20.0, 50.0, 0.0, 20.0, 40.0], + "azimuth": [0.0, 90.0, 180.0, 270.0, 0.0, 315.0, 90.0], + "dip": [60.0, 60.0, 60.0, 60.0, 60.0, 60.0, 60.0], + } + ) + + data = _make_example_data() + data = dataclasses.replace(data, path=path) + bbox = data.compute_bounding_box() + + # Expected geometry, based on having spiraled and zig-zagged with 30/60/90 and 45/45/90 dips/azimuths + xmin = 100.0 - 10 + xmax = 200.0 - 10 / math.sqrt(2) + 10 + ymin = 150.0 - 5 + ymax = 300.0 + 10 / math.sqrt(2) + zmin = (-50.0 / 2) * math.sqrt(3) + zmax = 50.0 + + self._assert_bounding_box_equal(bbox, xmin, xmax, ymin, ymax, zmin, zmax) + + def test_bounding_box_with_nans(self): + """Azimuth nans -> 0.0, dip nans -> 90.0""" + path = pd.DataFrame( + { + "distance": [0.0, 10.0, 20.0, 50.0, 0.0, 20.0, 40.0], + "azimuth": [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], + "dip": [np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan], + } + ) + + data = _make_example_data() + data = dataclasses.replace(data, path=path) + bbox = data.compute_bounding_box() + + self._assert_bounding_box_equal(bbox, 100.0, 200.0, 150.0, 300.0, -50.0, 50.0) + + def test_compute_bounding_box_np_unsorted_depths_raises(self): + with self.assertRaises(ObjectValidationError): + DownholeCollectionData._compute_bounding_box_np( + depths=np.array([10.0, 5.0, 20.0]), + dips=np.array([90.0, 90.0, 90.0]), + azimuths=np.array([0.0, 0.0, 0.0]), + ) + + def test_compute_bounding_box_np_length_mismatch_raises(self): + with self.assertRaises(ObjectValidationError): + DownholeCollectionData._compute_bounding_box_np( + depths=np.array([0.0, 10.0]), + dips=np.array([90.0]), + azimuths=np.array([0.0, 0.0]), + ) + + async def test_description_and_tags(self): + data = _make_example_data( + description="A test downhole collection", + tags={"site": "alpha", "status": "active"}, + ) + with self._mock_geoscience_objects(): + result = await DownholeCollection.create(context=self.context, data=data) + self.assertEqual(result.description, "A test downhole collection") + self.assertEqual(result.tags, {"site": "alpha", "status": "active"}) + + def test_attributes_length_raises(self): + """attributes length must match holes 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( + { + "hole_id": ["H1", "H2"], + "x": [0.0, 1.0], + "y": [0.0, 1.0], + "z": [0.0, 0.0], + "final": [10.0, 10.0], + "target": [10.0, 10.0], + "current": [10.0, 10.0], + } + ) + # attributes has 3 rows, but holes has 2 - should assert + bad_attributes = pd.DataFrame({"a": [1, 2, 3]}) + with self.assertRaises(ObjectValidationError): + DownholeCollectionData( + name="Bad", + path=path, + holes=holes, + properties=properties, + attributes=bad_attributes, + collections=[], + distance_unit=None, + desurvey=None, + ) + + def test_collection_chunk_overlap_gap_and_invalid_index_raise(self): + base = _make_example_data(collections=[]) + table = pd.DataFrame({"distance": [0.0, 1.0]}) + for holes in ( + pd.DataFrame({"hole_index": [0, 1], "offset": [0, 0], "count": [1, 1]}), + pd.DataFrame({"hole_index": [0], "offset": [1], "count": [1]}), + pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), + ): + with self.assertRaises(ObjectValidationError): + dataclasses.replace( + base, + collections=[DistanceCollection(name="invalid", holes=holes, distance_table=table)], + ) + + async def test_collection_add_rejects_duplicate_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]}), + distance_table=pd.DataFrame({"distance": [0.0]}), + ) + await result.collections.add(collection) + replacement = dataclasses.replace(collection, distance_table=pd.DataFrame({"distance": [1.0]})) + with self.assertRaises(ValueError): + 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) + invalid = dataclasses.replace( + collection, + holes=pd.DataFrame({"hole_index": [0, 0], "offset": [0, 1], "count": [1, 0]}), + ) + with self.assertRaises(ObjectValidationError): + await result.collections.add(invalid, replace=True) + + 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_update_dataframe_after_creation(self): + """Test updating the path DataFrame after downhole collection creation.""" + with self._mock_geoscience_objects(): + data = _make_example_data() + obj = await DownholeCollection.create(context=self.context, data=data) + + new_path = pd.DataFrame( + { + "distance": [0.0, 10.0, 20.0, 50.0, 0.0, 20.0, 40.0], + "azimuth": [0.0, 90.0, 180.0, 270.0, 0.0, 315.0, 90.0], + "dip": [60.0, 60.0, 60.0, 60.0, 60.0, 60.0, 60.0], + } + ) + await obj.location.path.from_dataframe(new_path) + + # Verify the data was updated + await obj.update() + expected = dataclasses.replace(data, path=new_path) + await self._check_dhc(expected, obj) + + async def test_json(self): + data = _make_example_data() + with self._mock_geoscience_objects() as mock_client: + obj = await DownholeCollection.create(context=self.context, data=data) + object_json = mock_client.objects[str(obj.metadata.url.object_id)] + + # Verify schema + self.assertIn("/objects/downhole-collection/", object_json["schema"]) + + # Verify base properties + self.assertEqual(object_json["name"], "Test DHC") + self.assertIn("uuid", object_json) + self.assertIn("bounding_box", object_json) + self.assertEqual(object_json["coordinate_reference_system"], "unspecified") + + # Verify DHC top level properties + self.assertEqual(object_json["type"], "downhole") + self.assertIn("distance_unit", object_json) + self.assertIn("desurvey", object_json) + + # Verify location structure + self.assertIn("location", object_json) + location = object_json["location"] + self.assertIn("path", location) + self.assertIn("holes", location) + self.assertIn("coordinates", location) + self.assertIn("distances", location) + self.assertIn("hole_id", location) + self.assertIn("collections", object_json) + collection = object_json["collections"][0] + self.assertIn("name", collection) + self.assertIn("collection_type", collection) + self.assertEqual(collection["collection_type"], "distance") + self.assertIn("holes", collection) + self.assertIn("distance", collection) diff --git a/packages/evo-objects/tests/typed/test_model.py b/packages/evo-objects/tests/typed/test_model.py index 1243508c..b600d106 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,6 +59,18 @@ 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 Second(SchemaModel): + pass + + with self.assertRaisesRegex(TypeError, r"must override _resolve_item_type\(\) and _data_to_schema\(\)"): + + class InvalidUnionList(SchemaList[First | Second]): + pass + class TestTable(DataTable): table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_3 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..911f7230 --- /dev/null +++ b/packages/evo-objects/tests/typed/test_prefetch.py @@ -0,0 +1,96 @@ +# 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. + +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) From c9c4561c9745d68f6e14219e915964892eed8796 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Wed, 5 Aug 2026 11:25:40 +0100 Subject: [PATCH 06/30] feat(objects): enhance collection handling and add new tests for attribute behavior --- .../evo-objects/tests/test_downhole_utils.py | 77 +++++++++++++++++++ packages/evo-objects/tests/typed/helpers.py | 2 + .../tests/typed/test_downhole_collection.py | 47 +++++++++++ 3 files changed, 126 insertions(+) diff --git a/packages/evo-objects/tests/test_downhole_utils.py b/packages/evo-objects/tests/test_downhole_utils.py index 08593121..b884af8f 100644 --- a/packages/evo-objects/tests/test_downhole_utils.py +++ b/packages/evo-objects/tests/test_downhole_utils.py @@ -17,6 +17,83 @@ class TestDownholeUtils(unittest.TestCase): + def _chunks(self, values, categories): + return hole_chunks_from_ids(pd.Series(values), dtype=pd.CategoricalDtype(categories=categories)) + + def test_single_hole_uses_zero_based_code(self): + chunks = self._chunks(["A"], ["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"], ["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"], ["A", "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"], ["A", "B", "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_category_absent_from_data_emits_zero_count_chunk(self): + chunks = self._chunks(["A", "A"], ["A", "B"]) + self.assertListEqual( + chunks.to_dict("records"), + [{"hole_index": 0, "offset": 0, "count": 2}, {"hole_index": 1, "offset": 0, "count": 0}], + ) + + def test_id_absent_from_categories_raises(self): + with self.assertRaises(ValueError): + self._chunks(["A", "C"], ["A", "B"]) + + def test_empty_input_emits_all_zero_count_chunks(self): + chunks = self._chunks([], ["A", "B"]) + self.assertListEqual( + chunks.to_dict("records"), + [{"hole_index": 0, "offset": 0, "count": 0}, {"hole_index": 1, "offset": 0, "count": 0}], + ) + + def test_hole_index_dtype_is_int32(self): + self.assertEqual(str(self._chunks(["A"], ["A"])["hole_index"].dtype), "int32") + + def test_offset_and_count_dtypes_are_uint64(self): + chunks = self._chunks(["A"], ["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): + categories = [f"H{index:02d}" for index in range(50)] + values = [hole_id for hole_id in categories for _ in range(200)] + chunks = self._chunks(values, categories) + 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_category_order_is_preserved_not_lexicographically_sorted(self): + chunks = self._chunks(["M", "M", "Z"], ["Z", "A", "M"]) + self.assertListEqual( + chunks.to_dict("records"), + [ + {"hole_index": 0, "offset": 2, "count": 1}, + {"hole_index": 1, "offset": 0, "count": 0}, + {"hole_index": 2, "offset": 0, "count": 2}, + ], + ) + def test_chunks_use_dtype_codes_and_round_trip(self): dtype = pd.CategoricalDtype(categories=["Z", "A", "M"]) values = pd.Series(["M", "M", "Z"]) diff --git a/packages/evo-objects/tests/typed/helpers.py b/packages/evo-objects/tests/typed/helpers.py index 630766a3..8860c6ea 100644 --- a/packages/evo-objects/tests/typed/helpers.py +++ b/packages/evo-objects/tests/typed/helpers.py @@ -84,6 +84,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) diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 1b575c32..8d329606 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -284,6 +284,11 @@ async def test_attribute_descriptions_round_trip_to_dataframe_metadata(self): 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) @@ -465,6 +470,19 @@ def test_collection_chunk_overlap_gap_and_invalid_index_raise(self): collections=[DistanceCollection(name="invalid", holes=holes, distance_table=table)], ) + 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]}), + distance_table=pd.DataFrame({"distance": [0.0, 1.0]}), + ) + 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_rejects_duplicate_hole_chunks_and_replaces_in_place(self): with self._mock_geoscience_objects(): result = await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[])) @@ -487,6 +505,35 @@ async def test_collection_add_rejects_duplicate_hole_chunks_and_replaces_in_plac with self.assertRaises(ObjectValidationError): await result.collections.add(invalid, 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]}), + distance_table=pd.DataFrame({"distance": [0.0]}), + ) + middle = DistanceCollection( + name="middle", + holes=pd.DataFrame({"hole_index": [1], "offset": [0], "count": [1]}), + distance_table=pd.DataFrame({"distance": [1.0]}), + ) + last = DistanceCollection( + name="last", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + distance_table=pd.DataFrame({"distance": [2.0]}), + ) + replacement = dataclasses.replace(middle, distance_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()) From 8ecca39663eae7ba840a2a21cbbe30e85647789a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Wed, 5 Aug 2026 11:33:02 +0100 Subject: [PATCH 07/30] feat(objects-tests): update categorical data generation and improve table loading functions - removal of depreaction warnings --- packages/evo-objects/tests/helpers.py | 2 +- packages/evo-objects/tests/test_tables.py | 28 ++++++++++++------- .../evo-objects/tests/typed/test_model.py | 4 +-- .../tests/typed/test_regular_grid.py | 2 +- .../tests/typed/test_regular_masked_grid.py | 2 +- .../tests/typed/test_tensor_grid.py | 2 +- 6 files changed, 24 insertions(+), 16 deletions(-) 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_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/test_model.py b/packages/evo-objects/tests/typed/test_model.py index b600d106..a432590b 100644 --- a/packages/evo-objects/tests/typed/test_model.py +++ b/packages/evo-objects/tests/typed/test_model.py @@ -72,13 +72,13 @@ class InvalidUnionList(SchemaList[First | Second]): pass -class TestTable(DataTable): +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_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( From 2a8a2d215d805cc78d1841597291279607849d45 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Wed, 5 Aug 2026 11:50:28 +0100 Subject: [PATCH 08/30] feat(objects-tests): add unit precedence and attribute type validation tests for interval collections --- .../tests/typed/test_downhole_collection.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 8d329606..7168eb93 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -247,6 +247,87 @@ async def test_interval_only_collection_uses_attribute_unit_when_not_explicit(se self.assertEqual(collection.from_to.unit, "ft") self.assertListEqual((await collection.to_dataframe()).columns.tolist(), ["from", "to", "grade"]) + @parameterized.expand( + [ + ("explicit", "m", "ft", "m"), + ("metadata", None, "ft", "ft"), + ("omitted", None, None, None), + ] + ) + async def test_interval_collection_unit_precedence(self, _name, explicit_unit, metadata_unit, expected_unit): + table = pd.DataFrame({"from": [0.0], "to": [1.0]}) + if metadata_unit is not None: + table.attrs["attribute_descriptions"] = {"from": AttributeDescription(unit=metadata_unit)} + collection = IntervalCollection( + name="intervals", + holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), + interval_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]}), + interval_table=table, + ) + 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]}), + interval_table=pd.DataFrame(table_data), + ) + with self._mock_geoscience_objects(), self.assertRaises(ObjectValidationError): + await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[collection])) + async def test_explicit_collection_unit_overrides_dataframe_metadata(self): table = pd.DataFrame({"distance": [0.0], "grade": [1.0]}) table.attrs["attribute_descriptions"] = {"distance": AttributeDescription(unit="ft")} From 045ae5b0146ba54ff65c33c565e6f4fc3f16e4cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Thu, 6 Aug 2026 07:55:27 +0100 Subject: [PATCH 09/30] chore(objects-test): update license information in test_prefetch.py --- packages/evo-objects/tests/typed/test_prefetch.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/packages/evo-objects/tests/typed/test_prefetch.py b/packages/evo-objects/tests/typed/test_prefetch.py index 911f7230..262f758d 100644 --- a/packages/evo-objects/tests/typed/test_prefetch.py +++ b/packages/evo-objects/tests/typed/test_prefetch.py @@ -1,6 +1,13 @@ # 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 From fba1243f93e6061234545a79e616795946f0c4d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 08:48:36 +0100 Subject: [PATCH 10/30] feat(objects): update AttributeDescription to use default empty strings --- .../evo-objects/src/evo/objects/typed/attributes.py | 13 ++++++------- packages/evo-objects/tests/typed/test_attributes.py | 8 ++++++-- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 6e1defbb..153fcbca 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -88,8 +88,8 @@ def _infer_attribute_type_from_series(series: pd.Series) -> str: @dataclass class AttributeDescription: - discipline: str | None = None - type: str | None = None + discipline: str = "" + type: str = "" unit: Any | None = None scale: str | None = None extensions: dict[str, typing.Any] | None = None @@ -100,11 +100,10 @@ def __post_init__(self) -> None: self.unit = str(self.unit.value) def to_schema(self) -> dict[str, Any]: - result: dict[str, Any] = {} - if self.discipline: - result["discipline"] = self.discipline - if self.type: - result["type"] = self.type + result: dict[str, Any] = { + "discipline": self.discipline, + "type": self.type, + } if self.unit: result["unit"] = self.unit if self.scale: diff --git a/packages/evo-objects/tests/typed/test_attributes.py b/packages/evo-objects/tests/typed/test_attributes.py index ad2004ff..a0ad59da 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -73,8 +73,12 @@ def test_pending_attribute_repr(self): class TestAttributeDescription(TestCase): - def test_empty_description_is_omitted(self): - self.assertEqual(AttributeDescription().to_schema(), {}) + 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_normalizes_value_units(self): class Unit: From dc4ee30866a493f15feffabb12b01195e7ae291e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 09:04:23 +0100 Subject: [PATCH 11/30] refactor(objects): unify table attribute naming in DistanceCollection and IntervalCollection --- packages/evo-objects/README.md | 2 +- .../evo/objects/typed/downhole_collection.py | 21 ++++------ .../tests/typed/test_downhole_collection.py | 38 +++++++++---------- 3 files changed, 27 insertions(+), 34 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index d3721e3b..c1c3414b 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -88,7 +88,7 @@ collections = [ IntervalCollection( name="geology", holes=hole_chunks_from_ids(pd.Series(["DH-01"]), dtype=hole_dtype), - interval_table=intervals, + table=intervals, unit="m", # Explicit collection units override DataFrame metadata. ) ] 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 41447558..c5845b6a 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -71,7 +71,7 @@ class DistanceCollection: name: str holes: HoleChunks - distance_table: Depths + table: Depths collection_type: str = "distance" unit: str | None = None @@ -80,7 +80,7 @@ class DistanceCollection: class IntervalCollection: name: str holes: HoleChunks - interval_table: Intervals + table: Intervals collection_type: str = "interval" unit: str | None = None @@ -134,9 +134,7 @@ def __post_init__(self): self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) for collection in self.collections: - table = ( - collection.distance_table if isinstance(collection, DistanceCollection) else collection.interval_table - ) + table = collection.table self._validate_hole_chunks(collection.holes, len(table), require_coverage=False) def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require_coverage: bool) -> None: @@ -357,7 +355,7 @@ async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: 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")] class DownholeDistanceTable(DistanceTable): @@ -394,7 +392,7 @@ async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: class DownholeIntervalTable(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("interval_table")] + from_to: Annotated[IntervalTableFromTo, SchemaLocation("from_to"), DataLocation("table")] holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: @@ -416,15 +414,12 @@ async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: result = [] for collection in data: model = DownholeIntervalTable if isinstance(collection, IntervalCollection) else DownholeDistanceTable - table = ( - collection.interval_table if isinstance(collection, IntervalCollection) else collection.distance_table - ) + table = collection.table table = table.copy() table.attrs = dict(table.attrs) if collection.unit is not None: table.attrs["unit"] = collection.unit - field = "interval_table" if isinstance(collection, IntervalCollection) else "distance_table" - result.append(await model._data_to_schema(replace(collection, **{field: table}), context)) + result.append(await model._data_to_schema(replace(collection, table=table), context)) return result def get(self, name: str) -> DownholeDistanceTable | DownholeIntervalTable | None: @@ -444,7 +439,7 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals valid_indices = set(location_holes["hole_index"].astype(int)) if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): raise ObjectValidationError("Collection hole_index is not present in the location holes table") - table = collection.distance_table if isinstance(collection, DistanceCollection) else collection.interval_table + table = collection.table _validate_chunk_ranges(collection.holes, len(table)) schema = await self._data_to_schema([collection], self._obj) if collection.name in existing: diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 7168eb93..b5d39372 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -67,7 +67,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"], @@ -169,12 +169,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: @@ -214,7 +212,7 @@ async def test_mixed_collections_round_trip_and_mutation(self): interval = IntervalCollection( name="intervals", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), - interval_table=pd.DataFrame({"from": [0.0, 1.0], "to": [1.0, 2.0], "lithology": ["a", "b"]}), + 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]) @@ -237,7 +235,7 @@ async def test_interval_only_collection_uses_attribute_unit_when_not_explicit(se interval = IntervalCollection( name="intervals", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), - interval_table=table, + table=table, ) data = _make_example_data(collections=[interval]) with self._mock_geoscience_objects(): @@ -261,7 +259,7 @@ async def test_interval_collection_unit_precedence(self, _name, explicit_unit, m collection = IntervalCollection( name="intervals", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - interval_table=table, + table=table, unit=explicit_unit, ) with self._mock_geoscience_objects() as mock_client: @@ -292,7 +290,7 @@ async def test_interval_collection_round_trips_all_attribute_types(self): collection = IntervalCollection( name="intervals", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), - interval_table=table, + table=table, ) expected_types = { "scalar": "scalar", @@ -323,7 +321,7 @@ async def test_interval_collection_requires_from_and_to_columns(self, _name, tab collection = IntervalCollection( name="invalid", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - interval_table=pd.DataFrame(table_data), + table=pd.DataFrame(table_data), ) with self._mock_geoscience_objects(), self.assertRaises(ObjectValidationError): await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[collection])) @@ -334,7 +332,7 @@ async def test_explicit_collection_unit_overrides_dataframe_metadata(self): collection = DistanceCollection( name="distances", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - distance_table=table, + table=table, unit="m", ) with self._mock_geoscience_objects(): @@ -356,7 +354,7 @@ async def test_attribute_descriptions_round_trip_to_dataframe_metadata(self): collection = DistanceCollection( name="grades", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - distance_table=table, + table=table, ) with self._mock_geoscience_objects(): result = await DownholeCollection.create( @@ -548,7 +546,7 @@ def test_collection_chunk_overlap_gap_and_invalid_index_raise(self): with self.assertRaises(ObjectValidationError): dataclasses.replace( base, - collections=[DistanceCollection(name="invalid", holes=holes, distance_table=table)], + collections=[DistanceCollection(name="invalid", holes=holes, table=table)], ) async def test_collection_allows_zero_row_hole(self): @@ -556,7 +554,7 @@ async def test_collection_allows_zero_row_hole(self): collection = DistanceCollection( name="sparse", holes=pd.DataFrame({"hole_index": [0, 1], "offset": [0, 0], "count": [2, 0]}), - distance_table=pd.DataFrame({"distance": [0.0, 1.0]}), + table=pd.DataFrame({"distance": [0.0, 1.0]}), ) data = dataclasses.replace(base, collections=[collection]) with self._mock_geoscience_objects(): @@ -570,10 +568,10 @@ async def test_collection_add_rejects_duplicate_hole_chunks_and_replaces_in_plac collection = DistanceCollection( name="measurements", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - distance_table=pd.DataFrame({"distance": [0.0]}), + table=pd.DataFrame({"distance": [0.0]}), ) await result.collections.add(collection) - replacement = dataclasses.replace(collection, distance_table=pd.DataFrame({"distance": [1.0]})) + replacement = dataclasses.replace(collection, table=pd.DataFrame({"distance": [1.0]})) with self.assertRaises(ValueError): await result.collections.add(replacement) await result.collections.add(replacement, replace=True) @@ -590,19 +588,19 @@ async def test_collection_replacement_preserves_position_and_persists_after_upda first = DistanceCollection( name="first", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - distance_table=pd.DataFrame({"distance": [0.0]}), + table=pd.DataFrame({"distance": [0.0]}), ) middle = DistanceCollection( name="middle", holes=pd.DataFrame({"hole_index": [1], "offset": [0], "count": [1]}), - distance_table=pd.DataFrame({"distance": [1.0]}), + table=pd.DataFrame({"distance": [1.0]}), ) last = DistanceCollection( name="last", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), - distance_table=pd.DataFrame({"distance": [2.0]}), + table=pd.DataFrame({"distance": [2.0]}), ) - replacement = dataclasses.replace(middle, distance_table=pd.DataFrame({"distance": [10.0]})) + 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) From 2dbb04b2e5597c4a91d875a0e9aaac7ff6ed081f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 09:47:41 +0100 Subject: [PATCH 12/30] feat(objects): Keep dense zero-based categorical codes as the SDK write convention, but treat persisted hole_index values as keys from the hole_id lookup table when reading existing objects. --- .../src/evo/objects/typed/attributes.py | 11 +++ .../evo/objects/typed/downhole_collection.py | 48 ++++++------ packages/evo-objects/tests/typed/helpers.py | 7 +- .../tests/typed/test_downhole_collection.py | 75 +++++++++++++++---- 4 files changed, 105 insertions(+), 36 deletions(-) diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 153fcbca..96566dd3 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -621,3 +621,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/downhole_collection.py b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py index c5845b6a..873330c9 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -94,8 +94,10 @@ class DownholeCollectionData(BaseSpatialObjectData): :param name: The name of the object. :param holes: A DataFrame describing which parts of `path` belong to which holes. - Columns: hole_index, offset, count. ``hole_index`` is a code in ``properties.hole_id``. - :param properties: DataFrame for the properties of the holes, joined to ``holes`` by the categorical hole-id code. + Columns: hole_index, offset, count. For object creation, ``hole_index`` is the zero-based categorical code + in ``properties.hole_id``. + :param properties: DataFrame for the properties of the holes. Its categorical ``hole_id`` codes are the lookup + keys referenced by creation-time ``holes`` tables. Mandatory columns: hole_id, final, target, current, x, y, z :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 @@ -120,7 +122,7 @@ class DownholeCollectionData(BaseSpatialObjectData): @property def hole_id_dtype(self) -> pd.CategoricalDtype: - """The categorical dtype whose codes are used by every ``hole_index``.""" + """The categorical dtype used to encode hole indices when creating this object.""" hole_ids = self.properties["hole_id"] if isinstance(hole_ids.dtype, pd.CategoricalDtype): return hole_ids.dtype @@ -145,7 +147,7 @@ def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require valid = set(range(len(self.hole_id_dtype.categories))) if not set(indices).issubset(valid): raise ObjectValidationError("hole_index must be a code in properties['hole_id'] categorical dtype") - if indices.duplicated().any(): + if require_coverage and indices.duplicated().any(): raise ObjectValidationError("Each hole_index may occur only once in a holes table") if require_coverage and set(indices) != valid: raise ObjectValidationError("Location holes must cover every hole_id categorical code exactly once") @@ -153,6 +155,9 @@ def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require 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_coverage: + return + non_empty = sorted(zip(offsets[counts > 0], counts[counts > 0], strict=True)) expected_offset = 0 for offset, count in non_empty: @@ -166,7 +171,8 @@ def compute_bounding_box(self) -> BoundingBox: bboxes = [] collars = self.properties.copy() - collars["_hole_index"] = collars["hole_id"].astype(self.hole_id_dtype).cat.codes + hole_indices = {hole_id: index for index, hole_id in enumerate(self.hole_id_dtype.categories)} + collars["_hole_index"] = collars["hole_id"].astype(object).map(hole_indices) collars_by_index = collars.set_index("_hole_index") for chunk in self.holes.itertuples(index=False): offset = int(chunk.offset) @@ -315,7 +321,7 @@ class DownholeLocation(SchemaModel): attributes: Annotated[Attributes, SchemaLocation("attributes"), DataLocation("attributes")] async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: - """Return collars with a categorical ``hole_id`` aligned to hole-index codes.""" + """Return collars with a categorical ``hole_id`` column.""" parts = [ await self.hole_id.to_dataframe(fb=fb), await self.coordinates.to_dataframe(fb=fb), @@ -487,29 +493,29 @@ def _validate_chunk_ranges(holes: HoleChunks, table_length: int) -> None: required = {"hole_index", "offset", "count"} if missing := required - set(holes.columns): raise ObjectValidationError(f"Hole chunks are missing columns: {sorted(missing)}") - if holes["hole_index"].astype(int).duplicated().any(): - raise ObjectValidationError("Each hole_index may occur only once in a holes table") 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") - expected_offset = 0 - for offset, count in sorted(zip(offsets[counts > 0], counts[counts > 0], strict=True)): - 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") async def _table_by_hole( table: DownholeDistanceTable | DownholeIntervalTable, data: pd.DataFrame, *, fb: IFeedback ) -> dict[str, pd.DataFrame]: root = table._context.root_model - collar_ids = await root.location.hole_id.to_dataframe(fb=fb) - categories = collar_ids.iloc[:, 0] - result: dict[str, pd.DataFrame] = {} + lookup = await root.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 table.holes.to_dataframe(fb=fb)).itertuples(index=False): - hole_id = str(categories.cat.categories[int(chunk.hole_index)]) - result[hole_id] = data.iloc[int(chunk.offset) : int(chunk.offset) + int(chunk.count)].reset_index(drop=True) - return result + 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() + } diff --git a/packages/evo-objects/tests/typed/helpers.py b/packages/evo-objects/tests/typed/helpers.py index 8860c6ea..4bd07589 100644 --- a/packages/evo-objects/tests/typed/helpers.py +++ b/packages/evo-objects/tests/typed/helpers.py @@ -112,8 +112,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_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index b5d39372..ad2dc8ae 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -535,19 +535,38 @@ def test_attributes_length_raises(self): desurvey=None, ) - def test_collection_chunk_overlap_gap_and_invalid_index_raise(self): + def test_location_chunks_require_each_creation_code_and_exact_path_coverage(self): base = _make_example_data(collections=[]) - table = pd.DataFrame({"distance": [0.0, 1.0]}) for holes in ( - pd.DataFrame({"hole_index": [0, 1], "offset": [0, 0], "count": [1, 1]}), - pd.DataFrame({"hole_index": [0], "offset": [1], "count": [1]}), - pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), + 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": [0, 1], "offset": [0, 5], "count": [4, 2]}), + pd.DataFrame({"hole_index": [0, 1], "offset": [0, 3], "count": [4, 3]}), ): - with self.assertRaises(ObjectValidationError): - dataclasses.replace( - base, - collections=[DistanceCollection(name="invalid", holes=holes, table=table)], - ) + 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, + ) + dataclasses.replace(base, collections=[collection]) + + with self.assertRaises(ObjectValidationError): + dataclasses.replace( + base, + collections=[ + DistanceCollection( + name="invalid", + holes=pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), + table=table, + ) + ], + ) async def test_collection_allows_zero_row_hole(self): base = _make_example_data(collections=[]) @@ -562,7 +581,7 @@ async def test_collection_allows_zero_row_hole(self): 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_rejects_duplicate_hole_chunks_and_replaces_in_place(self): + 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( @@ -577,12 +596,11 @@ async def test_collection_add_rejects_duplicate_hole_chunks_and_replaces_in_plac 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) - invalid = dataclasses.replace( + repeated = dataclasses.replace( collection, holes=pd.DataFrame({"hole_index": [0, 0], "offset": [0, 1], "count": [1, 0]}), ) - with self.assertRaises(ObjectValidationError): - await result.collections.add(invalid, replace=True) + await result.collections.add(repeated, replace=True) async def test_collection_replacement_preserves_position_and_persists_after_update(self): first = DistanceCollection( @@ -625,6 +643,35 @@ async def test_location_path_and_collection_reads(self): self.assertListEqual(list(by_hole), ["H001"]) self.assertEqual(len(by_hole["H001"]), 4) + 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]}), + ) + 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(): From 73945ff840b3af469d754fc3b01a980a22651e2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 10:21:57 +0100 Subject: [PATCH 13/30] feat(objects): Require explicit nullable domain choices for units and desurvey settings. --- .../create-downhole-collection.ipynb | 3 ++- .../evo/objects/typed/downhole_collection.py | 8 +++--- .../tests/typed/test_downhole_collection.py | 25 +++++++++++++++++++ 3 files changed, 31 insertions(+), 5 deletions(-) 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/src/evo/objects/typed/downhole_collection.py b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py index 873330c9..b739ddcf 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -72,8 +72,8 @@ class DistanceCollection: name: str holes: HoleChunks table: Depths + unit: str | None collection_type: str = "distance" - unit: str | None = None @dataclass @@ -81,8 +81,8 @@ class IntervalCollection: name: str holes: HoleChunks table: Intervals + unit: str | None collection_type: str = "interval" - unit: str | None = None DownholeCollectionEntry: TypeAlias = DistanceCollection | IntervalCollection @@ -117,8 +117,8 @@ class DownholeCollectionData(BaseSpatialObjectData): properties: HoleProperties attributes: HoleAttributes | None collections: list[DownholeCollectionEntry] - distance_unit: str | None = None - desurvey: str | None = None + distance_unit: str | None + desurvey: str | None @property def hole_id_dtype(self) -> pd.CategoricalDtype: diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index ad2dc8ae..1a66aa8a 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -13,6 +13,7 @@ import contextlib import dataclasses +import inspect import math import uuid from datetime import date @@ -75,6 +76,7 @@ def _make_example_data( "attr_num": [1.1, 2.2, 3.3, 4.4], } ), + unit=None, ) holes = pd.DataFrame( @@ -115,6 +117,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) @@ -236,6 +249,7 @@ async def test_interval_only_collection_uses_attribute_unit_when_not_explicit(se 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(): @@ -291,6 +305,7 @@ async def test_interval_collection_round_trips_all_attribute_types(self): name="intervals", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [2]}), table=table, + unit=None, ) expected_types = { "scalar": "scalar", @@ -322,6 +337,7 @@ async def test_interval_collection_requires_from_and_to_columns(self, _name, tab 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])) @@ -355,6 +371,7 @@ async def test_attribute_descriptions_round_trip_to_dataframe_metadata(self): 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( @@ -553,6 +570,7 @@ def test_collection_chunk_ranges_are_not_required_to_partition_the_table(self): name="repeated", holes=pd.DataFrame({"hole_index": [0, 0], "offset": [0, 2], "count": [1, 1]}), table=table, + unit=None, ) dataclasses.replace(base, collections=[collection]) @@ -564,6 +582,7 @@ def test_collection_chunk_ranges_are_not_required_to_partition_the_table(self): name="invalid", holes=pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), table=table, + unit=None, ) ], ) @@ -574,6 +593,7 @@ async def test_collection_allows_zero_row_hole(self): 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(): @@ -588,6 +608,7 @@ async def test_collection_add_allows_repeated_hole_chunks_and_replaces_in_place( 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]})) @@ -607,16 +628,19 @@ async def test_collection_replacement_preserves_position_and_persists_after_upda 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(): @@ -662,6 +686,7 @@ async def test_repeated_collection_chunks_are_concatenated_in_table_order(self): 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( From 0efb0067a1873380efab1520bf40bf37ef4119ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 11:31:39 +0100 Subject: [PATCH 14/30] feat(objects): enhance to_dataframe methods to accept attribute selection keys --- .../evo/objects/typed/downhole_collection.py | 30 ++++++++------- .../tests/typed/test_downhole_collection.py | 37 +++++++++++++++++++ 2 files changed, 54 insertions(+), 13 deletions(-) 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 b739ddcf..b2566a3c 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -320,20 +320,20 @@ class DownholeLocation(SchemaModel): coordinates: Annotated[CollarCoordinates, SchemaLocation("coordinates"), DataLocation("properties")] attributes: Annotated[Attributes, SchemaLocation("attributes"), DataLocation("attributes")] - async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: - """Return collars with a categorical ``hole_id`` column.""" + async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return collars with a categorical ``hole_id`` column 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(fb=fb)) + parts.append(await self.attributes.to_dataframe(*keys, fb=fb)) return pd.concat(parts, axis=1) - async def path_to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: + 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(fb=fb) + return await self.path.to_dataframe(*keys, fb=fb) class _Distances(DataTable): @@ -367,11 +367,13 @@ class DistanceTable(SchemaModel): class DownholeDistanceTable(DistanceTable): holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] - async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: - return await self.distance.to_dataframe(fb=fb) + async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: + """Return collection values and selected attributes.""" + return await self.distance.to_dataframe(*keys, fb=fb) - async def to_dataframe_by_hole(self, fb: IFeedback = NoFeedback) -> dict[str, pd.DataFrame]: - return await _table_by_hole(self, await self.to_dataframe(fb=fb), fb=fb) + 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 _table_by_hole(self, await self.to_dataframe(*keys, fb=fb), fb=fb) class _Intervals(DataTable): @@ -401,11 +403,13 @@ class DownholeIntervalTable(SchemaModel): from_to: Annotated[IntervalTableFromTo, SchemaLocation("from_to"), DataLocation("table")] holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] - async def to_dataframe(self, fb: IFeedback = NoFeedback) -> pd.DataFrame: - return await self.from_to.to_dataframe(fb=fb) + 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) - async def to_dataframe_by_hole(self, fb: IFeedback = NoFeedback) -> dict[str, pd.DataFrame]: - return await _table_by_hole(self, await self.to_dataframe(fb=fb), fb=fb) + 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 _table_by_hole(self, await self.to_dataframe(*keys, fb=fb), fb=fb) class DownholeCollectionTables(SchemaList[DownholeDistanceTable | DownholeIntervalTable]): diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 1a66aa8a..cb69f46e 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -667,6 +667,43 @@ async def test_location_path_and_collection_reads(self): self.assertListEqual(list(by_hole), ["H001"]) self.assertEqual(len(by_hole["H001"]), 4) + 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()) From 346325dbcd5a5789dc391cbb2e8475bf3ea420ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 11:49:06 +0100 Subject: [PATCH 15/30] feat(objects): require explicit units for downhole distance and interval collections. --- .../evo/objects/typed/downhole_collection.py | 46 ++++--------------- .../tests/typed/test_downhole_collection.py | 37 +++++++++------ 2 files changed, 32 insertions(+), 51 deletions(-) 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 b2566a3c..e58a713a 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -11,7 +11,7 @@ from __future__ import annotations -from dataclasses import dataclass, replace +from dataclasses import dataclass from typing import Annotated, Any, ClassVar, TypeAlias import numpy as np @@ -25,10 +25,7 @@ from evo.objects.typed._data import DataTable, DataTableAndAttributes from evo.objects.typed._downhole import HoleIdCategory from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel -from evo.objects.typed.attributes import ( - AttributeDescription, - Attributes, -) +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 @@ -58,11 +55,6 @@ 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_descriptions': {: }, ...} Depths: TypeAlias = pd.DataFrame # [ distance | ] Intervals: TypeAlias = pd.DataFrame # [ from | to | ] @@ -345,18 +337,6 @@ 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) - unit = data.attrs.get("unit") - attr_desc: AttributeDescription = data.attrs.get("attribute_descriptions", {}).get("distance") - if unit is None and attr_desc is not None: - unit = attr_desc.unit - if unit is not None: - # "unit" can be missing, but it must not be `None` - result["unit"] = unit - return result - class DistanceTable(SchemaModel): name: Annotated[str, SchemaLocation("name"), DataLocation("name")] @@ -385,17 +365,6 @@ class IntervalTableFromTo(DataTableAndAttributes): _table: Annotated[_Intervals, SchemaLocation("intervals.start_and_end"), 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) - unit = data.attrs.get("unit") - description = data.attrs.get("attribute_descriptions", {}).get("from") - if unit is None and description is not None: - unit = description.unit - if unit is not None: - result["unit"] = unit - return result - class DownholeIntervalTable(SchemaModel): name: Annotated[str, SchemaLocation("name"), DataLocation("name")] @@ -424,12 +393,13 @@ async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: result = [] for collection in data: model = DownholeIntervalTable if isinstance(collection, IntervalCollection) else DownholeDistanceTable - table = collection.table - table = table.copy() - table.attrs = dict(table.attrs) + schema = await model._data_to_schema(collection, context) if collection.unit is not None: - table.attrs["unit"] = collection.unit - result.append(await model._data_to_schema(replace(collection, table=table), context)) + if isinstance(collection, IntervalCollection): + schema["from_to"]["unit"] = collection.unit + else: + schema["distance"]["unit"] = collection.unit + result.append(schema) return result def get(self, name: str) -> DownholeDistanceTable | DownholeIntervalTable | None: diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index cb69f46e..8d4dfac9 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -242,7 +242,7 @@ async def test_mixed_collections_round_trip_and_mutation(self): 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_only_collection_uses_attribute_unit_when_not_explicit(self): + 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( @@ -252,24 +252,23 @@ async def test_interval_only_collection_uses_attribute_unit_when_not_explicit(se unit=None, ) data = _make_example_data(collections=[interval]) - with self._mock_geoscience_objects(): + 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.assertEqual(collection.from_to.unit, "ft") + 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", "ft", "m"), - ("metadata", None, "ft", "ft"), - ("omitted", None, None, None), + ("explicit", "m", "m"), + ("omitted", None, None), ] ) - async def test_interval_collection_unit_precedence(self, _name, explicit_unit, metadata_unit, expected_unit): + async def test_interval_collection_uses_explicit_unit_only(self, _name, explicit_unit, expected_unit): table = pd.DataFrame({"from": [0.0], "to": [1.0]}) - if metadata_unit is not None: - table.attrs["attribute_descriptions"] = {"from": AttributeDescription(unit=metadata_unit)} collection = IntervalCollection( name="intervals", holes=pd.DataFrame({"hole_index": [0], "offset": [0], "count": [1]}), @@ -342,20 +341,32 @@ async def test_interval_collection_requires_from_and_to_columns(self, _name, tab with self._mock_geoscience_objects(), self.assertRaises(ObjectValidationError): await DownholeCollection.create(context=self.context, data=_make_example_data(collections=[collection])) - async def test_explicit_collection_unit_overrides_dataframe_metadata(self): + @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="m", + unit=collection_unit, ) - with self._mock_geoscience_objects(): + 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, "m") + 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]}) From e4b6ae761abadd9ad4ef1ecf4eff5fef1c53dcd1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 11:59:42 +0100 Subject: [PATCH 16/30] feat(objects): moved method to parent class --- .../src/evo/objects/typed/downhole_collection.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 e58a713a..88a4dcca 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -343,14 +343,14 @@ class DistanceTable(SchemaModel): collection_type: Annotated[str, SchemaLocation("collection_type"), DataLocation("collection_type")] 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 DownholeDistanceTable(DistanceTable): holes: Annotated[HoleChunksTable, SchemaLocation("holes"), DataLocation("holes")] - async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: - """Return collection values and selected attributes.""" - return await self.distance.to_dataframe(*keys, fb=fb) - 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 _table_by_hole(self, await self.to_dataframe(*keys, fb=fb), fb=fb) From ffc4d3dd790c4c998d0622c7d7c5ac7537a4ff81 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 12:05:26 +0100 Subject: [PATCH 17/30] feat(objects): update interval table to use DepthIntervalsTable --- .../src/evo/objects/typed/downhole_collection.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) 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 88a4dcca..1150e6f8 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -23,7 +23,7 @@ from evo.common.utils import NoFeedback from evo.objects import SchemaVersion 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 Attributes from evo.objects.typed.exceptions import ObjectValidationError @@ -32,7 +32,6 @@ from evo.objects.utils.table_formats import ( DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, - FLOAT_ARRAY_2, FLOAT_ARRAY_3, KnownTableFormat, ) @@ -356,13 +355,8 @@ async def to_dataframe_by_hole(self, *keys: str, fb: IFeedback = NoFeedback) -> return await _table_by_hole(self, await self.to_dataframe(*keys, fb=fb), fb=fb) -class _Intervals(DataTable): - table_format: ClassVar[KnownTableFormat] = FLOAT_ARRAY_2 - data_columns: ClassVar[list[str]] = ["from", "to"] - - class IntervalTableFromTo(DataTableAndAttributes): - _table: Annotated[_Intervals, SchemaLocation("intervals.start_and_end"), DataLocation("")] + _table: Annotated[DepthIntervalsTable, SchemaLocation("intervals.start_and_end"), DataLocation("")] unit: Annotated[str | None, SchemaLocation("unit")] From 97c220d2f51dabb1b4e6a2b4dc270cbdbe02025a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 12:22:21 +0100 Subject: [PATCH 18/30] feat(objects): enforce unique collection names and handle legacy duplicates --- .../evo/objects/typed/downhole_collection.py | 37 +++++++++++++++---- .../tests/typed/test_downhole_collection.py | 27 +++++++++++++- 2 files changed, 56 insertions(+), 8 deletions(-) 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 1150e6f8..42f4a17b 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 @@ -125,6 +126,10 @@ def __post_init__(self): assert self.attributes is None or len(self.holes) == len(self.attributes) + names = [collection.name for collection in self.collections] + if len(names) != len(set(names)): + raise ObjectValidationError("Collection names must be unique") + self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) for collection in self.collections: table = collection.table @@ -396,19 +401,37 @@ async def _data_to_schema(cls, data: Any, context: IContext) -> list[Any]: 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 next((collection for collection in self if collection.name == name), 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 self.get(name) is not None + return isinstance(name, str) and bool(self._indices_for_name(name)) def names(self) -> list[str]: return [collection.name for collection in self] async def add(self, collection: DownholeCollectionEntry, *, replace: bool = False) -> None: - existing = self.names() - if collection.name in existing and not replace: - raise ValueError(f"Collection '{collection.name}' already exists") + 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") location_holes = await self._context.root_model.location.holes.to_dataframe() valid_indices = set(location_holes["hole_index"].astype(int)) if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): @@ -416,8 +439,8 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals table = collection.table _validate_chunk_ranges(collection.holes, len(table)) schema = await self._data_to_schema([collection], self._obj) - if collection.name in existing: - self._document[existing.index(collection.name)] = schema[0] + if existing_indices: + self._document[existing_indices[0]] = schema[0] else: self._document.append(schema[0]) diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 8d4dfac9..b41b4b31 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -12,10 +12,12 @@ 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 @@ -623,7 +625,7 @@ async def test_collection_add_allows_repeated_hole_chunks_and_replaces_in_place( ) await result.collections.add(collection) replacement = dataclasses.replace(collection, table=pd.DataFrame({"distance": [1.0]})) - with self.assertRaises(ValueError): + with self.assertRaises(ObjectValidationError): await result.collections.add(replacement) await result.collections.add(replacement, replace=True) self.assertEqual(result.collections.names(), ["measurements"]) @@ -634,6 +636,29 @@ async def test_collection_add_allows_repeated_hole_chunks_and_replaces_in_place( ) 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", From fca9d149c20671e4b6c557af4c95bb1b5a4cd280 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 12:41:03 +0100 Subject: [PATCH 19/30] feat(objects): enhance prefetch_collections method with proper names for keywords arguments --- .../evo/objects/typed/downhole_collection.py | 10 ++++++++-- .../tests/typed/test_downhole_collection.py | 20 ++++++++++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) 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 42f4a17b..2023ff27 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -465,7 +465,13 @@ class DownholeCollection(BaseSpatialObject): type: ClassVar[Annotated[str, SchemaLocation("type")]] = "downhole" - async def prefetch_collections(self, *names: str, include_location: bool = True, **kwargs: Any) -> None: + async def prefetch_collections( + self, + *names: str, + include_location: bool = True, + max_concurrent: int = 100, + fb: IFeedback = NoFeedback, + ) -> None: """Prefetch data referenced by named collections and optionally location data.""" from evo.objects.typed._prefetch import collect_data_ids @@ -477,7 +483,7 @@ async def prefetch_collections(self, *names: str, include_location: bool = True, if collection is None: raise KeyError(f"Unknown collection '{name}'") documents.append(collection.as_dict()) - await self.prefetch(data_ids=collect_data_ids(documents), **kwargs) + await self.prefetch(data_ids=collect_data_ids(documents), max_concurrent=max_concurrent, fb=fb) def _validate_chunk_ranges(holes: HoleChunks, table_length: int) -> None: diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index b41b4b31..725beb87 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -19,7 +19,7 @@ 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 @@ -28,6 +28,7 @@ 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 @@ -703,6 +704,23 @@ async def test_location_path_and_collection_reads(self): 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, 100) + self.assertIs(parameters["fb"].default, NoFeedback) + + 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", From c82166ded2d2735d243cf2396fc3a64e90828946 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 12:53:24 +0100 Subject: [PATCH 20/30] feat(objects): enhance hole chunk validation with valid indices and remove redundant method --- .../evo/objects/typed/downhole_collection.py | 35 ++++++++----------- 1 file changed, 15 insertions(+), 20 deletions(-) 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 2023ff27..e05f2382 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -130,22 +130,27 @@ def __post_init__(self): if len(names) != len(set(names)): raise ObjectValidationError("Collection names must be unique") - self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) + valid_indices = set(range(len(self.hole_id_dtype.categories))) + self._validate_hole_chunks(self.holes, len(self.path), valid_indices=valid_indices, require_coverage=True) for collection in self.collections: table = collection.table - self._validate_hole_chunks(collection.holes, len(table), require_coverage=False) + self._validate_hole_chunks( + collection.holes, len(table), valid_indices=valid_indices, require_coverage=False + ) - def _validate_hole_chunks(self, holes: HoleChunks, table_length: int, *, require_coverage: bool) -> None: + @staticmethod + def _validate_hole_chunks( + holes: HoleChunks, table_length: int, *, valid_indices: set[int], require_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) - valid = set(range(len(self.hole_id_dtype.categories))) - if not set(indices).issubset(valid): - raise ObjectValidationError("hole_index must be a code in properties['hole_id'] categorical dtype") + if not set(indices).issubset(valid_indices): + raise ObjectValidationError("hole_index must reference a valid hole") if require_coverage and indices.duplicated().any(): raise ObjectValidationError("Each hole_index may occur only once in a holes table") - if require_coverage and set(indices) != valid: + if require_coverage and set(indices) != valid_indices: raise ObjectValidationError("Location holes must cover every hole_id categorical code exactly once") offsets = holes["offset"].astype(int) counts = holes["count"].astype(int) @@ -434,10 +439,10 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals raise ObjectValidationError(f"Collection '{collection.name}' already exists") location_holes = await self._context.root_model.location.holes.to_dataframe() valid_indices = set(location_holes["hole_index"].astype(int)) - if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): - raise ObjectValidationError("Collection hole_index is not present in the location holes table") table = collection.table - _validate_chunk_ranges(collection.holes, len(table)) + DownholeCollectionData._validate_hole_chunks( + collection.holes, len(table), valid_indices=valid_indices, require_coverage=False + ) schema = await self._data_to_schema([collection], self._obj) if existing_indices: self._document[existing_indices[0]] = schema[0] @@ -486,16 +491,6 @@ async def prefetch_collections( await self.prefetch(data_ids=collect_data_ids(documents), max_concurrent=max_concurrent, fb=fb) -def _validate_chunk_ranges(holes: HoleChunks, table_length: int) -> None: - required = {"hole_index", "offset", "count"} - if missing := required - set(holes.columns): - raise ObjectValidationError(f"Hole chunks are missing columns: {sorted(missing)}") - 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") - - async def _table_by_hole( table: DownholeDistanceTable | DownholeIntervalTable, data: pd.DataFrame, *, fb: IFeedback ) -> dict[str, pd.DataFrame]: From d96d54b3fe0e90352d65bceb9aa118450e5a72df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Mon, 10 Aug 2026 13:06:26 +0100 Subject: [PATCH 21/30] feat(objects): refactor DownholeCollection classes to share common behavior and improve hole data handling --- .../evo/objects/typed/downhole_collection.py | 67 +++++++++++-------- 1 file changed, 39 insertions(+), 28 deletions(-) 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 e05f2382..87c21337 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -357,12 +357,45 @@ async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataF return await self.distance.to_dataframe(*keys, fb=fb) -class DownholeDistanceTable(DistanceTable): +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 containing DownholeCollection's 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 _table_by_hole(self, await self.to_dataframe(*keys, fb=fb), fb=fb) + return await self._table_by_hole(await self.to_dataframe(*keys, fb=fb), fb=fb) class IntervalTableFromTo(DataTableAndAttributes): @@ -370,7 +403,7 @@ class IntervalTableFromTo(DataTableAndAttributes): unit: Annotated[str | None, SchemaLocation("unit")] -class DownholeIntervalTable(SchemaModel): +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")] @@ -382,10 +415,10 @@ async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataF 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 _table_by_hole(self, await self.to_dataframe(*keys, fb=fb), fb=fb) + return await self._table_by_hole(await self.to_dataframe(*keys, fb=fb), fb=fb) -class DownholeCollectionTables(SchemaList[DownholeDistanceTable | DownholeIntervalTable]): +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 @@ -437,7 +470,7 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals ) if existing_indices and not replace: raise ObjectValidationError(f"Collection '{collection.name}' already exists") - location_holes = await self._context.root_model.location.holes.to_dataframe() + location_holes = await self._downhole_collection.location.holes.to_dataframe() valid_indices = set(location_holes["hole_index"].astype(int)) table = collection.table DownholeCollectionData._validate_hole_chunks( @@ -489,25 +522,3 @@ async def prefetch_collections( 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) - - -async def _table_by_hole( - table: DownholeDistanceTable | DownholeIntervalTable, data: pd.DataFrame, *, fb: IFeedback -) -> dict[str, pd.DataFrame]: - root = table._context.root_model - lookup = await root.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 table.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() - } From 21d2e6b45ad930ea66cf69f7e51dc56cb9ab7bff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 11 Aug 2026 07:48:56 +0100 Subject: [PATCH 22/30] feat(objects): moves import to module-level --- .../evo-objects/src/evo/objects/typed/downhole_collection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 87c21337..3f26de6b 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -26,6 +26,7 @@ from evo.objects.typed._data import DataTable, DataTableAndAttributes from evo.objects.typed._downhole import DepthIntervalsTable, HoleIdCategory from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel +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 @@ -511,7 +512,6 @@ async def prefetch_collections( fb: IFeedback = NoFeedback, ) -> None: """Prefetch data referenced by named collections and optionally location data.""" - from evo.objects.typed._prefetch import collect_data_ids documents = [] if include_location: From dbbd9985d05efc9bed1fb7fb396aff775540f7c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 11 Aug 2026 07:55:35 +0100 Subject: [PATCH 23/30] feat(objects): lower the default max concurrency for pre-fetching --- packages/evo-objects/src/evo/objects/typed/_prefetch.py | 2 +- packages/evo-objects/src/evo/objects/typed/base.py | 2 +- .../evo-objects/src/evo/objects/typed/downhole_collection.py | 2 +- packages/evo-objects/tests/typed/test_downhole_collection.py | 5 ++++- 4 files changed, 7 insertions(+), 4 deletions(-) diff --git a/packages/evo-objects/src/evo/objects/typed/_prefetch.py b/packages/evo-objects/src/evo/objects/typed/_prefetch.py index bb8b0c83..d0edc839 100644 --- a/packages/evo-objects/src/evo/objects/typed/_prefetch.py +++ b/packages/evo-objects/src/evo/objects/typed/_prefetch.py @@ -46,7 +46,7 @@ async def prefetch_object_data( obj: DownloadedObject, *, data_ids: Sequence[str] | None = None, - max_concurrent: int = 100, + max_concurrent: int = 8, fb: IFeedback = NoFeedback, ) -> None: """Warm cache entries referenced by an object, downloading each ID at most once.""" diff --git a/packages/evo-objects/src/evo/objects/typed/base.py b/packages/evo-objects/src/evo/objects/typed/base.py index 55675fb3..534365df 100644 --- a/packages/evo-objects/src/evo/objects/typed/base.py +++ b/packages/evo-objects/src/evo/objects/typed/base.py @@ -423,7 +423,7 @@ async def prefetch( self, *, data_ids: Sequence[str] | None = None, - max_concurrent: int = 100, + max_concurrent: int = 8, fb: IFeedback = NoFeedback, ) -> None: """Warm cached data files referenced by this object.""" 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 3f26de6b..83579cd8 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -508,7 +508,7 @@ async def prefetch_collections( self, *names: str, include_location: bool = True, - max_concurrent: int = 100, + max_concurrent: int = 8, fb: IFeedback = NoFeedback, ) -> None: """Prefetch data referenced by named collections and optionally location data.""" diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 725beb87..7ee0253b 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -710,9 +710,12 @@ async def test_prefetch_collections_exposes_and_forwards_explicit_options(self): parameters = inspect.signature(DownholeCollection.prefetch_collections).parameters self.assertNotIn("kwargs", parameters) - self.assertEqual(parameters["max_concurrent"].default, 100) + 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) From 4ce86844afff6be0df242cb91043bf59b47c941d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 11 Aug 2026 08:08:51 +0100 Subject: [PATCH 24/30] feat(objects): enforces stricted unit type checking --- .../src/evo/objects/typed/__init__.py | 2 ++ .../src/evo/objects/typed/attributes.py | 19 +++++++++++++++---- .../tests/typed/test_attributes.py | 9 +++++++++ 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/packages/evo-objects/src/evo/objects/typed/__init__.py b/packages/evo-objects/src/evo/objects/typed/__init__.py index 6f591b3b..501cce39 100644 --- a/packages/evo-objects/src/evo/objects/typed/__init__.py +++ b/packages/evo-objects/src/evo/objects/typed/__init__.py @@ -18,6 +18,7 @@ BlockModelAttributes, BlockModelPendingAttribute, PendingAttribute, + Unit, ) from .base import BaseObject, object_from_path, object_from_reference, object_from_uuid from .block_model_ref import ( @@ -119,6 +120,7 @@ "SpheroidalStructure", "Tensor3DGrid", "Tensor3DGridData", + "Unit", "Variogram", "VariogramCurveData", "VariogramData", diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 96566dd3..447be3b3 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 @@ -47,6 +47,7 @@ "BlockModelAttributes", "BlockModelPendingAttribute", "PendingAttribute", + "Unit", ] @@ -54,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. @@ -90,14 +98,17 @@ def _infer_attribute_type_from_series(series: pd.Series) -> str: class AttributeDescription: discipline: str = "" type: str = "" - unit: Any | 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 __post_init__(self) -> None: - if self.unit is not None and not isinstance(self.unit, str): - self.unit = str(self.unit.value) + 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] = { diff --git a/packages/evo-objects/tests/typed/test_attributes.py b/packages/evo-objects/tests/typed/test_attributes.py index a0ad59da..1a33c32c 100644 --- a/packages/evo-objects/tests/typed/test_attributes.py +++ b/packages/evo-objects/tests/typed/test_attributes.py @@ -11,6 +11,7 @@ from __future__ import annotations +from typing import Any, cast from unittest import TestCase import pandas as pd @@ -80,9 +81,17 @@ 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())) From 1a7a2d609b599bda2628b98a1d252e1107fb2950 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 11 Aug 2026 08:44:12 +0100 Subject: [PATCH 25/30] fix(objects): treat downhole hole_index as an opaque lookup key, hole_index is the integer key in the object's hole_id lookup table, not a positional categorical code. DownholeCollectionData and DownholeCollection no longer assume it is zero-based or contiguous. --- packages/evo-objects/README.md | 9 +- .../evo/objects/typed/downhole_collection.py | 143 ++++++++++++------ .../src/evo/objects/utils/downhole.py | 45 +++--- .../evo-objects/tests/test_downhole_utils.py | 90 +++++------ packages/evo-objects/tests/typed/helpers.py | 8 +- .../tests/typed/test_downhole_collection.py | 57 +++++-- 6 files changed, 225 insertions(+), 127 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index c1c3414b..720f564f 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -72,8 +72,10 @@ Check out the other methods on the `ObjectAPIClient` for more details on how to ### 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 chunk tables always use a zero-based `hole_index`: it is the code in the -shared categorical `properties["hole_id"]` dtype, not the row position of a collar. +distance tables, interval tables, or both. `hole_index` is an integer key in the persisted `hole_id` lookup table; it is +not a row position and need not be zero-based or contiguous. `hole_chunks_from_ids()` is a convenience helper that emits +dense zero-based keys for new objects. Pass `hole_indices={"DH-01": 1}` when writing against an existing or explicit +lookup table; the helper creates chunks only for IDs present in its input. ```python import pandas as pd @@ -82,12 +84,11 @@ 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 -hole_dtype = pd.CategoricalDtype(categories=["DH-01"]) intervals = pd.DataFrame({"from": [0.0], "to": [1.5], "lithology": ["sandstone"]}) collections = [ IntervalCollection( name="geology", - holes=hole_chunks_from_ids(pd.Series(["DH-01"]), dtype=hole_dtype), + holes=hole_chunks_from_ids(pd.Series(["DH-01"])), table=intervals, unit="m", # Explicit collection units override DataFrame metadata. ) 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 83579cd8..c1e18f2f 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -23,6 +23,7 @@ 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 DepthIntervalsTable, HoleIdCategory from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel @@ -35,6 +36,8 @@ DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, FLOAT_ARRAY_3, + INTEGER_ARRAY_1_INT32, + LOOKUP_TABLE_INT32, KnownTableFormat, ) @@ -83,14 +86,12 @@ class IntervalCollection: @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_index, offset, count. For object creation, ``hole_index`` is the zero-based categorical code - in ``properties.hole_id``. - :param properties: DataFrame for the properties of the holes. Its categorical ``hole_id`` codes are the lookup - keys referenced by creation-time ``holes`` tables. + 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. Mandatory columns: hole_id, final, target, current, x, y, z :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 @@ -113,14 +114,6 @@ class DownholeCollectionData(BaseSpatialObjectData): distance_unit: str | None desurvey: str | None - @property - def hole_id_dtype(self) -> pd.CategoricalDtype: - """The categorical dtype used to encode hole indices when creating this object.""" - hole_ids = self.properties["hole_id"] - if isinstance(hole_ids.dtype, pd.CategoricalDtype): - return hole_ids.dtype - return pd.CategoricalDtype(categories=sorted(hole_ids.dropna().unique())) - 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") @@ -131,28 +124,16 @@ def __post_init__(self): if len(names) != len(set(names)): raise ObjectValidationError("Collection names must be unique") - valid_indices = set(range(len(self.hole_id_dtype.categories))) - self._validate_hole_chunks(self.holes, len(self.path), valid_indices=valid_indices, require_coverage=True) + self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) for collection in self.collections: table = collection.table - self._validate_hole_chunks( - collection.holes, len(table), valid_indices=valid_indices, require_coverage=False - ) + self._validate_hole_chunks(collection.holes, len(table), require_coverage=False) @staticmethod - def _validate_hole_chunks( - holes: HoleChunks, table_length: int, *, valid_indices: set[int], require_coverage: bool - ) -> None: + def _validate_hole_chunks(holes: HoleChunks, table_length: int, *, require_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") - if require_coverage and indices.duplicated().any(): - raise ObjectValidationError("Each hole_index may occur only once in a holes table") - if require_coverage and set(indices) != valid_indices: - raise ObjectValidationError("Location holes must cover every hole_id categorical code 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(): @@ -169,21 +150,62 @@ def _validate_hole_chunks( if expected_offset != table_length: raise ObjectValidationError("Hole chunk ranges must cover the associated table exactly once") - def compute_bounding_box(self) -> BoundingBox: - bboxes = [] + def _collars_by_hole_index(self) -> dict[int, tuple[float, float, float]] | None: + """Map hole lookup keys to collar coordinates, or ``None`` when the mapping is not known. - collars = self.properties.copy() - hole_indices = {hole_id: index for index, hole_id in enumerate(self.hole_id_dtype.categories)} - collars["_hole_index"] = collars["hole_id"].astype(object).map(hole_indices) - collars_by_index = collars.set_index("_hole_index") + Keys are taken from an explicit ``properties["hole_index"]`` column when present. Otherwise the SDK creation + convention (dense zero-based keys in sorted hole-ID order) is used, and only when it accounts for every key + referenced by ``holes``. + """ + hole_ids = self.properties["hole_id"].astype(str) + if "hole_index" in self.properties: + keys = self.properties["hole_index"].astype(int) + else: + dense = {hole_id: index for index, hole_id in enumerate(sorted(hole_ids.unique()))} + keys = hole_ids.map(dense).astype(int) + coordinates = self.properties[_COORDINATE_COLUMNS].astype(float) + collars = { + int(key): coordinate + for key, coordinate in zip(keys, coordinates.itertuples(index=False, name=None), strict=True) + } + referenced = {int(chunk.hole_index) for chunk in self.holes.itertuples(index=False)} + return collars if referenced <= set(collars) else None + + def compute_bounding_box(self) -> BoundingBox: + collars_by_hole_index = self._collars_by_hole_index() + if collars_by_hole_index is not None: + bboxes = [ + self._compute_hole_bounding_box( + self.path[int(chunk.offset) : int(chunk.offset) + int(chunk.count)], + collars_by_hole_index[int(chunk.hole_index)], + ) + for chunk in self.holes.itertuples(index=False) + if int(chunk.count) + ] + if bboxes: + return BoundingBox.combine(bboxes) + + # The collar for each chunk cannot be identified, so fall back to an envelope that contains every hole. + collars = BoundingBox.from_points(self.properties[_COORDINATE_COLUMNS].astype(float).to_numpy()) + relative_bboxes = [] for chunk in self.holes.itertuples(index=False): offset = int(chunk.offset) count = int(chunk.count) - collar = tuple(collars_by_index.loc[int(chunk.hole_index), _COORDINATE_COLUMNS]) - path_table = self.path[offset : offset + count] - bboxes.append(self._compute_hole_bounding_box(path_table, collar)) - - return BoundingBox.combine(bboxes) + if count: + relative_bboxes.append( + self._compute_hole_bounding_box(self.path[offset : offset + count], (0.0, 0.0, 0.0)) + ) + if not relative_bboxes: + return collars + relative = BoundingBox.combine(relative_bboxes) + return BoundingBox( + min_x=collars.min_x + relative.min_x, + max_x=collars.max_x + relative.max_x, + min_y=collars.min_y + relative.min_y, + max_y=collars.max_y + relative.max_y, + min_z=collars.min_z + relative.min_z, + max_z=collars.max_z + relative.max_z, + ) @staticmethod def _compute_bounding_box_np( @@ -314,8 +336,39 @@ 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. + + ``properties`` may provide a ``hole_index`` column to persist specific lookup keys, for example when rewriting a + downloaded object. When it is absent, dense zero-based keys are generated in sorted hole-ID order as an SDK + creation convenience. + """ + + @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") + if "hole_index" in data: + keys = data["hole_index"].astype("int32") + if keys.duplicated().any(): + raise ObjectValidationError("properties hole_index values must be unique") + else: + dense = {hole_id: index for index, hole_id in enumerate(sorted(hole_ids.unique()))} + keys = hole_ids.map(dense).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")] @@ -323,7 +376,7 @@ class DownholeLocation(SchemaModel): attributes: Annotated[Attributes, SchemaLocation("attributes"), DataLocation("attributes")] async def to_dataframe(self, *keys: str, fb: IFeedback = NoFeedback) -> pd.DataFrame: - """Return collars with a categorical ``hole_id`` column and selected attributes.""" + """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), @@ -472,11 +525,11 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals if existing_indices and not replace: raise ObjectValidationError(f"Collection '{collection.name}' already exists") location_holes = await self._downhole_collection.location.holes.to_dataframe() - valid_indices = set(location_holes["hole_index"].astype(int)) table = collection.table - DownholeCollectionData._validate_hole_chunks( - collection.holes, len(table), valid_indices=valid_indices, require_coverage=False - ) + DownholeCollectionData._validate_hole_chunks(collection.holes, len(table), require_coverage=False) + valid_indices = set(location_holes["hole_index"].astype(int)) + if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): + raise ObjectValidationError("hole_index must reference a key in the location holes table") schema = await self._data_to_schema([collection], self._obj) if existing_indices: self._document[existing_indices[0]] = schema[0] diff --git a/packages/evo-objects/src/evo/objects/utils/downhole.py b/packages/evo-objects/src/evo/objects/utils/downhole.py index 0dedea8f..81337ef2 100644 --- a/packages/evo-objects/src/evo/objects/utils/downhole.py +++ b/packages/evo-objects/src/evo/objects/utils/downhole.py @@ -13,42 +13,51 @@ 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, *, dtype: pd.CategoricalDtype) -> pd.DataFrame: - """Run-length encode contiguous hole IDs using their categorical codes. +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. - A zero-count entry is emitted for categories absent from ``hole_ids``. An ID - outside ``dtype`` or a repeated non-contiguous run is rejected. + ``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. """ - unknown_mask = hole_ids.notna() & ~hole_ids.isin(dtype.categories) - if unknown_mask.any(): - unknown = hole_ids[unknown_mask].unique().tolist() - raise ValueError(f"hole_ids contains values absent from dtype: {unknown}") - categorical = hole_ids.astype(dtype) - codes = categorical.cat.codes.to_numpy(dtype=np.int32) - chunks: dict[int, tuple[int, int]] = {} + 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]) - if code < 0: - raise ValueError("hole_ids cannot contain missing values") end = start + 1 while end < len(codes) and codes[end] == code: end += 1 - if code in chunks: + if any(chunk_code == code for chunk_code, _, _ in chunks): raise ValueError("Rows for each hole_id must be contiguous") - chunks[code] = (start, end - start) + chunks.append((code, start, end - start)) start = end return pd.DataFrame( { - "hole_index": np.arange(len(dtype.categories), dtype=np.int32), - "offset": np.array([chunks.get(code, (0, 0))[0] for code in range(len(dtype.categories))], dtype=np.uint64), - "count": np.array([chunks.get(code, (0, 0))[1] for code in range(len(dtype.categories))], dtype=np.uint64), + "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), } ) diff --git a/packages/evo-objects/tests/test_downhole_utils.py b/packages/evo-objects/tests/test_downhole_utils.py index b884af8f..fd3e7233 100644 --- a/packages/evo-objects/tests/test_downhole_utils.py +++ b/packages/evo-objects/tests/test_downhole_utils.py @@ -17,26 +17,26 @@ class TestDownholeUtils(unittest.TestCase): - def _chunks(self, values, categories): - return hole_chunks_from_ids(pd.Series(values), dtype=pd.CategoricalDtype(categories=categories)) + 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"], ["A"]) + 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"], ["A"]) + 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"], ["A", "B"]) + 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"], ["A", "B", "C"]) + chunks = self._chunks(["A", "B", "B", "B", "C", "C"]) self.assertListEqual( chunks.to_dict("records"), [ @@ -46,36 +46,30 @@ def test_varying_hole_counts_are_preserved(self): ], ) - def test_category_absent_from_data_emits_zero_count_chunk(self): - chunks = self._chunks(["A", "A"], ["A", "B"]) - self.assertListEqual( - chunks.to_dict("records"), - [{"hole_index": 0, "offset": 0, "count": 2}, {"hole_index": 1, "offset": 0, "count": 0}], - ) + 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_categories_raises(self): + def test_id_absent_from_hole_indices_raises(self): with self.assertRaises(ValueError): - self._chunks(["A", "C"], ["A", "B"]) + self._chunks(["A", "C"], {"A": 1, "B": 42}) - def test_empty_input_emits_all_zero_count_chunks(self): - chunks = self._chunks([], ["A", "B"]) - self.assertListEqual( - chunks.to_dict("records"), - [{"hole_index": 0, "offset": 0, "count": 0}, {"hole_index": 1, "offset": 0, "count": 0}], - ) + 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"], ["A"])["hole_index"].dtype), "int32") + self.assertEqual(str(self._chunks(["A"])["hole_index"].dtype), "int32") def test_offset_and_count_dtypes_are_uint64(self): - chunks = self._chunks(["A"], ["A"]) + 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): - categories = [f"H{index:02d}" for index in range(50)] - values = [hole_id for hole_id in categories for _ in range(200)] - chunks = self._chunks(values, categories) + 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) @@ -83,37 +77,41 @@ def test_fifty_holes_with_two_hundred_rows_each_round_trip(self): expand_hole_index(chunks, len(values)).tolist(), [index for index in range(50) for _ in range(200)] ) - def test_category_order_is_preserved_not_lexicographically_sorted(self): - chunks = self._chunks(["M", "M", "Z"], ["Z", "A", "M"]) + 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": 2, "count": 1}, - {"hole_index": 1, "offset": 0, "count": 0}, - {"hole_index": 2, "offset": 0, "count": 2}, + {"hole_index": 0, "offset": 0, "count": 2}, + {"hole_index": 1, "offset": 2, "count": 1}, ], ) - def test_chunks_use_dtype_codes_and_round_trip(self): - dtype = pd.CategoricalDtype(categories=["Z", "A", "M"]) + def test_chunks_use_dense_creation_keys_and_round_trip(self): values = pd.Series(["M", "M", "Z"]) - chunks = hole_chunks_from_ids(values, dtype=dtype) - self.assertListEqual(chunks["hole_index"].tolist(), [0, 1, 2]) - self.assertListEqual(chunks["count"].tolist(), [1, 0, 2]) - self.assertListEqual(expand_hole_index(chunks, len(values)).tolist(), [2, 2, 0]) + 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): - dtype = pd.CategoricalDtype(categories=["A", "B"]) with self.assertRaises(ValueError): - hole_chunks_from_ids(pd.Series(["A", "B", "A"]), dtype=dtype) + 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"]), dtype=dtype) + hole_chunks_from_ids(pd.Series(["C"]), hole_indices={"A": 1, "B": 42}) - def test_empty_input_emits_zero_count_entries_with_schema_dtypes(self): - chunks = hole_chunks_from_ids(pd.Series([], dtype="string"), dtype=pd.CategoricalDtype(categories=["A", "B"])) - self.assertListEqual(chunks["hole_index"].tolist(), [0, 1]) - self.assertListEqual(chunks["offset"].tolist(), [0, 0]) - self.assertListEqual(chunks["count"].tolist(), [0, 0]) + 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") @@ -121,3 +119,7 @@ def test_empty_input_emits_zero_count_entries_with_schema_dtypes(self): 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_sparse_persisted_lookup_keys(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/typed/helpers.py b/packages/evo-objects/tests/typed/helpers.py index 4bd07589..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.""" diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 7ee0253b..075fb45b 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -170,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()) @@ -566,11 +567,9 @@ def test_attributes_length_raises(self): desurvey=None, ) - def test_location_chunks_require_each_creation_code_and_exact_path_coverage(self): + def test_location_chunks_require_exact_path_coverage_without_key_constraints(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": [0, 1], "offset": [0, 5], "count": [4, 2]}), pd.DataFrame({"hole_index": [0, 1], "offset": [0, 3], "count": [4, 3]}), ): @@ -588,18 +587,46 @@ def test_collection_chunk_ranges_are_not_required_to_partition_the_table(self): ) dataclasses.replace(base, collections=[collection]) - with self.assertRaises(ObjectValidationError): - dataclasses.replace( - base, - collections=[ - DistanceCollection( - name="invalid", - holes=pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), - table=table, - unit=None, - ) - ], - ) + dataclasses.replace( + base, + collections=[ + DistanceCollection( + name="opaque-key", + holes=pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), + table=table, + unit=None, + ) + ], + ) + + 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_bounding_box_falls_back_to_envelope_for_undeclared_keys(self): + base = _make_example_data(collections=[]) + data = dataclasses.replace(base, holes=pd.DataFrame({"hole_index": [7, 9], "offset": [0, 4], "count": [4, 3]})) + exact = base.compute_bounding_box() + bbox = data.compute_bounding_box() + self.assertLessEqual(bbox.min_x, exact.min_x) + self.assertGreaterEqual(bbox.max_x, exact.max_x) + self.assertLessEqual(bbox.min_z, exact.min_z) + self.assertGreaterEqual(bbox.max_z, exact.max_z) async def test_collection_allows_zero_row_hole(self): base = _make_example_data(collections=[]) From a02ac93b04eacf408c5124ae93aa29dd74d2bda0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 11 Aug 2026 09:01:20 +0100 Subject: [PATCH 26/30] feat(objects): enhance DownholeCollectionTable methods with detailed docstrings --- .../src/evo/objects/typed/downhole_collection.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) 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 c1e18f2f..d41296ba 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -514,9 +514,15 @@ 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( @@ -537,6 +543,12 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals 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] From 41451256b5b1551a63603987a8e52fbaf410b558 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Tue, 11 Aug 2026 09:31:51 +0100 Subject: [PATCH 27/30] docs(objects): clarify unit descriptions for collection coordinates and attribute handling --- packages/evo-objects/README.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index 720f564f..054a7e18 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -90,7 +90,7 @@ collections = [ name="geology", holes=hole_chunks_from_ids(pd.Series(["DH-01"])), table=intervals, - unit="m", # Explicit collection units override DataFrame metadata. + unit="m", # Collection coordinate units are set explicitly here. ) ] @@ -101,8 +101,9 @@ collections = [ 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 round-trip through `DataFrame.attrs["attribute_descriptions"]`; an explicit collection `unit` -takes precedence over unit metadata on the distance or `from` column. +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 From 1b48a4ee7029fe365bb4b206ba77885be56a400b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Fri, 14 Aug 2026 08:09:14 +0100 Subject: [PATCH 28/30] feat(objects): refine hole index handling and validation in DownholeCollectionData --- packages/evo-objects/README.md | 8 +- .../evo/objects/typed/downhole_collection.py | 170 ++++++++---------- .../src/evo/objects/utils/downhole.py | 4 +- .../tests/typed/test_downhole_collection.py | 52 +++--- 4 files changed, 107 insertions(+), 127 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index 054a7e18..baf2c0ec 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -72,10 +72,10 @@ Check out the other methods on the `ObjectAPIClient` for more details on how to ### 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 persisted `hole_id` lookup table; it is -not a row position and need not be zero-based or contiguous. `hole_chunks_from_ids()` is a convenience helper that emits -dense zero-based keys for new objects. Pass `hole_indices={"DH-01": 1}` when writing against an existing or explicit -lookup table; the helper creates chunks only for IDs present in its input. +distance tables, interval tables, or both. For `DownholeCollectionData`, `hole_index` is a dense zero-based key in the +sorted `hole_id` lookup. Location holes must contain each key exactly once, in any row order. Collection holes may +contain a subset or repeated chunks, but every key must reference a location hole. Persisted objects are read using their +actual lookup keys, which need not be zero-based or contiguous. ```python import pandas as pd 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 d41296ba..5bce5828 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -23,7 +23,6 @@ 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 DepthIntervalsTable, HoleIdCategory from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel @@ -36,8 +35,6 @@ DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, FLOAT_ARRAY_3, - INTEGER_ARRAY_1_INT32, - LOOKUP_TABLE_INT32, KnownTableFormat, ) @@ -90,8 +87,9 @@ class DownholeCollectionData(BaseSpatialObjectData): :param name: The name of the object. :param holes: A DataFrame describing which parts of `path` belong to which 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. + Columns: hole_index, offset, count. ``hole_index`` is a dense zero-based hole-ID lookup key. + Each hole must appear exactly once, but rows may be in any order. + :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 :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 @@ -114,31 +112,70 @@ class DownholeCollectionData(BaseSpatialObjectData): distance_unit: str | None desurvey: str | None + @property + def hole_ids(self) -> list[str]: + """Return hole IDs in the order used to assign dense lookup keys.""" + return sorted(self.properties["hole_id"].astype(str).tolist()) + + @property + def hole_indices(self) -> dict[str, int]: + """Map each hole ID to its dense zero-based lookup key.""" + return {hole_id: index for index, hole_id in enumerate(self.hole_ids)} + 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") - assert self.attributes is None or len(self.holes) == len(self.attributes) + 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") - self._validate_hole_chunks(self.holes, len(self.path), require_coverage=True) + valid_indices = set(range(len(self.hole_ids))) + 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), require_coverage=False) + self._validate_hole_chunks( + collection.holes, + len(table), + valid_indices=valid_indices, + require_hole_coverage=False, + require_table_coverage=False, + ) @staticmethod - def _validate_hole_chunks(holes: HoleChunks, table_length: int, *, require_coverage: bool) -> None: + 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") + if require_hole_coverage and (indices.duplicated().any() or set(indices) != valid_indices): + raise ObjectValidationError("Location holes must contain each hole_index 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_coverage: + if not require_table_coverage: return non_empty = sorted(zip(offsets[counts > 0], counts[counts > 0], strict=True)) @@ -150,62 +187,25 @@ def _validate_hole_chunks(holes: HoleChunks, table_length: int, *, require_cover 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]] | None: - """Map hole lookup keys to collar coordinates, or ``None`` when the mapping is not known. - - Keys are taken from an explicit ``properties["hole_index"]`` column when present. Otherwise the SDK creation - convention (dense zero-based keys in sorted hole-ID order) is used, and only when it accounts for every key - referenced by ``holes``. - """ - hole_ids = self.properties["hole_id"].astype(str) - if "hole_index" in self.properties: - keys = self.properties["hole_index"].astype(int) - else: - dense = {hole_id: index for index, hole_id in enumerate(sorted(hole_ids.unique()))} - keys = hole_ids.map(dense).astype(int) - coordinates = self.properties[_COORDINATE_COLUMNS].astype(float) - collars = { - int(key): coordinate - for key, coordinate in zip(keys, coordinates.itertuples(index=False, name=None), strict=True) - } - referenced = {int(chunk.hole_index) for chunk in self.holes.itertuples(index=False)} - return collars if referenced <= set(collars) else None - def compute_bounding_box(self) -> BoundingBox: - collars_by_hole_index = self._collars_by_hole_index() - if collars_by_hole_index is not None: - bboxes = [ - self._compute_hole_bounding_box( - self.path[int(chunk.offset) : int(chunk.offset) + int(chunk.count)], - collars_by_hole_index[int(chunk.hole_index)], - ) - for chunk in self.holes.itertuples(index=False) - if int(chunk.count) - ] - if bboxes: - return BoundingBox.combine(bboxes) - - # The collar for each chunk cannot be identified, so fall back to an envelope that contains every hole. - collars = BoundingBox.from_points(self.properties[_COORDINATE_COLUMNS].astype(float).to_numpy()) - relative_bboxes = [] - for chunk in self.holes.itertuples(index=False): - offset = int(chunk.offset) - count = int(chunk.count) + collars = self.properties.copy() + collars["_hole_index"] = collars["hole_id"].astype(str).map(self.hole_indices) + collars_by_index = collars.set_index("_hole_index") + + bboxes = [] + for hole_index, offset, count in self.holes[["hole_index", "offset", "count"]].itertuples( + index=False, name=None + ): + coordinates = collars_by_index.loc[int(hole_index), _COORDINATE_COLUMNS].to_numpy(dtype=np.float64) + collar = (float(coordinates[0]), float(coordinates[1]), float(coordinates[2])) + offset = int(offset) + count = int(count) if count: - relative_bboxes.append( - self._compute_hole_bounding_box(self.path[offset : offset + count], (0.0, 0.0, 0.0)) - ) - if not relative_bboxes: - return collars - relative = BoundingBox.combine(relative_bboxes) - return BoundingBox( - min_x=collars.min_x + relative.min_x, - max_x=collars.max_x + relative.max_x, - min_y=collars.min_y + relative.min_y, - max_y=collars.max_y + relative.max_y, - min_z=collars.min_z + relative.min_z, - max_z=collars.max_z + relative.max_z, - ) + 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 def _compute_bounding_box_np( @@ -337,34 +337,12 @@ async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: class LocationHoleIdCategory(HoleIdCategory): - """Collar hole IDs persisted with explicit integer lookup keys. - - ``properties`` may provide a ``hole_index`` column to persist specific lookup keys, for example when rewriting a - downloaded object. When it is absent, dense zero-based keys are generated in sorted hole-ID order as an SDK - creation convenience. - """ + """Collar hole IDs persisted with dense zero-based keys in sorted ID order.""" @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") - if "hole_index" in data: - keys = data["hole_index"].astype("int32") - if keys.duplicated().any(): - raise ObjectValidationError("properties hole_index values must be unique") - else: - dense = {hole_id: index for index, hole_id in enumerate(sorted(hole_ids.unique()))} - keys = hole_ids.map(dense).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 - ), - } + normalized = data.assign(hole_id=data["hole_id"].astype(str)) + return await super()._data_to_schema(normalized, context) class DownholeLocation(SchemaModel): @@ -532,10 +510,14 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals raise ObjectValidationError(f"Collection '{collection.name}' already exists") location_holes = await self._downhole_collection.location.holes.to_dataframe() table = collection.table - DownholeCollectionData._validate_hole_chunks(collection.holes, len(table), require_coverage=False) valid_indices = set(location_holes["hole_index"].astype(int)) - if not set(collection.holes["hole_index"].astype(int)).issubset(valid_indices): - raise ObjectValidationError("hole_index must reference a key in the location holes table") + 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] diff --git a/packages/evo-objects/src/evo/objects/utils/downhole.py b/packages/evo-objects/src/evo/objects/utils/downhole.py index 81337ef2..70f82188 100644 --- a/packages/evo-objects/src/evo/objects/utils/downhole.py +++ b/packages/evo-objects/src/evo/objects/utils/downhole.py @@ -27,7 +27,9 @@ def hole_chunks_from_ids(hole_ids: pd.Series, *, hole_indices: Mapping[str, int] ``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. + produce chunks; an empty collection therefore produces no chunks. Explicit + mappings are intended for collection chunks used with persisted objects; + location chunks in ``DownholeCollectionData`` must use dense zero-based keys. """ if hole_ids.isna().any(): raise ValueError("hole_ids cannot contain missing values") diff --git a/packages/evo-objects/tests/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 075fb45b..65948da4 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -539,7 +539,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( @@ -567,9 +567,12 @@ def test_attributes_length_raises(self): desurvey=None, ) - def test_location_chunks_require_exact_path_coverage_without_key_constraints(self): + def test_location_chunks_require_each_dense_index_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]}), ): @@ -587,46 +590,39 @@ def test_collection_chunk_ranges_are_not_required_to_partition_the_table(self): ) dataclasses.replace(base, collections=[collection]) - dataclasses.replace( - base, - collections=[ - DistanceCollection( - name="opaque-key", - holes=pd.DataFrame({"hole_index": [2], "offset": [0], "count": [2]}), - table=table, - unit=None, - ) - ], + 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"): + dataclasses.replace(base, collections=[invalid_collection]) - async def test_data_and_reads_support_one_based_sparse_lookup_keys(self): + async def test_location_holes_may_be_out_of_hole_index_order(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]}) + holes = pd.DataFrame({"hole_index": [1, 0], "offset": [4, 0], "count": [3, 4]}) collection = DistanceCollection( name="measurements", - holes=pd.DataFrame({"hole_index": [42], "offset": [0], "count": [1]}), + holes=pd.DataFrame({"hole_index": [1], "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: + data = dataclasses.replace(data, holes=holes, collections=[collection]) + with self._mock_geoscience_objects(): 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_bounding_box_falls_back_to_envelope_for_undeclared_keys(self): + def test_bounding_box_includes_collar_for_hole_without_path_rows(self): base = _make_example_data(collections=[]) - data = dataclasses.replace(base, holes=pd.DataFrame({"hole_index": [7, 9], "offset": [0, 4], "count": [4, 3]})) - exact = base.compute_bounding_box() + 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.assertLessEqual(bbox.min_x, exact.min_x) - self.assertGreaterEqual(bbox.max_x, exact.max_x) - self.assertLessEqual(bbox.min_z, exact.min_z) - self.assertGreaterEqual(bbox.max_z, exact.max_z) + 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=[]) From ea86f6ebbb2337e20c7502d7358338d8e462e795 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Fri, 28 Aug 2026 13:55:54 +0100 Subject: [PATCH 29/30] docs(objects): update documentation for hole_chunks_from_ids to clarify key assignment behavior --- packages/evo-objects/src/evo/objects/utils/downhole.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/packages/evo-objects/src/evo/objects/utils/downhole.py b/packages/evo-objects/src/evo/objects/utils/downhole.py index 70f82188..663e1b41 100644 --- a/packages/evo-objects/src/evo/objects/utils/downhole.py +++ b/packages/evo-objects/src/evo/objects/utils/downhole.py @@ -24,12 +24,10 @@ 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. Explicit - mappings are intended for collection chunks used with persisted objects; - location chunks in ``DownholeCollectionData`` must use dense zero-based keys. + ``hole_indices`` maps hole IDs to their lookup-table keys. When omitted, the + keys are dense, zero-based, and assigned in sorted ID order. Explicit mappings + may use other unique 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") From 032b3e16eff13aed7740b4e4f65c7bb41e8caed4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Peixoto?= Date: Wed, 2 Sep 2026 10:06:46 +0100 Subject: [PATCH 30/30] feat(downhole): enhance hole index handling and validation, update documentation --- packages/evo-objects/README.md | 8 +- .../src/evo/objects/typed/attributes.py | 2 +- .../evo/objects/typed/downhole_collection.py | 84 +++++++++++-------- .../src/evo/objects/utils/downhole.py | 8 +- .../evo-objects/tests/test_downhole_utils.py | 2 +- .../tests/typed/test_downhole_collection.py | 34 +++++++- 6 files changed, 92 insertions(+), 46 deletions(-) diff --git a/packages/evo-objects/README.md b/packages/evo-objects/README.md index baf2c0ec..ab93e4d3 100644 --- a/packages/evo-objects/README.md +++ b/packages/evo-objects/README.md @@ -72,10 +72,10 @@ Check out the other methods on the `ObjectAPIClient` for more details on how to ### 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. For `DownholeCollectionData`, `hole_index` is a dense zero-based key in the -sorted `hole_id` lookup. Location holes must contain each key exactly once, in any row order. Collection holes may -contain a subset or repeated chunks, but every key must reference a location hole. Persisted objects are read using their -actual lookup keys, which need not be zero-based or contiguous. +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 diff --git a/packages/evo-objects/src/evo/objects/typed/attributes.py b/packages/evo-objects/src/evo/objects/typed/attributes.py index 447be3b3..3f35a4ff 100644 --- a/packages/evo-objects/src/evo/objects/typed/attributes.py +++ b/packages/evo-objects/src/evo/objects/typed/attributes.py @@ -637,7 +637,7 @@ 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 + 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): 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 5bce5828..d5d2c83e 100644 --- a/packages/evo-objects/src/evo/objects/typed/downhole_collection.py +++ b/packages/evo-objects/src/evo/objects/typed/downhole_collection.py @@ -23,6 +23,7 @@ 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 DepthIntervalsTable, HoleIdCategory from evo.objects.typed._model import DataLocation, SchemaList, SchemaLocation, SchemaModel @@ -35,6 +36,8 @@ DOWNHOLE_COLLECTION_LOCATION_HOLES, FLOAT_ARRAY_1, FLOAT_ARRAY_3, + INTEGER_ARRAY_1_INT32, + LOOKUP_TABLE_INT32, KnownTableFormat, ) @@ -81,16 +84,27 @@ class IntervalCollection: 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 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_index, offset, count. ``hole_index`` is a dense zero-based hole-ID lookup key. - Each hole must appear exactly once, but rows may be in any order. + 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 + 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. @@ -112,16 +126,6 @@ class DownholeCollectionData(BaseSpatialObjectData): distance_unit: str | None desurvey: str | None - @property - def hole_ids(self) -> list[str]: - """Return hole IDs in the order used to assign dense lookup keys.""" - return sorted(self.properties["hole_id"].astype(str).tolist()) - - @property - def hole_indices(self) -> dict[str, int]: - """Map each hole ID to its dense zero-based lookup key.""" - return {hole_id: index for index, hole_id in enumerate(self.hole_ids)} - def __post_init__(self): hole_ids = self.properties["hole_id"] if hole_ids.isna().any(): @@ -136,7 +140,7 @@ def __post_init__(self): if len(names) != len(set(names)): raise ObjectValidationError("Collection names must be unique") - valid_indices = set(range(len(self.hole_ids))) + valid_indices = set(_property_hole_indices(self.properties)) self._validate_hole_chunks( self.holes, len(self.path), @@ -168,9 +172,9 @@ def _validate_hole_chunks( 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") + 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_index exactly once") + 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(): @@ -187,19 +191,22 @@ def _validate_hole_chunks( if expected_offset != table_length: raise ObjectValidationError("Hole chunk ranges must cover the associated table exactly once") - def compute_bounding_box(self) -> BoundingBox: - collars = self.properties.copy() - collars["_hole_index"] = collars["hole_id"].astype(str).map(self.hole_indices) - collars_by_index = collars.set_index("_hole_index") + 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 hole_index, offset, count in self.holes[["hole_index", "offset", "count"]].itertuples( - index=False, name=None - ): - coordinates = collars_by_index.loc[int(hole_index), _COORDINATE_COLUMNS].to_numpy(dtype=np.float64) - collar = (float(coordinates[0]), float(coordinates[1]), float(coordinates[2])) - offset = int(offset) - count = int(count) + 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: @@ -337,12 +344,23 @@ async def _data_to_schema(cls, data: HoleAttributes, context: IContext) -> Any: class LocationHoleIdCategory(HoleIdCategory): - """Collar hole IDs persisted with dense zero-based keys in sorted ID order.""" + """Collar hole IDs persisted with explicit integer lookup keys.""" @classmethod async def _data_to_schema(cls, data: pd.DataFrame, context: IContext) -> Any: - normalized = data.assign(hole_id=data["hole_id"].astype(str)) - return await super()._data_to_schema(normalized, context) + 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): @@ -403,7 +421,7 @@ def _downhole_collection(self) -> DownholeCollection: return root async def _table_by_hole(self, data: pd.DataFrame, *, fb: IFeedback) -> dict[str, pd.DataFrame]: - """Group table rows by hole using the containing DownholeCollection's hole-id lookup.""" + """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]]] = {} @@ -508,9 +526,9 @@ async def add(self, collection: DownholeCollectionEntry, *, replace: bool = Fals ) if existing_indices and not replace: raise ObjectValidationError(f"Collection '{collection.name}' already exists") - location_holes = await self._downhole_collection.location.holes.to_dataframe() + lookup = await self._downhole_collection.location.hole_id.to_indexed_dataframe() table = collection.table - valid_indices = set(location_holes["hole_index"].astype(int)) + valid_indices = set(lookup["key"].astype(int)) DownholeCollectionData._validate_hole_chunks( collection.holes, len(table), diff --git a/packages/evo-objects/src/evo/objects/utils/downhole.py b/packages/evo-objects/src/evo/objects/utils/downhole.py index 663e1b41..81337ef2 100644 --- a/packages/evo-objects/src/evo/objects/utils/downhole.py +++ b/packages/evo-objects/src/evo/objects/utils/downhole.py @@ -24,10 +24,10 @@ 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, the - keys are dense, zero-based, and assigned in sorted ID order. Explicit mappings - may use other unique keys. Only IDs present in ``hole_ids`` produce chunks; an - empty collection therefore produces no 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") diff --git a/packages/evo-objects/tests/test_downhole_utils.py b/packages/evo-objects/tests/test_downhole_utils.py index fd3e7233..a4a8ab9d 100644 --- a/packages/evo-objects/tests/test_downhole_utils.py +++ b/packages/evo-objects/tests/test_downhole_utils.py @@ -120,6 +120,6 @@ 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_sparse_persisted_lookup_keys(self): + 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/typed/test_downhole_collection.py b/packages/evo-objects/tests/typed/test_downhole_collection.py index 65948da4..e31edc47 100644 --- a/packages/evo-objects/tests/typed/test_downhole_collection.py +++ b/packages/evo-objects/tests/typed/test_downhole_collection.py @@ -461,11 +461,12 @@ 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_hole_index_not_property_position(self): + 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( @@ -567,7 +568,7 @@ def test_attributes_length_raises(self): desurvey=None, ) - def test_location_chunks_require_each_dense_index_and_exact_path_coverage(self): + 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]}), @@ -596,7 +597,7 @@ def test_collection_chunk_ranges_are_not_required_to_partition_the_table(self): table=table, unit=None, ) - with self.assertRaisesRegex(ObjectValidationError, "hole_index must reference a valid hole"): + 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): @@ -614,6 +615,33 @@ async def test_location_holes_may_be_out_of_hole_index_order(self): 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(