Repository: TileDB-Inc/TileDB-Py
Pull request #2110 (2024-11-18 07:40:42)

Implement serialization functionality for `Config` and `VFS` objects
Also removes unnecessary `PatchedConfig` and `PatchedCtx`.

---

[sc-55042]


Patch:
diff --git a/pyproject.toml b/pyproject.toml
index 8ee52a50..0e822ea3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -76,7 +76,7 @@ sdist.include = ["tiledb/_generated_version.py"]
 TILEDB_PATH = {env="TILEDB_PATH"}
 TILEDB_VERSION = {env="TILEDB_VERSION"}
 TILEDB_HASH = {env="TILEDB_HASH"}
-TILEDB_REMOVE_DEPRECATIONS = "ON"
+TILEDB_REMOVE_DEPRECATIONS = "OFF"
 TILEDB_SERIALIZATION = "OFF"
 
 [tool.pytest.ini_options]
diff --git a/tiledb/array_schema.py b/tiledb/array_schema.py
index 09cf727f..d44d4b51 100644
--- a/tiledb/array_schema.py
+++ b/tiledb/array_schema.py
@@ -396,11 +396,9 @@ class ArraySchema(CtxMixin, lt.ArraySchema):
 
             :rtype: tiledb.CurrentDomain
             """
-            curr_dom = CurrentDomain.from_pybind11(
+            return CurrentDomain.from_pybind11(
                 self._ctx, self._current_domain(self._ctx)
             )
-            curr_dom._set_domain(self.domain)
-            return curr_dom
 
         def set_current_domain(self, current_domain):
             """Set the current domain
diff --git a/tiledb/cc/current_domain.cc b/tiledb/cc/current_domain.cc
index f45b1c33..2436fbf7 100644
--- a/tiledb/cc/current_domain.cc
+++ b/tiledb/cc/current_domain.cc
@@ -118,87 +118,86 @@ void init_current_domain(py::module &m) {
           py::arg("dim_name"), py::arg("start"), py::arg("end"))
 
       .def("_range",
-           [](NDRectangle &ndrect, const std::string &dim_name,
-              const py::dtype &n_type) -> py::tuple {
-             if (n_type.is(py::dtype::of<uint64_t>())) {
+           [](NDRectangle &ndrect, const std::string &dim_name) -> py::tuple {
+             const tiledb_datatype_t n_type = ndrect.range_dtype(dim_name);
+             if (n_type == TILEDB_UINT64) {
                auto range = ndrect.range<uint64_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int64_t>())) {
+             } else if (n_type == TILEDB_INT64) {
                auto range = ndrect.range<int64_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<uint32_t>())) {
+             } else if (n_type == TILEDB_UINT32) {
                auto range = ndrect.range<uint32_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int32_t>())) {
+             } else if (n_type == TILEDB_INT32) {
                auto range = ndrect.range<int32_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<uint16_t>())) {
+             } else if (n_type == TILEDB_UINT16) {
                auto range = ndrect.range<uint16_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int16_t>())) {
+             } else if (n_type == TILEDB_INT16) {
                auto range = ndrect.range<int16_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<uint8_t>())) {
+             } else if (n_type == TILEDB_UINT8) {
                auto range = ndrect.range<uint8_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int8_t>())) {
+             } else if (n_type == TILEDB_INT8) {
                auto range = ndrect.range<int8_t>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<double>())) {
+             } else if (n_type == TILEDB_FLOAT64) {
                auto range = ndrect.range<double>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<float>())) {
+             } else if (n_type == TILEDB_FLOAT32) {
                auto range = ndrect.range<float>(dim_name);
                return py::make_tuple(range[0], range[1]);
-             } else if (py::getattr(n_type, "kind").is(py::str("S")) ||
-                        py::getattr(n_type, "kind").is(py::str("U"))) {
+             } else if (n_type == TILEDB_STRING_ASCII ||
+                        n_type == TILEDB_STRING_UTF8) {
                auto range = ndrect.range<std::string>(dim_name);
                return py::make_tuple(range[0], range[1]);
              } else {
                TPY_ERROR_LOC("Unsupported type for NDRectangle's range");
              }
            })
