apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-28 06:31:03 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

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

7
tests/integration/test_rows_cassandra_integration.py
View file

@ -108,6 +108,7 @@ class TestRowsCassandraIntegration:
processor.ensure_tables = Processor.ensure_tables.__get__(processor, Processor)
processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor)
processor.get_index_names = Processor.get_index_names.__get__(processor, Processor)
processor.get_row_id = Processor.get_row_id.__get__(processor, Processor)
processor.build_index_value = Processor.build_index_value.__get__(processor, Processor)
processor.register_partitions = Processor.register_partitions.__get__(processor, Processor)
processor._apply_schema_config = Processor._apply_schema_config.__get__(processor, Processor)
@ -477,9 +478,9 @@ class TestRowsCassandraIntegration:
# Check that data is passed as a dict (will be map in Cassandra)
insert_call = rows_insert_calls[0]
values = insert_call[0][1]
# Values are: (collection, schema_name, index_name, index_value, data, source)
# values[4] should be the data map
data_map = values[4]
# Values are: (collection, schema_name, index_name, index_value, row_id, data, source)
# values[5] should be the data map
data_map = values[5]
assert isinstance(data_map, dict)
assert data_map["id"] == "123"
assert data_map["name"] == "Test Item"

17
tests/unit/test_cli/test_error_handling_edge_cases.py
View file

@ -310,22 +310,7 @@ Bob,bob@email.com,"age with quote,42'''
def test_websocket_connection_failure(self):
"""Test WebSocket connection failure handling"""
input_file = self.create_temp_file("name,email\nJohn,john@email.com", '.csv')
descriptor_file = self.create_temp_file(json.dumps(self.valid_descriptor), '.json')
try:
# Test with invalid URL
with pytest.raises(Exception):
load_structured_data(
api_url="http://invalid-host:9999",
input_file=input_file,
descriptor_file=descriptor_file,
batch_size=1,
flow='obj-ex'
)
finally:
self.cleanup_temp_file(input_file)
self.cleanup_temp_file(descriptor_file)
skip_internal_tests()
# Edge Case Data Tests
def test_extremely_long_lines(self):

10
tests/unit/test_embeddings/test_row_embeddings_processor.py
View file

@ -64,8 +64,8 @@ class TestRowEmbeddingsProcessor(IsolatedAsyncioTestCase):
index_names = processor.get_index_names(schema)
# Should include primary key and indexed field
assert 'id' in index_names
# Should include only indexed fields, not primary key
assert 'id' not in index_names
assert 'name' in index_names
assert 'email' not in index_names
@ -356,12 +356,10 @@ class TestRowEmbeddingsProcessor(IsolatedAsyncioTestCase):
# Mock the flow
mock_embeddings_request = AsyncMock()
# Return batch of vector sets (one per text)
# 4 unique texts: CUST001, John Doe, CUST002, Jane Smith
# 2 unique texts: John Doe, Jane Smith (only 'name' is indexed)
mock_embeddings_request.embed.return_value = [
[[0.1, 0.2, 0.3]], # vectors for text 1
[[0.2, 0.3, 0.4]], # vectors for text 2
[[0.3, 0.4, 0.5]], # vectors for text 3
[[0.4, 0.5, 0.6]], # vectors for text 4
]
mock_output = AsyncMock()
@ -383,7 +381,7 @@ class TestRowEmbeddingsProcessor(IsolatedAsyncioTestCase):
# Verify it was called with a list of texts
call_args = mock_embeddings_request.embed.call_args
assert 'texts' in call_args.kwargs
assert len(call_args.kwargs['texts']) == 4
assert len(call_args.kwargs['texts']) == 2
# Should have sent output
mock_output.send.assert_called()

13
tests/unit/test_storage/test_rows_cassandra_storage.py
View file

@ -195,6 +195,7 @@ class TestRowsCassandraStorageLogic:
processor.session = MagicMock()
processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor)
processor.get_index_names = Processor.get_index_names.__get__(processor, Processor)
processor.get_row_id = Processor.get_row_id.__get__(processor, Processor)
processor.build_index_value = Processor.build_index_value.__get__(processor, Processor)
processor.ensure_tables = MagicMock()
processor.register_partitions = MagicMock()
@ -231,13 +232,14 @@ class TestRowsCassandraStorageLogic:
# Verify using unified rows table
assert "INSERT INTO default.rows" in insert_cql
# Values should be: (collection, schema_name, index_name, index_value, data, source)
# Values should be: (collection, schema_name, index_name, index_value, row_id, data, source)
assert values[0] == "test_collection" # collection
assert values[1] == "test_schema" # schema_name
assert values[2] == "id" # index_name (primary key field)
assert values[3] == ["123"] # index_value as list
assert values[4] == {"id": "123", "value": "test_data"} # data map
assert values[5] == "" # source
assert values[4] == "123" # row_id (primary key value)
assert values[5] == {"id": "123", "value": "test_data"} # data map
assert values[6] == "" # source
@pytest.mark.asyncio
@patch('trustgraph.storage.rows.cassandra.write.async_execute', new_callable=AsyncMock)
@ -261,6 +263,7 @@ class TestRowsCassandraStorageLogic:
processor.session = MagicMock()
processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor)
processor.get_index_names = Processor.get_index_names.__get__(processor, Processor)
processor.get_row_id = Processor.get_row_id.__get__(processor, Processor)
processor.build_index_value = Processor.build_index_value.__get__(processor, Processor)
processor.ensure_tables = MagicMock()
processor.register_partitions = MagicMock()
@ -321,6 +324,7 @@ class TestRowsCassandraStorageBatchLogic:
processor.session = MagicMock()
processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor)
processor.get_index_names = Processor.get_index_names.__get__(processor, Processor)
processor.get_row_id = Processor.get_row_id.__get__(processor, Processor)
processor.build_index_value = Processor.build_index_value.__get__(processor, Processor)
processor.ensure_tables = MagicMock()
processor.register_partitions = MagicMock()
@ -378,6 +382,7 @@ class TestRowsCassandraStorageBatchLogic:
processor.session = MagicMock()
processor.sanitize_name = Processor.sanitize_name.__get__(processor, Processor)
processor.get_index_names = Processor.get_index_names.__get__(processor, Processor)
processor.get_row_id = Processor.get_row_id.__get__(processor, Processor)
processor.build_index_value = Processor.build_index_value.__get__(processor, Processor)
processor.ensure_tables = MagicMock()
processor.register_partitions = MagicMock()
@ -432,7 +437,7 @@ class TestUnifiedTableStructure:
assert "index_value frozen<list<text>>" in rows_cql
assert "data map<text, text>" in rows_cql
assert "source text" in rows_cql
assert "PRIMARY KEY ((collection, schema_name, index_name), index_value)" in rows_cql
assert "PRIMARY KEY ((collection, schema_name, index_name), index_value, row_id)" in rows_cql
# Check row_partitions table creation
partitions_cql = processor.session.execute.call_args_list[1][0][0]