Back to using strings instead of enum for index type

lancedb · Nov 17, 2023 · e160ad8 · e160ad8
1 parent aa8308a
commit e160ad8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 13 deletions.
diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py
@@ -22,7 +22,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime, timedelta
-from enum import Enum
 from functools import lru_cache
 from pathlib import Path
 from typing import (
@@ -32,6 +31,7 @@
     Iterable,
     Iterator,
     List,
+    Literal,
     Optional,
     TypedDict,
     Union,
@@ -78,10 +78,6 @@
     from .progress import FragmentWriteProgress
 
 
-class ScalarIndexType(Enum):
-    BTREE = 0
-
-
 class LanceDataset(pa.dataset.Dataset):
     """A dataset in Lance format where the data is stored at the given uri."""
 
@@ -702,7 +698,7 @@ def cleanup_old_versions(
     def create_scalar_index(
         self,
         column: str,
-        index_type: ScalarIndexType,
+        index_type: Literal["BTREE"],
         name: Optional[str] = None,
         *,
         replace: bool = True,
@@ -783,12 +779,11 @@ def create_scalar_index(
         .. code-block:: python
 
             import lance
-            from lance.dataset import ScalarIndexType
 
             dataset = lance.dataset("/tmp/images.lance")
             dataset.create_index(
                 "category",
-                ScalarIndexType.BTREE,
+                "BTREE",
             )
 
         Experimental Status:
@@ -822,15 +817,16 @@ def create_scalar_index(
                 f"Scalar index column {column} must be int, float, bool, or str"
             )
 
-        if index_type != ScalarIndexType.BTREE:
+        index_type = index_type.upper()
+        if index_type != "BTREE":
             raise NotImplementedError(
                 (
-                    "Only 'ScalarIndexType.BTREE' is supported for ",
+                    'Only "BTREE" is supported for ',
                     f"index_type.  Received {index_type}",
                 )
             )
 
-        self._ds.create_index([column], index_type.name, name, replace)
+        self._ds.create_index([column], index_type, name, replace)
 
     def create_index(
         self,

diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py
@@ -18,7 +18,6 @@
 import numpy as np
 import pyarrow as pa
 import pytest
-from lance.dataset import ScalarIndexType
 from lance.vector import vec_to_table
 
 
@@ -55,7 +54,7 @@ def indexed_dataset(tmp_path):
         num_partitions=4,
         num_sub_vectors=2,
     )
-    dataset.create_scalar_index("meta", index_type=ScalarIndexType.BTREE)
+    dataset.create_scalar_index("meta", index_type="BTREE")
     return dataset