-      .def("_range",
-           [](NDRectangle &ndrect, unsigned dim_idx,
-              const py::dtype &n_type) -> py::tuple {
-             if (n_type.is(py::dtype::of<uint64_t>())) {
-               auto range = ndrect.range<uint64_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int64_t>())) {
-               auto range = ndrect.range<int64_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<uint32_t>())) {
-               auto range = ndrect.range<uint32_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int32_t>())) {
-               auto range = ndrect.range<int32_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<uint16_t>())) {
-               auto range = ndrect.range<uint16_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int16_t>())) {
-               auto range = ndrect.range<int16_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<uint8_t>())) {
-               auto range = ndrect.range<uint8_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<int8_t>())) {
-               auto range = ndrect.range<int8_t>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<double>())) {
-               auto range = ndrect.range<double>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (n_type.is(py::dtype::of<float>())) {
-               auto range = ndrect.range<float>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else if (py::getattr(n_type, "kind").is(py::str("S")) ||
-                        py::getattr(n_type, "kind").is(py::str("U"))) {
-               auto range = ndrect.range<std::string>(dim_idx);
-               return py::make_tuple(range[0], range[1]);
-             } else {
-               TPY_ERROR_LOC("Unsupported type for NDRectangle's range");
-             }
-           });
+      .def("_range", [](NDRectangle &ndrect, unsigned dim_idx) -> py::tuple {
+        const tiledb_datatype_t n_type = ndrect.range_dtype(dim_idx);
+        if (n_type == TILEDB_UINT64) {
+          auto range = ndrect.range<uint64_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_INT64) {
+          auto range = ndrect.range<int64_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_UINT32) {
+          auto range = ndrect.range<uint32_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_INT32) {
+          auto range = ndrect.range<int32_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_UINT16) {
+          auto range = ndrect.range<uint16_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_INT16) {
+          auto range = ndrect.range<int16_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_UINT8) {
+          auto range = ndrect.range<uint8_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_INT8) {
+          auto range = ndrect.range<int8_t>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_FLOAT64) {
+          auto range = ndrect.range<double>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_FLOAT32) {
+          auto range = ndrect.range<float>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else if (n_type == TILEDB_STRING_ASCII ||
+                   n_type == TILEDB_STRING_UTF8) {
+          auto range = ndrect.range<std::string>(dim_idx);
+          return py::make_tuple(range[0], range[1]);
+        } else {
+          TPY_ERROR_LOC("Unsupported type for NDRectangle's range");
+        }
+      });
 
   py::class_<CurrentDomain>(m, "CurrentDomain")
       .def(py::init<CurrentDomain>())
@@ -210,14 +209,14 @@ void init_current_domain(py::module &m) {
              return py::capsule(curr_dom.ptr().get(), "curr_dom");
            })
 
+      .def_property_readonly("_is_empty", &CurrentDomain::is_empty)
+
       .def_property_readonly("_type", &CurrentDomain::type)
 
       .def("_set_ndrectangle", &CurrentDomain::set_ndrectangle,
            py::arg("ndrect"))
 
-      .def("_ndrectangle", &CurrentDomain::ndrectangle)
-
-      .def("_is_empty", &CurrentDomain::is_empty);
+      .def("_ndrectangle", &CurrentDomain::ndrectangle);
 #endif
 }
 
diff --git a/tiledb/ctx.py b/tiledb/ctx.py
index b4898f37..cc1a6314 100644
--- a/tiledb/ctx.py
+++ b/tiledb/ctx.py
@@ -293,6 +293,20 @@ class Config(lt.Config):
         """
         self.save_to_file(uri)
 
+    def __reduce__(self):
+        """
+        Customize the pickling process by defining how to serialize
+        and reconstruct the Config object.
+        """
+        state = self.dict()
+        return (self.__class__, (), state)
+
+    def __setstate__(self, state):
+        """
+        Customize how the Config object is restored from its serialized state.
+        """
+        self.__init__(state)
+
 
 class ConfigKeys:
     """
diff --git a/tiledb/current_domain.py b/tiledb/current_domain.py
index f00755f8..b842ad73 100644
--- a/tiledb/current_domain.py
+++ b/tiledb/current_domain.py
@@ -18,9 +18,6 @@ class CurrentDomain(CtxMixin, lt.CurrentDomain):
         """
         super().__init__(ctx)
 
-    def _set_domain(self, domain: Domain):
-        self._domain = domain
-
     @property
     def type(self):
         """The type of the current domain.
@@ -35,7 +32,7 @@ class CurrentDomain(CtxMixin, lt.CurrentDomain):
 
         :rtype: bool
         """
-        return self._is_empty()
+        return self._is_empty
 
     def set_ndrectangle(self, ndrect: NDRectangle):
         """Sets an N-dimensional rectangle representation on a current domain.
@@ -44,7 +41,6 @@ class CurrentDomain(CtxMixin, lt.CurrentDomain):
         :raises tiledb.TileDBError:
         """
         self._set_ndrectangle(ndrect)
-        self._domain = ndrect._get_domain()
 
     @property
     def ndrectangle(self):
@@ -53,6 +49,4 @@ class CurrentDomain(CtxMixin, lt.CurrentDomain):
         :rtype: NDRectangle
         :raises tiledb.TileDBError:
         """
-        ndrect = NDRectangle.from_pybind11(self._ctx, self._ndrectangle())
-        ndrect._set_domain(self._domain)
-        return ndrect
+        return NDRectangle.from_pybind11(self._ctx, self._ndrectangle())
diff --git a/tiledb/ndrectangle.py b/tiledb/ndrectangle.py
index a8194091..e071e403 100644
--- a/tiledb/ndrectangle.py
+++ b/tiledb/ndrectangle.py
@@ -19,7 +19,6 @@ class NDRectangle(CtxMixin, lt.NDRectangle):
         :raises tiledb.TileDBError:
         """
         super().__init__(ctx, domain)
-        self._set_domain(domain)
 
     def __str__(self) -> str:
         dimensions_str = ", ".join(
@@ -28,12 +27,6 @@ class NDRectangle(CtxMixin, lt.NDRectangle):
         )
         return f"NDRectangle({dimensions_str})"
 
-    def _set_domain(self, domain: Domain):
-        self._domain = domain
-
-    def _get_domain(self) -> Domain:
-        return self._domain
-
     def set_range(
         self,
         dim: Union[str, int],
@@ -58,4 +51,4 @@ class NDRectangle(CtxMixin, lt.NDRectangle):
         :return: Range as a tuple (start, end)
         :raises tiledb.TileDBError:
         """
-        return self._range(dim, self._domain.dim(dim).dtype)
+        return self._range(dim)
diff --git a/tiledb/vfs.py b/tiledb/vfs.py
index 7a703fae..5b271319 100644
--- a/tiledb/vfs.py
+++ b/tiledb/vfs.py
@@ -25,7 +25,7 @@ class VFS(lt.VFS):
     """
 
     def __init__(self, config: Union[Config, dict] = None, ctx: Optional[Ctx] = None):
-        ctx = ctx or default_ctx()
+        self.ctx = ctx or default_ctx()
 
         if config:
             from .libtiledb import Config
@@ -39,12 +39,12 @@ class VFS(lt.VFS):
                     raise ValueError("`config` argument must be of type Config or dict")
 
             # Convert all values to strings
-            config = {k: str(v) for k, v in config.items()}
+            self.config_dict = {k: str(v) for k, v in config.items()}
 
-            ccfg = tiledb.Config(config)
-            super().__init__(ctx, ccfg)
+            ccfg = tiledb.Config(self.config_dict)
+            super().__init__(self.ctx, ccfg)
         else:
-            super().__init__(ctx)
+            super().__init__(self.ctx)
 
     def ctx(self) -> Ctx:
         """
@@ -329,6 +329,21 @@ class VFS(lt.VFS):
     isfile = is_file
     size = file_size
 
