apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-27 22:21:04 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

feat: batch row imports and split index configuration (#1055)

Browse source 
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:
cybermaggedon 2026-07-22 10:20:06 +01:00 • committed by GitHub
parent bdfdbb3c1f
commit 1740e315f4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 190 additions and 54 deletions
Download patch file Download diff file Expand all files Collapse all files

10
trustgraph-flow/trustgraph/embeddings/row_embeddings/embeddings.py
View file

@ -122,10 +122,16 @@ class Processor(CollectionConfigHandler, FlowProcessor):
fields=fields
)
vector_indexes = set(schema_def.get("vector-indexes", []))
for field in fields:
if field.name in vector_indexes:
field.indexed = True
ws_schemas[schema_name] = row_schema
logger.info(
f"Loaded schema: {schema_name} with "
f"{len(fields)} fields for {workspace}"
f"{len(fields)} fields, {len(vector_indexes)} vector-indexed "
f"for {workspace}"
)
except Exception as e:
@ -140,7 +146,7 @@ class Processor(CollectionConfigHandler, FlowProcessor):
"""Get all index names for a schema."""
index_names = []
for field in schema.fields:
if field.primary or field.indexed:
if field.indexed:
index_names.append(field.name)
return index_names

8
trustgraph-flow/trustgraph/query/rows/cassandra/service.py
View file

@ -190,11 +190,17 @@ class Processor(FlowProcessor):
fields=fields
)
query_indexes = set(schema_def.get("query-indexes", []))
for field in fields:
if field.name in query_indexes:
field.indexed = True
ws_schemas[schema_name] = row_schema
builder.add_schema(schema_name, row_schema)
logger.info(
f"Loaded schema: {schema_name} with "
f"{len(fields)} fields for {workspace}"
f"{len(fields)} fields, {len(query_indexes)} query-indexed "
f"for {workspace}"
)
except Exception as e:

32
trustgraph-flow/trustgraph/storage/rows/cassandra/write.py
View file

@ -7,10 +7,13 @@ Uses a single 'rows' table with the schema:
- schema_name: text
- index_name: text
- index_value: frozen<list<text>>
- row_id: text
- data: map<text, text>
- source: text
Each row is written multiple times - once per indexed field defined in the schema.
Each row is written multiple times - once per query-indexed field defined
in the schema. The row_id (primary key value) is included as a clustering
column to allow multiple rows with the same index value.
"""
import asyncio
@ -183,10 +186,16 @@ class Processor(CollectionConfigHandler, FlowProcessor):
fields=fields
)
query_indexes = set(schema_def.get("query-indexes", []))
for field in fields:
if field.name in query_indexes:
field.indexed = True
ws_schemas[schema_name] = row_schema
logger.info(
f"Loaded schema: {schema_name} with "
f"{len(fields)} fields for {workspace}"
f"{len(fields)} fields, {len(query_indexes)} query-indexed "
f"for {workspace}"
)
except Exception as e:
@ -263,9 +272,10 @@ class Processor(CollectionConfigHandler, FlowProcessor):
schema_name text,
index_name text,
index_value frozen<list<text>>,
row_id text,
data map<text, text>,
source text,
PRIMARY KEY ((collection, schema_name, index_name), index_value)
PRIMARY KEY ((collection, schema_name, index_name), index_value, row_id)
)
"""
@ -351,6 +361,14 @@ class Processor(CollectionConfigHandler, FlowProcessor):
self.registered_partitions.add(cache_key)
logger.info(f"Registered partitions for {collection}/{schema_name}: {index_names}")
def get_row_id(self, schema: RowSchema, value_map: Dict[str, str]) -> str:
"""Get the primary key value for a row."""
for field in schema.fields:
if field.primary:
value = value_map.get(field.name)
return str(value) if value is not None else ""
return ""
def build_index_value(self, value_map: Dict[str, str], index_name: str) -> List[str]:
"""
Build the index_value list for a given index.
@ -420,8 +438,8 @@ class Processor(CollectionConfigHandler, FlowProcessor):
# Prepare insert statement
insert_cql = f"""
INSERT INTO {safe_keyspace}.rows
(collection, schema_name, index_name, index_value, data, source)
VALUES (%s, %s, %s, %s, %s, %s)
(collection, schema_name, index_name, index_value, row_id, data, source)
VALUES (%s, %s, %s, %s, %s, %s, %s)
"""
# Process each row in the batch
@ -434,6 +452,8 @@ class Processor(CollectionConfigHandler, FlowProcessor):
if raw_value is not None:
data_map[field.name] = str(raw_value)
row_id = self.get_row_id(schema, value_map)
# Write one copy per index
for index_name in index_names:
index_value = self.build_index_value(value_map, index_name)
@ -450,7 +470,7 @@ class Processor(CollectionConfigHandler, FlowProcessor):
await async_execute(
self.session,
insert_cql,
(collection, schema_name, index_name, index_value, data_map, source),
(collection, schema_name, index_name, index_value, row_id, data_map, source),
)
rows_written += 1
except Exception as e: