Merge pull request #7 from msoedov/ralph/improvements

improvements

Author: msoedov

.flake8

@@ -1,5 +1,7 @@
 [flake8]
     max-line-length = 160
+    # E203 is whitespace before ':' - conflicts with black's slice formatting
+    extend-ignore = E203
     per-file-ignores =
         # Ignore docstring lints for tests
         *: D100, D101, D102, D103, D104, D107, D105, D401, D400, D205, D202, D209

.pre-commit-config.yaml

@@ -1,59 +1,57 @@
-
 default_language_version:
-        python: python3
+        python: python3.11
 
 repos:
 
   - repo: https://github.com/asottile/pyupgrade
-    rev: v2.31.1
+    rev: v3.21.2
     hooks:
       - id: pyupgrade
-        args: [--py39-plus]
+        args: [--py311-plus]
 
   - repo: https://github.com/psf/black
-    rev: 22.8.0
+    rev: 26.1.0
     hooks:
       - id: black
-        language_version: python3.9
+        language_version: python3.11
 
   - repo: https://github.com/pycqa/flake8
-    rev: 5.0.4
+    rev: 7.3.0
     hooks:
       - id: flake8
         language_version: python3
         additional_dependencies: [flake8-docstrings]
 
   - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
+    rev: 7.0.0
     hooks:
       - id: isort
         args: [--profile, black]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v6.0.0
     hooks:
       - id: end-of-file-fixer
       - id: check-yaml
       - id: trailing-whitespace
       - id: check-merge-conflict
 
   - repo: https://github.com/executablebooks/mdformat
-    rev: 0.7.14
+    rev: 1.0.0
     hooks:
       - id: mdformat
         name: mdformat
         entry: mdformat .
         language_version: python3
 
-
-  - repo: https://github.com/myint/docformatter
-    rev: v1.4
+  - repo: https://github.com/PyCQA/docformatter
+    rev: v1.7.7
     hooks:
       - id: docformatter
         args: [--in-place]
 
   - repo: https://github.com/hadialqattan/pycln
-    rev: v2.1.1 # Possible releases: https://github.com/hadialqattan/pycln/releases
+    rev: v2.6.0
     hooks:
       - id: pycln
 

pyproject.toml

@@ -3,33 +3,33 @@ name = "vector_lake"
 version = "0.0.5"
 description = "S3 vector database for bigdata"
 readme = "Readme.md"
-requires-python = ">=3.9,<4.0"
+requires-python = ">=3.11,<4.0"
 license = { text = "MIT" }
 keywords = ["vector", "database", "bigdata"]
 authors = [{ name = "Alexander Miasoiedov", email = "msoedov@gmail.com" }]
 maintainers = [{ name = "Alexander Miasoiedov", email = "msoedov@gmail.com" }]
 dependencies = [
-    "fire>=0.5.0,<0.6.0",
-    "loguru>=0.7.0,<0.8.0",
-    "numpy>=1.26.4,<2.0.0",
-    "pandas>=2.2.2,<3.0.0",
-    "pyarrow>=16.1.0,<17.0.0",
-    "pydantic>=1.10.15,<2.0.0",
-    "pytz>=2024.1,<2025.0",
-    "scikit-learn>=1.4.2,<2.0.0",
+    "fire>=0.7.1,<0.8.0",
+    "loguru>=0.7.3,<0.8.0",
+    "numpy>=2.0.0,<3.0.0",
+    "pandas>=3.0.0,<4.0.0",
+    "pyarrow>=23.0.0,<24.0.0",
+    "pydantic>=1.10.24,<2.0.0",
+    "pytz>=2025.2",
+    "scikit-learn>=1.8.0,<2.0.0",
 ]
 urls = { "Repository" = "https://github.com/msoedov/vector_lake" }
 
 [project.optional-dependencies]
-s3 = ["boto3>=1.34.115,<2.0.0"]
+s3 = ["boto3>=1.42.0,<2.0.0"]
 dev = [
-    "black>=24.4.2,<25.0.0",
-    "mypy>=1.10.0,<2.0.0",
-    "httpx>=0.27.0,<0.28.0",
-    "pytest>=8.2.0,<9.0.0",
-    "types-requests>=2.31.0,<3.0.0",
-    "pre-commit>=3.7.1,<4.0.0",
-    "hypothesis>=6.99.13,<7.0.0",
+    "black>=26.1.0,<27.0.0",
+    "mypy>=1.19.0,<2.0.0",
+    "httpx>=0.28.0,<0.29.0",
+    "pytest>=9.0.0,<10.0.0",
+    "types-requests>=2.32.0,<3.0.0",
+    "pre-commit>=4.5.0,<5.0.0",
+    "hypothesis>=6.151.0,<7.0.0",
     "pytest-benchmark==4.0.0",
 ]
 
@@ -54,25 +54,25 @@ packages = [{ include = "vector_lake", from = "." }]
 vector_lake = "vector_lake.__main__:entrypoint"
 
 [tool.poetry.dependencies]