+    # pickling support
+    def __getstate__(self):
+        # self.config_dict might not exist. In that case use the config from ctx.
+        if hasattr(self, "config_dict"):
+            config_dict = self.config_dict
+        else:
+            config_dict = self.config().dict()
+        return (config_dict,)
+
+    def __setstate__(self, state):
+        config_dict = state[0]
+        config = Config(params=config_dict)
+        ctx = Ctx(config)
+        self.__init__(config=config, ctx=ctx)
+
 
 class FileIO(io.RawIOBase):
     """TileDB FileIO class that encapsulates files opened by tiledb.VFS. The file


Test patch:
diff --git a/tiledb/tests/conftest.py b/tiledb/tests/conftest.py
index 6565dd89..ecce429e 100644
--- a/tiledb/tests/conftest.py
+++ b/tiledb/tests/conftest.py
@@ -52,54 +52,6 @@ def pytest_configure(config):
     # default must be set here rather than globally
     pytest.tiledb_vfs = "file"
 
-    vfs_config(config)
-
-
-def vfs_config(pytestconfig):
-    vfs_config_override = {}
-
-    vfs = pytestconfig.getoption("vfs")
-    if vfs == "s3":
-        pytest.tiledb_vfs = "s3"
-
-        vfs_config_override.update(
-            {
-                "vfs.s3.endpoint_override": "localhost:9999",
-                "vfs.s3.aws_access_key_id": "minio",
-                "vfs.s3.aws_secret_access_key": "miniosecretkey",
-                "vfs.s3.scheme": "https",
-                "vfs.s3.verify_ssl": False,
-                "vfs.s3.use_virtual_addressing": False,
-            }
-        )
-
-    vfs_config_arg = pytestconfig.getoption("vfs-config", None)
-    if vfs_config_arg:
-        pass
-
-    tiledb._orig_ctx = tiledb.Ctx
-
-    def get_config(config):
-        final_config = {}
-        if isinstance(config, tiledb.Config):
-            final_config = config.dict()
-        elif config:
-            final_config = config
-
-        final_config.update(vfs_config_override)
-        return final_config
-
-    class PatchedCtx(tiledb.Ctx):
-        def __init__(self, config=None):
-            super().__init__(get_config(config))
-
-    class PatchedConfig(tiledb.Config):
-        def __init__(self, params=None):
-            super().__init__(get_config(params))
-
-    tiledb.Ctx = PatchedCtx
-    tiledb.Config = PatchedConfig
-
 
 @pytest.fixture(scope="function", autouse=True)
 def isolate_os_fork(original_os_fork):
diff --git a/tiledb/tests/test_context_and_config.py b/tiledb/tests/test_context_and_config.py
index 05433899..13ce7b72 100644
--- a/tiledb/tests/test_context_and_config.py
+++ b/tiledb/tests/test_context_and_config.py
@@ -1,4 +1,6 @@
+import io
 import os
+import pickle
 import subprocess
 import sys
 import xml
@@ -261,3 +263,20 @@ class TestConfig(DiskTestCase):
             pytest.fail(
                 f"Could not parse config._repr_html_(). Saw {config._repr_html_()}"
             )
+
+    def test_config_pickle(self):
+        # test that Config can be pickled and unpickled
+        config = tiledb.Config(
+            {
+                "rest.use_refactored_array_open": "false",
+                "rest.use_refactored_array_open_and_query_submit": "true",
+                "vfs.azure.storage_account_name": "myaccount",
+            }
+        )
+        with io.BytesIO() as buf:
+            pickle.dump(config, buf)
+            buf.seek(0)
+            config2 = pickle.load(buf)
+
+            self.assertIsInstance(config2, tiledb.Config)
+            self.assertEqual(config2.dict(), config.dict())
diff --git a/tiledb/tests/test_vfs.py b/tiledb/tests/test_vfs.py
index 8219539a..da710577 100644
--- a/tiledb/tests/test_vfs.py
+++ b/tiledb/tests/test_vfs.py
@@ -1,6 +1,7 @@
 import io
 import os
 import pathlib
+import pickle
 import random
 import sys
 
@@ -239,6 +240,21 @@ class TestVFS(DiskTestCase):
             txtio = io.TextIOWrapper(f2, encoding="utf-8")
             self.assertEqual(txtio.readlines(), lines)
 
+    def test_pickle(self):
+        # test that vfs can be pickled and unpickled with config options
+        config = tiledb.Config(
+            {"vfs.s3.region": "eu-west-1", "vfs.max_parallel_ops": "1"}
+        )
+        vfs = tiledb.VFS(config)
+        with io.BytesIO() as buf:
+            pickle.dump(vfs, buf)
+            buf.seek(0)
+            vfs2 = pickle.load(buf)
+
+            self.assertIsInstance(vfs2, tiledb.VFS)
+            self.assertEqual(vfs2.config()["vfs.s3.region"], "eu-west-1")
+            self.assertEqual(vfs2.config()["vfs.max_parallel_ops"], "1")
+
     def test_sc42569_vfs_memoryview(self):
         # This test is to ensure that giving np.ndarray buffer to readinto works
         # when trying to write bytes that cannot be converted to float32 or int32