PM-pinou
252 lines
9.2 KiB
Python
252 lines
9.2 KiB
Python
"""Thread-safe in-memory session store for datasets and results.
|
||
|
||
Singleton pattern with RLock guard. Stores:
|
||
- Datasets: dataset_id -> (LazyFrame, schema, metadata)
|
||
- Cluster results: result_id -> (labels, n_clusters, params, parent_dataset_id)
|
||
- Feature results: result_id -> (features, parent_cluster_result_id)
|
||
"""
|
||
import threading
|
||
import gc
|
||
import uuid
|
||
from typing import Any, Optional
|
||
|
||
import polars as pl
|
||
|
||
|
||
# Approximate per-element byte sizes for memory estimation
|
||
_DTYPE_BYTE_SIZES = {
|
||
# Polars dtype strings
|
||
'Int8': 1, 'Int16': 2, 'Int32': 4, 'Int64': 8,
|
||
'UInt8': 1, 'UInt16': 2, 'UInt32': 4, 'UInt64': 8,
|
||
'Float32': 4, 'Float64': 8,
|
||
'Boolean': 1, 'Bool': 1,
|
||
'Utf8': 50, 'String': 50, 'Categorical': 8,
|
||
'Date': 4, 'Datetime': 8, 'Time': 8, 'Duration': 8,
|
||
'List': 16, 'Object': 16, 'Null': 0,
|
||
# Polars __str__ representations
|
||
'i32': 4, 'i64': 8, 'u32': 4, 'u64': 8,
|
||
'f32': 4, 'f64': 8, 'str': 50, 'bool': 1,
|
||
}
|
||
|
||
|
||
def _estimate_dtype_bytes(dtype: Any) -> int:
|
||
"""Guess per-element byte size from a Polars dtype."""
|
||
key = str(dtype)
|
||
for prefix, size in _DTYPE_BYTE_SIZES.items():
|
||
if key.startswith(prefix) or key == prefix:
|
||
return size
|
||
return 8 # fallback
|
||
|
||
|
||
def _generate_id(prefix: str = 'ds') -> str:
|
||
"""Generate a short unique identifier."""
|
||
return f"{prefix}_{uuid.uuid4().hex[:12]}"
|
||
|
||
|
||
class SessionStore:
|
||
"""Singleton in-memory store for datasets and analysis results.
|
||
|
||
Thread-safe via RLock. Datasets store Polars LazyFrames (query plans,
|
||
not materialised data), so clone/shallow-copy is effectively free.
|
||
"""
|
||
|
||
_instance: Optional['SessionStore'] = None
|
||
_lock: threading.RLock = threading.RLock()
|
||
|
||
def __new__(cls) -> 'SessionStore':
|
||
if cls._instance is None:
|
||
with cls._lock:
|
||
if cls._instance is None:
|
||
cls._instance = super().__new__(cls)
|
||
cls._instance._stores: dict[str, dict[str, Any]] = {}
|
||
cls._instance._datalock = threading.RLock()
|
||
return cls._instance
|
||
|
||
# ── dataset operations ──────────────────────────────────────────────
|
||
|
||
def store_dataset(
|
||
self,
|
||
dataset_id: str,
|
||
lazyframe: pl.LazyFrame,
|
||
schema: Optional[dict] = None,
|
||
metadata: Optional[dict] = None,
|
||
parent_id: Optional[str] = None,
|
||
) -> None:
|
||
"""Store a LazyFrame under *dataset_id*.
|
||
|
||
Args:
|
||
dataset_id: Unique key for the dataset.
|
||
lazyframe: Polars LazyFrame (query plan, not materialised).
|
||
schema: Column schema dict {name: dtype_string}.
|
||
metadata: Optional dict (row_count, file_count, …).
|
||
parent_id: If this is a derived dataset (filter/clone/…), the
|
||
source dataset_id.
|
||
"""
|
||
with self._datalock:
|
||
self._stores[dataset_id] = {
|
||
'type': 'dataset',
|
||
'lazyframe': lazyframe,
|
||
'schema': schema or {},
|
||
'metadata': metadata or {},
|
||
'parent_id': parent_id,
|
||
}
|
||
|
||
def get_dataset(self, dataset_id: str) -> Optional[dict]:
|
||
"""Retrieve a stored dataset entry (or *None*)."""
|
||
with self._datalock:
|
||
entry = self._stores.get(dataset_id)
|
||
if entry and entry.get('type') == 'dataset':
|
||
return entry
|
||
return None
|
||
|
||
def drop_dataset(self, dataset_id: str) -> bool:
|
||
"""Remove a dataset from the store and force garbage collection.
|
||
|
||
Returns True if the dataset existed.
|
||
"""
|
||
with self._datalock:
|
||
existed = dataset_id in self._stores and self._stores[dataset_id].get('type') == 'dataset'
|
||
if existed:
|
||
del self._stores[dataset_id]
|
||
if existed:
|
||
gc.collect()
|
||
return existed
|
||
|
||
def clone_dataset(self, dataset_id: str) -> Optional[str]:
|
||
"""Shallow clone — LazyFrame is a query plan so no data is copied.
|
||
|
||
Returns the new dataset_id, or *None* if source does not exist.
|
||
"""
|
||
with self._datalock:
|
||
entry = self.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return None
|
||
new_id = _generate_id('clone')
|
||
self._stores[new_id] = {
|
||
'type': 'dataset',
|
||
'lazyframe': entry['lazyframe'], # same plan pointer
|
||
'schema': dict(entry['schema']),
|
||
'metadata': dict(entry.get('metadata', {})),
|
||
'parent_id': dataset_id,
|
||
}
|
||
return new_id
|
||
|
||
def list_datasets(self) -> list[dict]:
|
||
"""Return summary dict for every stored dataset."""
|
||
with self._datalock:
|
||
results = []
|
||
for ds_id, entry in self._stores.items():
|
||
if entry.get('type') != 'dataset':
|
||
continue
|
||
meta = entry.get('metadata', {})
|
||
results.append({
|
||
'dataset_id': ds_id,
|
||
'parent_id': entry.get('parent_id'),
|
||
'row_count': meta.get('row_count'),
|
||
'file_count': meta.get('file_count'),
|
||
'columns': list(entry.get('schema', {}).keys()),
|
||
'column_count': len(entry.get('schema', {})),
|
||
})
|
||
return results
|
||
|
||
def memory_estimate_mb(self, dataset_id: str) -> float:
|
||
"""Rough memory estimate based on schema × row_count."""
|
||
with self._datalock:
|
||
entry = self.get_dataset(dataset_id)
|
||
if entry is None:
|
||
return 0.0
|
||
schema = entry.get('schema', {})
|
||
row_count = entry.get('metadata', {}).get('row_count', 0)
|
||
if not row_count or not schema:
|
||
return 0.0
|
||
bytes_per_row = sum(_estimate_dtype_bytes(dtype) for dtype in schema.values())
|
||
return (bytes_per_row * row_count) / (1024 * 1024)
|
||
|
||
# ── cluster-result operations ───────────────────────────────────────
|
||
|
||
def store_cluster_result(
|
||
self,
|
||
result_id: str,
|
||
labels: list[int],
|
||
n_clusters: int,
|
||
params: dict[str, Any],
|
||
parent_dataset_id: str,
|
||
) -> None:
|
||
"""Store a clustering result."""
|
||
with self._datalock:
|
||
self._stores[result_id] = {
|
||
'type': 'cluster_result',
|
||
'labels': labels,
|
||
'n_clusters': n_clusters,
|
||
'params': params,
|
||
'parent_dataset_id': parent_dataset_id,
|
||
}
|
||
|
||
def get_cluster_result(self, result_id: str) -> Optional[dict]:
|
||
"""Retrieve a stored cluster result."""
|
||
with self._datalock:
|
||
entry = self._stores.get(result_id)
|
||
if entry and entry.get('type') == 'cluster_result':
|
||
return entry
|
||
return None
|
||
|
||
def drop_cluster_result(self, result_id: str) -> bool:
|
||
"""Remove a cluster result."""
|
||
with self._datalock:
|
||
existed = result_id in self._stores and self._stores[result_id].get('type') == 'cluster_result'
|
||
if existed:
|
||
del self._stores[result_id]
|
||
return existed
|
||
|
||
# ── feature-result operations ───────────────────────────────────────
|
||
|
||
def store_feature_result(
|
||
self,
|
||
result_id: str,
|
||
features: list[dict[str, Any]],
|
||
parent_cluster_result_id: str,
|
||
) -> None:
|
||
"""Store extracted features per cluster."""
|
||
with self._datalock:
|
||
self._stores[result_id] = {
|
||
'type': 'feature_result',
|
||
'features': features,
|
||
'parent_cluster_result_id': parent_cluster_result_id,
|
||
}
|
||
|
||
def get_feature_result(self, result_id: str) -> Optional[dict]:
|
||
"""Retrieve a stored feature result."""
|
||
with self._datalock:
|
||
entry = self._stores.get(result_id)
|
||
if entry and entry.get('type') == 'feature_result':
|
||
return entry
|
||
return None
|
||
|
||
def drop_feature_result(self, result_id: str) -> bool:
|
||
"""Remove a feature result."""
|
||
with self._datalock:
|
||
existed = result_id in self._stores and self._stores[result_id].get('type') == 'feature_result'
|
||
if existed:
|
||
del self._stores[result_id]
|
||
return existed
|
||
|
||
# ── general ─────────────────────────────────────────────────────────
|
||
|
||
def drop_all(self) -> int:
|
||
"""Drop every entry and run gc. Returns count of removed entries."""
|
||
with self._datalock:
|
||
count = len(self._stores)
|
||
self._stores.clear()
|
||
if count:
|
||
gc.collect()
|
||
return count
|
||
|
||
def __repr__(self) -> str:
|
||
with self._datalock:
|
||
types = {}
|
||
for entry in self._stores.values():
|
||
t = entry.get('type', 'unknown')
|
||
types[t] = types.get(t, 0) + 1
|
||
parts = [f"{k}={v}" for k, v in sorted(types.items())]
|
||
return f"<SessionStore {'; '.join(parts)}>"
|