mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-27 22:21:04 +02:00
feat: batch row imports and split index configuration (#1055)
Adds batching to structured data row imports. Previously each row was sent as a separate WebSocket message, causing one embeddings service call per row. The bulk import clients now accumulate rows into batches (default 40) before sending, and the --batch-size CLI argument is wired through from tg-load-structured-data to the import call. Splits the schema "indexes" field into two separate configurations: - "query-indexes": fields for Cassandra exact-match retrieval - "vector-indexes": fields for vector embedding and semantic search This avoids wasting compute embedding primary key integers that only need exact lookup, while allowing each field to independently opt in to either or both retrieval patterns. Adds row_id (the primary key value) as a clustering column in the Cassandra rows table so multiple rows sharing the same index value no longer silently overwrite each other. Tests updated for new Cassandra schema.
This commit is contained in:
parent
bdfdbb3c1f
commit
1740e315f4
11 changed files with 190 additions and 54 deletions
|
|
@ -117,13 +117,30 @@ class AsyncBulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
async def import_rows(self, flow: str, rows: AsyncIterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
async def import_rows(
|
||||
self, flow: str, rows: AsyncIterator[Dict[str, Any]],
|
||||
batch_size: int = 40,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""Bulk import rows via WebSocket"""
|
||||
ws_url = self._build_ws_url(f"/api/v1/flow/{flow}/import/rows")
|
||||
|
||||
async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
|
||||
batch = []
|
||||
template = None
|
||||
async for row in rows:
|
||||
await websocket.send(json.dumps(row))
|
||||
if template is None:
|
||||
template = row
|
||||
batch.append(row.get("values", row))
|
||||
if len(batch) >= batch_size:
|
||||
message = dict(template)
|
||||
message["values"] = batch
|
||||
await websocket.send(json.dumps(message))
|
||||
batch = []
|
||||
if batch:
|
||||
message = dict(template)
|
||||
message["values"] = batch
|
||||
await websocket.send(json.dumps(message))
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""Close connections"""
|
||||
|
|
|
|||
|
|
@ -538,7 +538,11 @@ class BulkClient:
|
|||
async for raw_message in websocket:
|
||||
yield json.loads(raw_message)
|
||||
|
||||
def import_rows(self, flow: str, rows: Iterator[Dict[str, Any]], **kwargs: Any) -> None:
|
||||
def import_rows(
|
||||
self, flow: str, rows: Iterator[Dict[str, Any]],
|
||||
batch_size: int = 40,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
"""
|
||||
Bulk import structured rows into a flow.
|
||||
|
||||
|
|
@ -548,6 +552,7 @@ class BulkClient:
|
|||
Args:
|
||||
flow: Flow identifier
|
||||
rows: Iterator yielding row dictionaries
|
||||
batch_size: Number of rows per batch (default 40)
|
||||
**kwargs: Additional parameters (reserved for future use)
|
||||
|
||||
Example:
|
||||
|
|
@ -566,15 +571,31 @@ class BulkClient:
|
|||
)
|
||||
```
|
||||
"""
|
||||
self._run_async(self._import_rows_async(flow, rows))
|
||||
self._run_async(self._import_rows_async(flow, rows, batch_size))
|
||||
|
||||
async def _import_rows_async(self, flow: str, rows: Iterator[Dict[str, Any]]) -> None:
|
||||
async def _import_rows_async(
|
||||
self, flow: str, rows: Iterator[Dict[str, Any]],
|
||||
batch_size: int,
|
||||
) -> None:
|
||||
"""Async implementation of rows import"""
|
||||
ws_url = self._build_ws_url(f"/api/v1/flow/{flow}/import/rows")
|
||||
|
||||
async with websockets.connect(ws_url, ping_interval=20, ping_timeout=self.timeout) as websocket:
|
||||
batch = []
|
||||
template = None
|
||||
for row in rows:
|
||||
await websocket.send(json.dumps(row))
|
||||
if template is None:
|
||||
template = row
|
||||
batch.append(row.get("values", row))
|
||||
if len(batch) >= batch_size:
|
||||
message = dict(template)
|
||||
message["values"] = batch
|
||||
await websocket.send(json.dumps(message))
|
||||
batch = []
|
||||
if batch:
|
||||
message = dict(template)
|
||||
message["values"] = batch
|
||||
await websocket.send(json.dumps(message))
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close connections"""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue