Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion python/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ dependencies = [
"uvicorn>=0.24.0",
"pydantic>=2.0.0",
"openai>=1.52.0",
"lance>=0.17.0",
"pylance",
]
description = "Python bindings for the lance-graph Cypher engine"
authors = [{ name = "Lance Devs", email = "dev@lancedb.com" }]
Expand Down
58 changes: 20 additions & 38 deletions python/python/knowledge_graph/store.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
"""Persistent storage helpers built on Lance datasets (with Feather fallback)."""
"""Persistent storage helpers built on Lance datasets."""

from __future__ import annotations

Expand Down Expand Up @@ -46,42 +46,35 @@ def list_datasets(self) -> Dict[str, "Path"]:
datasets: Dict[str, Path] = {}
if not self._root.exists():
return datasets
valid_suffixes = {".lance", ".arrow"}
for child in self._root.iterdir():
if child.is_dir() and child.suffix == ".lance":
datasets[child.stem] = child
elif child.is_file() and child.suffix in valid_suffixes:
datasets[child.stem] = child
return datasets

def _dataset_path(self, name: str) -> "Path":
"""Create the canonical path for a dataset."""
safe_name = name.replace("/", "_")
suffix = ".lance" if self._get_lance() else ".arrow"
return self._root / f"{safe_name}{suffix}"
return self._root / f"{safe_name}.lance"

def _get_lance(self) -> Optional[ModuleType]:
def _get_lance(self) -> ModuleType:
if not self._lance_attempted:
self._lance_attempted = True
try:
module = import_module("lance")
except ImportError:
module = None
else:
has_writer = hasattr(module, "write_dataset")
has_loader = hasattr(module, "dataset")
if not (has_writer and has_loader):
LOGGER.warning(
"Installed `lance` package missing dataset APIs; "
"falling back to Feather storage."
)
module = None
self._lance = module
if module is None:
LOGGER.debug(
"Lance storage unavailable; using Feather files under %s.",
self._root,
except ImportError as e:
raise ImportError(
"Lance module is required but not installed. "
"Install it with: pip install pylance"
) from e

has_loader = hasattr(module, "dataset")
if not (has_loader):
raise ImportError(
"Installed `lance` package is missing required dataset APIs."
)
self._lance = module
if self._lance is None:
raise ImportError("Lance module failed to load")
return self._lance

def load_tables(
Expand All @@ -90,7 +83,6 @@ def load_tables(
) -> Mapping[str, "pa.Table"]:
"""Load Lance datasets as PyArrow tables."""
lance = self._get_lance()
use_lance = lance is not None

self.ensure_layout()
available = self.list_datasets()
Expand All @@ -101,13 +93,8 @@ def load_tables(
path = available.get(name, self._dataset_path(name))
if not path.exists():
raise FileNotFoundError(f"Dataset '{name}' not found at {path}")
if path.suffix == ".lance" and use_lance:
dataset = lance.dataset(str(path)) # type: ignore[union-attr]
table = dataset.scanner().to_table()
else:
import pyarrow.feather as feather

table = feather.read_table(str(path))
dataset = lance.dataset(str(path))
table = dataset.scanner().to_table()
tables[name] = table
return tables

Expand All @@ -123,10 +110,5 @@ def write_tables(self, tables: Mapping[str, "pa.Table"]) -> None:
f"Dataset '{name}' must be a pyarrow.Table (got {type(table)!r})"
)
path = self._dataset_path(name)
if path.suffix == ".lance" and lance is not None:
mode = "overwrite" if path.exists() else "create"
lance.write_dataset(table, str(path), mode=mode) # type: ignore[union-attr]
else:
import pyarrow.feather as feather

feather.write_feather(table, str(path))
mode = "overwrite" if path.exists() else "create"
lance.write_dataset(table, str(path), mode=mode)
Loading
Loading