From 969f1c4b9aac69cad510999a4b461d1163c6e2ef Mon Sep 17 00:00:00 2001
From: Tonio Fincke <tonio.fincke@brockmann-consult.de>
Date: Mon, 18 Nov 2024 08:15:12 +0100
Subject: [PATCH] work on integrating vector data cubes

---
 .../core/store/fs/impl/test_vectordatacube.py | 11 +++++++++++
 xcube/core/mldataset/abc.py                   | 19 +++++++++++++++++++
 xcube/core/store/datatype.py                  |  2 ++
 xcube/core/store/fs/store.py                  |  6 +++++-
 4 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/test/core/store/fs/impl/test_vectordatacube.py b/test/core/store/fs/impl/test_vectordatacube.py
index 4ba9f604a..7f46d765b 100644
--- a/test/core/store/fs/impl/test_vectordatacube.py
+++ b/test/core/store/fs/impl/test_vectordatacube.py
@@ -6,11 +6,22 @@
 
 from xcube.core.new import new_vector_data_cube
 
+from xcube.core.store import new_data_store
 from xcube.core.store.datatype import VECTOR_DATA_CUBE_TYPE
 from xcube.core.store.fs.impl.vectordatacube import VectorDataCubeZarrFsDataAccessor
 from xcube.core.store.fs.impl.vectordatacube import VectorDataCubeNetcdfFsDataAccessor
 
 
+class VectorDataCubeStoreTest(unittest.TestCase):
+
+    def test_write_to_and_read_from_store(self):
+        store = new_data_store("file")
+        vdc = new_vector_data_cube()
+        data_id = store.write_data(
+            vdc, data_id="vdc_test.zarr"
+        )
+        self.assertIsNotNone(data_id)
+
 class VectorDataCubeZarrFsDataAccessorTest(unittest.TestCase):
 
     @classmethod
diff --git a/xcube/core/mldataset/abc.py b/xcube/core/mldataset/abc.py
index 14450baf1..faa192581 100644
--- a/xcube/core/mldataset/abc.py
+++ b/xcube/core/mldataset/abc.py
@@ -9,6 +9,7 @@
 from collections.abc import Sequence
 
 import xarray as xr
+import xvec
 
 from xcube.core.gridmapping import GridMapping
 from xcube.core.tilingscheme import TilingScheme
@@ -148,3 +149,21 @@ def get_level_for_resolution(self, xy_res: ScalarOrPair[float]) -> int:
             if x_res > given_x_res and y_res > given_y_res:
                 return max(0, level - 1)
         return self.num_levels - 1
+
+
+class VectorDataCube(xr.Dataset):
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        for coord_name, coord_values in self.coords.items():
+            if isinstance(coord_values.to_index(), xvec.GeometryIndex):
+                return
+        raise ValueError(f"Dataset has no Coordinate with an xvec.GeometryIndex.")
+
+    @classmethod
+    def from_dataset(cls, dataset):
+        return cls(
+            data_vars=dataset.data_vars,
+            coords=dataset.coords,
+            attrs=dataset.attrs
+        )
diff --git a/xcube/core/store/datatype.py b/xcube/core/store/datatype.py
index 2f19a2ff6..0d1ba5846 100644
--- a/xcube/core/store/datatype.py
+++ b/xcube/core/store/datatype.py
@@ -214,6 +214,8 @@ def register_default_data_types():
         MULTI_LEVEL_DATASET_ITERATOR_TYPE,
         GEO_DATA_FRAME_TYPE,
         GEO_DATA_FRAME_ITERATOR_TYPE,
+        VECTOR_DATA_CUBE_TYPE,
+        VECTOR_DATA_CUBE_ITERATOR_TYPE,
     ]:
         DataType.register_data_type(data_type)
 
diff --git a/xcube/core/store/fs/store.py b/xcube/core/store/fs/store.py
index 2852130ab..0e400e6c3 100644
--- a/xcube/core/store/fs/store.py
+++ b/xcube/core/store/fs/store.py
@@ -51,6 +51,7 @@
 from ..datatype import DataTypeLike
 from ..datatype import GEO_DATA_FRAME_TYPE
 from ..datatype import MULTI_LEVEL_DATASET_TYPE
+from ..datatype import VECTOR_DATA_CUBE_TYPE
 from ..descriptor import DataDescriptor
 from ..descriptor import new_data_descriptor
 from ..error import DataStoreError
@@ -443,7 +444,10 @@ def _guess_writer_id(self, data, data_id: str = None):
                 data_type = accessor_id_parts[0]
                 format_id = accessor_id_parts[1]
         if isinstance(data, xr.Dataset):
-            data_type = DATASET_TYPE.alias
+            if len(data.xvec.geom_coords) > 0:
+                data_type = VECTOR_DATA_CUBE_TYPE.alias
+            else:
+                data_type = DATASET_TYPE.alias
             format_id = format_id or "zarr"
         elif isinstance(data, MultiLevelDataset):
             data_type = MULTI_LEVEL_DATASET_TYPE.alias