-python = "^3.9"
-fire = "^0.5.0"
-loguru = "^0.7.0"
-numpy = "^1.26.4"
-pandas = "^2.2.2"
-pyarrow = "^16.1.0"
-pydantic = "^1.10.15"
-pytz = "^2024.1"
-scikit-learn = "^1.4.2"
-boto3 = { version = "^1.34.115", optional = true }
+python = "^3.11"
+fire = "^0.7.1"
+loguru = "^0.7.3"
+numpy = "^2.0.0"
+pandas = "^3.0.0"
+pyarrow = "^23.0.0"
+pydantic = "^1.10.24"
+pytz = "^2025.2"
+scikit-learn = "^1.8.0"
+boto3 = { version = "^1.42.0", optional = true }
 
 [tool.poetry.group.dev.dependencies]
-black = "^24.4.2"
-mypy = "^1.10.0"
-httpx = "^0.27.0"
-pytest = "^8.2.0"
-types-requests = "^2.31.0"
-pre-commit = "^3.7.1"
-hypothesis = "^6.99.13"
+black = "^26.1.0"
+mypy = "^1.19.0"
+httpx = "^0.28.0"
+pytest = "^9.0.0"
+types-requests = "^2.32.0"
+pre-commit = "^4.5.0"
+hypothesis = "^6.151.0"
 pytest-benchmark = "4.0.0"
 
 [tool.poetry.extras]

tests/__init__.py

@@ -33,7 +33,9 @@ def index_factory():
 @pytest.fixture(scope="session")
 def s3_index_factory():
     if not os.environ.get("RUN_S3_TESTS"):
-        pytest.skip("RUN_S3_TESTS is not set; skipping S3 integration tests by default.")
+        pytest.skip(
+            "RUN_S3_TESTS is not set; skipping S3 integration tests by default."
+        )
     pytest.importorskip("boto3")
     os.environ.setdefault("LOCALSTACK_ENDPOINT_URL", "http://localhost:4566")
     os.environ.setdefault("AWS_ACCESS_KEY_ID", "foo")

tests/conftest.py

@@ -0,0 +1,30 @@
+import numpy as np
+
+
+class TestS3Integration:
+    def test_end_to_end_s3(self, s3_index_factory):
+        index = s3_index_factory(
+            location="s3://unittest-vector-lake",
+            dimension=5,
+            approx_shards=243,
+            size=0,
+            force_clean=True,
+        )
+        vector = np.random.rand(1, 5)[0]
+        closest_vectors = index.query(vector, 4)
+        assert len(closest_vectors) == 0
+        index.add(vector, metadata={"id": 1}, document="unit test")
+        closest_vectors = index.query(vector, 4)
+        assert len(closest_vectors) == 1
+        index.persist()
+        # Reload the index
+        index = s3_index_factory(
+            location="s3://unittest-vector-lake",
+            dimension=5,
+            approx_shards=243,
+            size=0,
+        )
+        closest_vectors = index.query(vector, 4)
+        assert len(closest_vectors) == 1
+        assert np.array_equal(closest_vectors[0]["vector"], vector)
+        index.delete_remote()

tests/integration/__init__.py

@@ -3,8 +3,7 @@
 from vector_lake.core.index import make_granularity, make_nodes
 
 
-class TestUnit:
-    # Tests that the function works when D is 1
+class TestMakeGranularity:
     def test_D_is_1(self):
         D = 1
         M = 10
@@ -16,7 +15,6 @@ def test_D_is_1(self):
         assert all(level <= M for level in levels)
         assert sum(levels) == M
 
-    # Tests that D and M are positive integers
     def test_positive_integers(self):
         D = 3
         M = 10
@@ -25,15 +23,15 @@ def test_positive_integers(self):
         assert all(isinstance(level, int) for level in levels)
         assert all(level > 0 for level in levels)
 
-    # Tests that the function works with levels and num_shards as 1
+
+class TestMakeNodes:
     def test_levels_and_num_shards_as_1(self):
         levels = [1, 1]
         num_shards = 1
         nodes = make_nodes(levels, num_shards)
         assert nodes.shape == (1, 2)
         assert np.allclose(nodes, np.array([[0.0, 0.0]]))
 
-    # Tests that the function works with default values for levels and num_shards
     def test_default_values(self):
         levels = [2, 2]
         num_shards = 4
@@ -43,6 +41,8 @@ def test_default_values(self):
             nodes, np.array([[0.0, 0.0], [0.0, 0.5], [0.5, 0.0], [0.5, 0.5]])
         )
 
+
+class TestEndToEnd:
     def test_end_to_end(self, index_factory):
         index = index_factory(
             location="/tmp/cosine/empty",
@@ -68,8 +68,3 @@ def test_end_to_end(self, index_factory):
         closest_vectors = index.query(vector, 4)
         assert len(closest_vectors) == 1
         assert np.array_equal(closest_vectors[0]["vector"], vector)
-        return index
-
-    def test_end_to_end_s3(self, s3_index_factory):
-        index = self.test_end_to_end(s3_index_factory)
-        index.delete_remote()