from typing import List, Optional, Union, Dict, Any
import os
from .hyperstreamdb import Device as _Device
from .hyperstreamdb import Table as _RustTable
from .hyperstreamdb import Session as _RustSession
from .hyperstreamdb import *
def _has_torch():
try:
import torch
return True
except ImportError:
return False
class Device:
"""
HyperStreamDB Compute Device (CPU, CUDA, MPS, ROCm, Intel).
- **Torch Alignment** - Automatically aliases `cuda` to `rocm` on AMD hardware if `torch.version.hip` is detected.
"""
def __new__(cls, device: str = "cpu", index: Optional[int] = None):
device = device.lower()
# 1. Handle Torch-style alignment
if device.startswith("cuda"):
# If Torch is present and on AMD, or if native probing finds only ROCm
if _has_torch():
import torch
if getattr(torch.version, 'hip', None):
return _Device("rocm", index=index)
# If not using torch but only AMD hardware is present, alias cuda to rocm
if not _Device.is_available("cuda") and _Device.is_available("rocm"):
return _Device("rocm", index=index)
if (device.startswith("xpu") or device == "intel"):
return _Device("intel", index=index)
# 2. Native direct mapping
return _Device(device, index=index)
@staticmethod
def is_available(device_type: str) -> bool:
device_type = device_type.lower()
if device_type == "cuda":
# Torch compatibility: 'cuda' is true if either NVIDIA or AMD is present
return _Device.is_available("cuda") or _Device.is_available("rocm")
if device_type == "xpu":
return _Device.is_available("intel")
return _Device.is_available(device_type)
@staticmethod
def list_available_backends():
backends = ['cpu']
# Use precise strings for native listing
for b in ['cuda', 'rocm', 'mps', 'intel']:
if _Device.is_available(b):
backends.append(b)
return backends
@staticmethod
def auto_detect():
# 1. Check Torch first | **Torch Alignment** | ❌ No | ✅ ROCm-as-CUDA |
if _has_torch():
import torch
if torch.cuda.is_available():
if getattr(torch.version, 'hip', None):
return _Device("rocm")
return _Device("cuda")
# Check for Intel IPEX
if hasattr(torch, "xpu") and torch.xpu.is_available():
return _Device("intel")
# 2. Fallback to native probing
for b in ['cuda', 'rocm', 'mps', 'intel']:
if _Device.is_available(b):
return _Device(b)
return _Device('cpu')
@staticmethod
def deactivate():
_Device.deactivate()
_Device.backend = property(lambda self: self.type_name)
_Device.device_id = property(lambda self: self.index)
ComputeContext = Device
GPUContext = Device
from .embeddings import registry, EmbeddingFunction
import pandas as pd
try:
import pyarrow as pa
except ImportError:
pa = None
try:
import polars as pl
except ImportError:
pl = None
class IndexType:
"""
HyperStreamDB Indexing Algorithms.
"""
HNSW = "hnsw"
BM25 = "bm25"
BLOOM = "bloom"
BITMAP = "bitmap"
INVERTED = "inverted"
def _resolve_uri(uri: str) -> str:
if not uri.startswith(("s3://", "file://", "az://", "gs://", "http://", "https://")):
return f"file://{os.abspath(uri)}" if hasattr(os, "abspath") else uri
return uri
def open_table(uri: str, **kwargs) -> Table:
"""Open an existing HyperStreamDB table."""
return Table(uri, **kwargs)
[docs]
class Query:
"""
Fluent Query interface for HyperStreamDB.
"""
def __init__(self, table, filter_expr: Optional[str] = None):
self._table = table
self._filter = filter_expr
self._vector_filter = None
self._columns = None
[docs]
def filter(self, expr: str) -> 'Query':
"""Apply a SQL-like filter expression."""
if self._filter:
self._filter = f"({self._filter}) AND ({expr})"
else:
self._filter = expr
return self
[docs]
def vector_search(self, query: Union[List[float], str], column: Optional[str] = None, k: int = 10, **kwargs) -> 'Query':
"""
Apply a vector search filter.
Args:
query: The query vector (list of floats) or a string to be vectorized.
column: The vector column to search against.
k: Number of nearest neighbors to return.
**kwargs: Additional parameters (e.g., n_probe).
"""
self._vector_filter = {
"column": column,
"query": query,
"k": k,
**kwargs
}
return self
[docs]
def select(self, columns: List[str]) -> 'Query':
"""Select specific columns to return."""
self._columns = columns
return self
[docs]
def to_pandas(self, device: Optional[Any] = None):
"""Execute the query and return results as a Pandas DataFrame."""
return self._table.to_pandas(
filter=self._filter,
vector_filter=self._vector_filter,
columns=self._columns,
device=device
)
[docs]
def to_arrow(self, device: Optional[Any] = None):
"""Execute the query and return results as an Arrow Table."""
return self._table.to_arrow(
filter=self._filter,
vector_filter=self._vector_filter,
columns=self._columns,
device=device
)
[docs]
def execute(self, device: Optional[Any] = None, to_arrow: bool = False):
"""Execute the query and return results as a Pandas DataFrame (default) or Arrow Table."""
if to_arrow:
return self.to_arrow(device)
return self.to_pandas(device)
[docs]
class Table:
"""
HyperStreamDB Table — Apache Iceberg/Parquet-compatible columnar vector store.
**Default behaviour (v0.4.1+)**
- ``index_all = False`` — Vector indexes are *not* built automatically.
Call ``table.index_all = True`` or ``table.add_index(column, 'hnsw')``
to enable indexing for a specific session or column.
- ``autocommit = False`` — Writes accumulate in an in-memory buffer.
Call ``table.commit()`` (or ``await table.commit_async()``) to persist
data to Parquet and advance the Iceberg snapshot.
These defaults exist for performance: automatic indexing previously caused
silent 15-18 s HNSW build latency on every ``commit()`` for tables with
vector columns, even when the user had not requested an index.
Args:
uri: Table location (``file:///path`` or cloud URI).
inner_table: Internal — do not pass directly.
device: Optional compute device for GPU-accelerated index builds.
index_all: Enable automatic indexing of all compatible columns.
Defaults to ``False``. Set ``True`` to restore legacy behaviour.
primary_key: Column name (or list) to use as primary key.
explain: If ``True``, return query plans instead of results.
"""
def __init__(self, uri: str, inner_table: Optional[_RustTable] = None, device: Optional[Any] = None, index_all: bool = False, primary_key: Optional[str] = None, explain: bool = False):
uri = _resolve_uri(uri)
self.explain = explain
if inner_table:
self._inner = inner_table
else:
self._inner = _RustTable(uri, device=device)
self._inner.set_index_all(index_all)
if primary_key:
if isinstance(primary_key, str):
self._inner.set_primary_key([primary_key])
else:
self._inner.set_primary_key(list(primary_key))
self._embedding_configs = {}
[docs]
@classmethod
def create(cls, uri: str, schema, device: Optional[Any] = None) -> 'Table':
"""Create a new table with an explicit schema."""
uri = _resolve_uri(uri)
return cls(uri, inner_table=_RustTable.create(uri, schema, device=device))
[docs]
@classmethod
def create_partitioned(cls, uri: str, schema, partition_spec: Dict[str, Any], device: Optional[Any] = None) -> 'Table':
"""Create a new table with an explicit schema and partitioning."""
uri = _resolve_uri(uri)
return cls(uri, inner_table=_RustTable.create_partitioned(uri, schema, partition_spec, device=device))
[docs]
@classmethod
def register_external(cls, uri: str, iceberg_metadata_uri: str, device: Optional[Any] = None) -> 'Table':
"""Register an existing Iceberg table."""
uri = _resolve_uri(uri)
return cls(uri, inner_table=_RustTable.register_external(uri, iceberg_metadata_uri), device=device)
[docs]
def define_embedding(self, column: str, function: Union[str, EmbeddingFunction], vector_column: Optional[str] = None):
"""
Link a source column to an embedding function for automatic vectorization.
Args:
column: The source text column.
function: Registered function name or EmbeddingFunction instance.
vector_column: Target vector column name (defaults to {column}_vector).
"""
self._embedding_configs[column] = {
"function": function,
"vector_column": vector_column or f"{column}_vector"
}
[docs]
def write(self, data: Any, device: Optional[Any] = None, mode: str = "append"):
"""
Write data to the table, automatically generating embeddings for configured columns.
Args:
data: pandas.DataFrame, pyarrow.Table, polars.DataFrame, numpy.ndarray, torch.Tensor, or List[Dict].
device: Optional Device for GPU acceleration.
mode: 'append' (default) or 'overwrite' (clears table first).
"""
if mode == "overwrite":
self.truncate()
if isinstance(data, pd.DataFrame):
return self._write_pandas(data, device=device)
elif pa and isinstance(data, pa.Table):
return self._write_arrow(data, device=device)
elif pl and isinstance(data, pl.DataFrame):
return self._write_polars(data, device=device)
elif isinstance(data, list):
return self._write_list(data, device=device)
else:
try:
import numpy as np
if isinstance(data, np.ndarray):
return self.write(pd.DataFrame(data), device=device)
import torch
if isinstance(data, torch.Tensor):
return self.write(pd.DataFrame(data.detach().cpu().numpy()), device=device)
except ImportError:
pass
raise TypeError(f"Unsupported data type for write: {type(data)}")
[docs]
def insert(self, data: Any, device: Optional[Any] = None):
"""Alias for write() for compatibility with common vector DB APIs."""
return self.write(data, device=device)
[docs]
def write_pandas(self, df: pd.DataFrame, device: Optional[Any] = None):
"""High-level Pandas ingestion with auto-vectorization."""
return self._write_pandas(df, device=device)
[docs]
def write_arrow(self, table: 'pa.Table', device: Optional[Any] = None):
"""High-level Arrow ingestion with auto-vectorization."""
return self._write_arrow(table, device=device)
[docs]
def upsert(self, data: Any, key_column: Union[str, List[str]], mode: str = "merge_on_read", device: Optional[Any] = None):
"""Update or insert data using a key column (or list of columns) to avoid duplicates."""
from .hyperstreamdb import PyMergeMode
# Map string mode to Enum
enum_mode = PyMergeMode.MergeOnRead
if mode.lower() == "merge_on_write":
enum_mode = PyMergeMode.MergeOnWrite
if isinstance(data, pd.DataFrame):
processed_df = self._auto_vectorize(data)
# If key_column is a list, join it with commas for the Rust side (or update Rust to take list)
if isinstance(key_column, list):
key_str = ",".join(key_column)
else:
key_str = key_column
return self._inner.merge_pandas(processed_df, key_str, enum_mode, device=device)
df = pd.DataFrame(data)
return self.upsert(df, key_column, mode, device=device)
[docs]
def commit(self):
"""Commit temporary segments to the table."""
return self._inner.commit()
[docs]
def truncate(self):
"""Clear all data from the table while keeping the schema."""
return self._inner.truncate()
[docs]
def vacuum(self, retention_versions: int = 1):
"""
Physically delete unreferenced data and manifest files to reclaim space.
Args:
retention_versions: Number of snapshots to keep (default 1).
"""
return self._inner.vacuum(retention_versions)
@property
def autocommit(self) -> bool:
"""Get or set the autocommit state of the table."""
return self._inner.autocommit
@autocommit.setter
def autocommit(self, value: bool):
self._inner.autocommit = value
[docs]
def wait_for_background_tasks(self):
"""Wait for all background tasks (like index building) to complete."""
return self._inner.wait_for_background_tasks()
[docs]
def delete(self, filter: str):
"""Delete rows matching the filter expression."""
return self._inner.delete(filter)
def _write_pandas(self, df: pd.DataFrame, device: Optional[Any] = None):
processed_df = self._auto_vectorize(df)
return self._inner.write_pandas(processed_df, device=device)
def _write_arrow(self, table: 'pa.Table', device: Optional[Any] = None):
if self._embedding_configs:
df = table.to_pandas()
return self._write_pandas(df, device=device)
if pa and isinstance(table, pa.RecordBatch):
from pyarrow import Table as paTable
table = paTable.from_batches([table])
return self._inner.write_arrow(table, device=device)
def _write_polars(self, df: 'pl.DataFrame', device: Optional[Any] = None):
if self._embedding_configs:
pandas_df = df.to_pandas()
return self._write_pandas(pandas_df, device=device)
return self._inner.write_arrow(df.to_arrow(), device=device)
def _write_list(self, data: List[Any], device: Optional[Any] = None):
if not data:
return
# Check if first element is an Arrow object
first = data[0]
if (pa and isinstance(first, (pa.RecordBatch, pa.Table))):
from pyarrow import Table as paTable
if isinstance(first, pa.RecordBatch):
combined = paTable.from_batches(data)
else:
combined = pa.concat_tables(data)
return self._write_arrow(combined, device=device)
# Default to pandas for List[Dict] or other types
df = pd.DataFrame(data)
return self._write_pandas(df, device=device)
def _auto_vectorize(self, data: Union[pd.DataFrame, List[Dict[str, Any]]]):
if not self._embedding_configs:
return data
if isinstance(data, pd.DataFrame):
import numpy as np
df = data.copy()
for col, config in self._embedding_configs.items():
if col in df.columns:
func = config["function"]
if isinstance(func, str):
func = registry.get(func)
if func:
vector_col = config["vector_column"]
embeddings = func(df[col].tolist())
# Enforce Float32 for vector compatibility
if isinstance(embeddings, np.ndarray):
embeddings = embeddings.astype(np.float32)
df[vector_col] = list(embeddings)
return df
# Large list branch omitted for brevity, logic is similar (use pandas path)
return data
def _prepare_vector_filter(self, vector_filter: Optional[Union[Dict[str, Any], List[float]]], **kwargs) -> Optional[Dict[str, Any]]:
if vector_filter is None:
return None
# 1. Handle vector_filter as a list (simplified search)
if not isinstance(vector_filter, dict):
column = "embedding"
if self._embedding_configs:
column = list(self._embedding_configs.values())[0]["vector_column"]
vector_filter = {"column": column, "query": vector_filter}
# 2. Add extra kwargs (k, n_probe) to vector_filter if present
if kwargs:
vector_filter.update(kwargs)
if "k" not in vector_filter:
vector_filter["k"] = 10
# Ensure column is set (e.g. if fluent API sent column=None)
if vector_filter.get("column") is None:
column = "embedding"
if self._embedding_configs:
column = list(self._embedding_configs.values())[0]["vector_column"]
vector_filter["column"] = column
# Auto-vectorize string query
if "query" in vector_filter and isinstance(vector_filter["query"], str):
# Try to find a matching embedding function
target_col = vector_filter.get("column")
func = None
# 1. Check if we have an explicit config for this vector column
for src_col, config in self._embedding_configs.items():
if config["vector_column"] == target_col:
func = config["function"]
break
# 2. If not, check if any registered function matches the column name
if not func:
func = registry.get(target_col)
if func:
if isinstance(func, str):
func = registry.get(func)
if func:
# Vectorize the query string
vector_filter["query"] = func([vector_filter["query"]])[0].tolist()
if self.explain:
print(f"[Explain] Vectorized query using device: {target_col}")
return vector_filter
[docs]
def to_pandas(self, filter: Optional[str] = None, vector_filter: Optional[Union[Dict[str, Any], List[float]]] = None, columns: Optional[List[str]] = None, device: Optional[Any] = None, **kwargs):
"""
Read table to Pandas with auto-vectorization of search queries and flexible parameters.
Parameters:
filter: Optional scalar WHERE clause (e.g., "category = 'news'")
vector_filter: Dict with vector search params:
- column: str (required) - vector column name
- query: list (required) - query vector
- k: int (required) - number of results
- metric: str (optional) - 'l2'|'cosine'|'innerproduct'|'l1'|'hamming'|'jaccard' (default: l2)
- ef_search: int (optional) - HNSW ef parameter for tuning
- probes: int (optional) - IVF probes parameter for tuning
columns: Optional list of column names to select
device: Optional compute device (GPU/CPU)
**kwargs: Extra params (merged into vector_filter if present)
Example::
# Vector search with cosine metric
df = table.to_pandas(vector_filter={
"column": "embedding",
"query": [1.0, 2.0, 3.0],
"k": 5,
"metric": "cosine",
"ef_search": 200 # Tune HNSW search quality
})
"""
vf = self._prepare_vector_filter(vector_filter, **kwargs)
if self.explain:
# Call native Rust explain logic
print(self._inner.explain(filter, vf))
filtered_kwargs = {k: v for k, v in kwargs.items() if k not in ["k", "n_probe", "column"]}
return self._inner.to_pandas(filter, vf, columns, device=device, **filtered_kwargs)
[docs]
def to_arrow(self, filter: Optional[str] = None, vector_filter: Optional[Union[Dict[str, Any], List[float]]] = None, columns: Optional[List[str]] = None, device: Optional[Any] = None, **kwargs):
"""
Read table to Arrow Table with auto-vectorization of search queries and flexible parameters.
Parameters:
filter: Optional scalar WHERE clause (e.g., "category = 'news'")
vector_filter: Dict with vector search params:
- column: str (required) - vector column name
- query: list (required) - query vector
- k: int (required) - number of results
- metric: str (optional) - 'l2'|'cosine'|'innerproduct'|'l1'|'hamming'|'jaccard' (default: l2)
- ef_search: int (optional) - HNSW ef parameter for tuning
- probes: int (optional) - IVF probes parameter for tuning
columns: Optional list of column names to select
device: Optional compute device (GPU/CPU)
**kwargs: Extra params (merged into vector_filter if present)
"""
if "filter" in kwargs and filter is None:
filter = kwargs.pop("filter")
vf = self._prepare_vector_filter(vector_filter, **kwargs)
# to_arrow in Rust doesn't currently take **kwargs
return self._inner.to_arrow(filter, vf, columns, device=device)
[docs]
def sql(self, query: str) -> Any:
"""
Execute a SQL query against the table.
The table is registered as 't'.
"""
return self._inner.execute_sql(query)
[docs]
def query(self) -> Query:
"""Start a fluent query."""
return Query(self)
[docs]
def read(self, filter: Optional[str] = None, vector_filter: Optional[Union[Dict[str, Any], List[float]]] = None, columns: Optional[List[str]] = None, device: Optional[Any] = None, **kwargs):
"""
Read table to Arrow Table (alias for to_arrow).
"""
return self.to_arrow(filter, vector_filter, columns, device=device, **kwargs)
[docs]
def vector_search(self, column: str, query: List[float], k: int = 10, filter: Optional[str] = None, columns: Optional[List[str]] = None, device: Optional[Any] = None, **kwargs):
"""Backward compatibility alias for to_pandas with vector filter."""
vf = {"column": column, "query": query, "k": k}
vf.update(kwargs)
return self.to_pandas(filter=filter, vector_filter=vf, columns=columns, device=device)
[docs]
def search(self, column: str, query: List[float], k: int = 10, filter: Optional[str] = None, columns: Optional[List[str]] = None, device: Optional[Any] = None, **kwargs):
"""Alias for vector_search."""
return self.vector_search(column, query, k, filter, columns, device, **kwargs)
[docs]
def filter(self, expr: Optional[str] = None, vector_filter: Optional[Union[Dict[str, Any], List[float]]] = None, **kwargs) -> 'Query':
"""
Start a fluent query or apply immediate filters.
"""
if "filter" in kwargs and expr is None:
expr = kwargs.pop("filter")
q = Query(self, expr)
if vector_filter is not None:
if isinstance(vector_filter, list):
q.vector_search(vector_filter, **kwargs)
elif isinstance(vector_filter, dict):
# Merge dict into Query state
q._vector_filter = vector_filter
if kwargs:
q._vector_filter.update(kwargs)
elif kwargs:
# Assume kwargs refer to search params if vector_filter was missing but k was provided?
# Actually better to be explicit: table.filter(vector_filter=v, k=5)
pass
return q
@property
def primary_key(self):
"""Get the current primary key column."""
return self._inner.get_primary_key()
@primary_key.setter
def primary_key(self, columns: Union[str, List[str]]):
"""Set the primary key column(s)."""
if isinstance(columns, str):
self._inner.set_primary_key([columns])
else:
self._inner.set_primary_key(list(columns))
@property
def index_all(self):
"""
Whether to build HNSW/BM25 indexes for all compatible columns on commit.
Defaults to ``False`` (opt-in). Setting this to ``True`` triggers
background index builds after every ``commit()`` call — useful when
you want fast ANN search but be aware of the additional commit latency
(~15 s per 100 K rows with 768-dim vectors on CPU).
For selective indexing, prefer ``table.add_index(column, 'hnsw')``.
"""
return self._inner.get_index_all()
@index_all.setter
def index_all(self, value):
self._inner.set_index_all(value)
@property
def row_count(self) -> int:
"""Get total row count in the table."""
return self._inner.get_table_statistics().row_count
@property
def statistics(self):
"""Get full table statistics."""
return self._inner.get_table_statistics()
[docs]
def add_index_columns(self, columns: List[str], tokenizer: Optional[str] = None):
"""
Add columns to the indexing configuration.
Args:
columns: List of column names to index.
tokenizer: Optional tokenizer name from the registry.
"""
return self._inner.add_index_columns(columns, tokenizer)
[docs]
def set_index_config(self, column: str, enabled: bool = True, tokenizer: Optional[str] = None, device: Optional[str] = None):
"""
Set indexing configuration for a specific column.
(Legacy compatibility wrapper)
"""
if not enabled:
return self.drop_index(column)
config = {"type": "hnsw"}
if tokenizer: config["tokenizer"] = tokenizer
if device: config["build_device"] = device
return self.add_index(column, config)
[docs]
def set_index_columns(self, config: Dict[str, Union[str, List[Union[str, Dict[str, Any]]], Dict[str, Any]]]):
"""
Update indexing specifications for multiple columns at once.
Supports both simple strings and advanced configuration dictionaries.
Example::
table.set_index_columns({
"embedding": IndexType.HNSW,
"content": ["hnsw", "bm25"],
"category": "bitmap"
})
"""
return self._inner.set_index_columns(config)
[docs]
def add_index(self, column: str, algorithm: Union[str, Dict[str, Any]] = "hnsw", **kwargs):
"""
Add an indexing strategy to a column.
"""
if isinstance(algorithm, str):
algorithm = {"type": algorithm}
if kwargs:
# Map 'device' to 'build_device' for consistency with set_index_config
if 'device' in kwargs:
kwargs['build_device'] = kwargs.pop('device')
algorithm.update(kwargs)
return self._inner.add_index(column, algorithm)
[docs]
def drop_index(self, column: str):
"""
Remove all indexing strategies from a column.
"""
return self._inner.drop_index(column)
[docs]
def add_primary_key(self, column: str):
"""
Atomically add a column to the primary key.
This performs a validation check for duplicates across all existing data.
If validation fails, the change is NOT committed.
"""
return self._inner.add_primary_key(column)
[docs]
def drop_primary_key(self, column: str):
"""
Atomically remove a column from the primary key.
"""
return self._inner.drop_primary_key(column)
[docs]
def set_sort_order(self, columns: List[str], ascending: List[bool]):
"""Set the table's default sort order for future data writes."""
return self._inner.replace_sort_order(columns, ascending)
[docs]
def set_partition_spec(self, spec: List[Dict[str, Any]]):
"""
Update the table's partition specification.
Args:
spec: List of partition fields, each being a dict with:
- source_id: int (or source_ids: List[int])
- name: str
- transform: str
- field_id: int (optional)
"""
from .hyperstreamdb import PartitionField
fields = []
for item in spec:
if isinstance(item, dict):
# Handle both 'source_id' (singular) and 'source_ids' (plural) for flexibility
source_ids = item.get("source_ids")
if source_ids is None:
sid = item.get("source_id")
source_ids = [sid] if sid is not None else []
fields.append(PartitionField(
source_ids=source_ids,
name=item["name"],
transform=item["transform"],
field_id=item.get("field_id")
))
else:
fields.append(item)
return self._inner.update_spec(fields)
def __getattr__(self, name):
"""Delegate other calls to the Rust implementation."""
return getattr(self._inner, name)
def __repr__(self):
return f"HyperStreamTable(uri={self._inner.table_uri()})"
[docs]
class Session:
"""
HyperStreamDB Query Session with integration for Python Table objects.
"""
def __init__(self, memory_mb: Optional[int] = None):
self._inner = _RustSession(memory_mb)
[docs]
def register(self, name: str, table: Union[Table, _RustTable]):
"""Register a table in the session for SQL queries."""
if hasattr(table, "_inner"):
# Unwrap Python Table to get the Rust implementation
return self._inner.register(name, table._inner)
return self._inner.register(name, table)
[docs]
def sql(self, query: str) -> Any:
"""Execute a SQL query against the table (registered as 't')."""
return self._inner.sql(query)