feat: add list_indices to the async api (#1074)

This commit is contained in:
Weston Pace
2024-03-12 14:41:21 -07:00
parent 9031ec6878
commit b6a522d483
14 changed files with 233 additions and 16 deletions

View File

@@ -1,4 +1,4 @@
from typing import Dict, Optional
from typing import Dict, List, Optional
import pyarrow as pa
@@ -39,6 +39,11 @@ class Table:
async def checkout(self, version): ...
async def checkout_latest(self): ...
async def restore(self): ...
async def list_indices(self) -> List[IndexConfig]: ...
class IndexConfig:
index_type: str
columns: List[str]
async def connect(
uri: str,

View File

@@ -3,6 +3,9 @@ from typing import Optional
from ._lancedb import (
Index as LanceDbIndex,
)
from ._lancedb import (
IndexConfig,
)
class BTree(object):
@@ -155,3 +158,6 @@ class IvfPq(object):
max_iterations=max_iterations,
sample_rate=sample_rate,
)
__all__ = ["BTree", "IvfPq", "IndexConfig"]

View File

@@ -59,7 +59,7 @@ if TYPE_CHECKING:
from ._lancedb import Table as LanceDBTable
from .db import LanceDBConnection
from .index import BTree, IvfPq
from .index import BTree, IndexConfig, IvfPq
pd = safe_import_pandas()
@@ -2409,3 +2409,9 @@ class AsyncTable:
out state and the read_consistency_interval, if any, will apply.
"""
await self._inner.restore()
async def list_indices(self) -> IndexConfig:
"""
List all indices that have been created with Self::create_index
"""
return await self._inner.list_indices()

View File

@@ -41,6 +41,10 @@ async def test_create_scalar_index(some_table: AsyncTable):
await some_table.create_index("id")
# Can recreate if replace=True
await some_table.create_index("id", replace=True)
indices = await some_table.list_indices()
assert len(indices) == 1
assert indices[0].index_type == "BTree"
assert indices[0].columns == ["id"]
# Can't recreate if replace=False
with pytest.raises(RuntimeError, match="already exists"):
await some_table.create_index("id", replace=False)
@@ -59,3 +63,7 @@ async def test_create_vector_index(some_table: AsyncTable):
await some_table.create_index("vector", replace=False)
# Can also specify index type
await some_table.create_index("vector", config=IvfPq(num_partitions=100))
indices = await some_table.list_indices()
assert len(indices) == 1
assert indices[0].index_type == "IvfPq"
assert indices[0].columns == ["vector"]

View File

@@ -85,3 +85,25 @@ impl Index {
})
}
}
#[pyclass(get_all)]
/// A description of an index currently configured on a column
pub struct IndexConfig {
/// The type of the index
pub index_type: String,
/// The columns in the index
///
/// Currently this is always a list of size 1. In the future there may
/// be more columns to represent composite indices.
pub columns: Vec<String>,
}
impl From<lancedb::index::IndexConfig> for IndexConfig {
fn from(value: lancedb::index::IndexConfig) -> Self {
let index_type = format!("{:?}", value.index_type);
Self {
index_type,
columns: value.columns,
}
}
}

View File

@@ -14,7 +14,7 @@
use connection::{connect, Connection};
use env_logger::Env;
use index::Index;
use index::{Index, IndexConfig};
use pyo3::{pymodule, types::PyModule, wrap_pyfunction, PyResult, Python};
use table::Table;
@@ -33,6 +33,7 @@ pub fn _lancedb(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_class::<Connection>()?;
m.add_class::<Table>()?;
m.add_class::<Index>()?;
m.add_class::<IndexConfig>()?;
m.add_function(wrap_pyfunction!(connect, m)?)?;
m.add("__version__", env!("CARGO_PKG_VERSION"))?;
Ok(())

View File

@@ -11,7 +11,10 @@ use pyo3::{
};
use pyo3_asyncio::tokio::future_into_py;
use crate::{error::PythonErrorExt, index::Index};
use crate::{
error::PythonErrorExt,
index::{Index, IndexConfig},
};
#[pyclass]
pub struct Table {
@@ -127,6 +130,19 @@ impl Table {
})
}
pub fn list_indices(self_: PyRef<'_, Self>) -> PyResult<&PyAny> {
let inner = self_.inner_ref()?.clone();
future_into_py(self_.py(), async move {
Ok(inner
.list_indices()
.await
.infer_error()?
.into_iter()
.map(IndexConfig::from)
.collect::<Vec<_>>())
})
}
pub fn __repr__(&self) -> String {
match &self.inner {
None => format!("ClosedTable({})", self.name),