feat: Vector Index & Columnar Engine Optimizations - #621
feat: Vector Index & Columnar Engine Optimizations#621madamsetty-pavan wants to merge 11 commits into
Conversation
product-auto-label
Bot
added
the
api: alloydb
dishaprakash
left a comment
dishaprakash
left a comment
There was a problem hiding this comment.
Can we also add functional tests that use the live database?
We can use fixtures to setup and teardown to bring them to original states.
… separate leaf search modes
|
Done. Added live database functional tests using setup and teardown fixtures:
|
|
/gcbrun |
1 similar comment
|
/gcbrun |
|
Please fix the integration tests |
| params = [] | ||
| if self.pct_leaves_to_search is not None: | ||
| params.append(f"scann.pct_leaves_to_search = {self.pct_leaves_to_search}") | ||
| if self.num_leaves_to_search is not None: |
There was a problem hiding this comment.
ScaNNQueryOptions() with no arguments used to emit scann.num_leaves_to_search = 1. After this change it emits scann.pct_leaves_to_search = 1 instead. T
hose aren't equivalent: one is an absolute leaf count, the other a proportion: so existing users get different recall/latency on upgrade with no error or warning. Could we keep num_leaves_to_search: int = 1 as the default and emit pct_leaves_to_search only when the user sets it? If the default change is intentional, it should be called out as a breaking change in the changelog.
|
|
||
| def index_options(self) -> str: | ||
| """Set index query options for vector store initialization.""" | ||
| if self.mode and self.mode.upper() == "AUTO": |
There was a problem hiding this comment.
nit: Should we add a validation here? If mode is not auto, it should emit a warning
| text(query), | ||
| { | ||
| "model_id": model_id, | ||
| "table_name": self.table_name, |
There was a problem hiding this comment.
Should we be adding a schema_name arg here as well? Vectorstore does support schema_name
| # First we need to get the spec ID for the current table | ||
| specs = await self.adefine_vector_assist_spec() | ||
| if not specs: | ||
| return [] |
There was a problem hiding this comment.
Should we have better error handling and logging here?
|
|
||
| spec_id = specs[0].get("vector_spec_id") | ||
| if not spec_id: | ||
| return [] |
There was a problem hiding this comment.
Should we have better error handling and logging here?
| params = { | ||
| "model_id": model_id, | ||
| "source_table": source_table, | ||
| "source_query": source_query, |
There was a problem hiding this comment.
Do we need a null check for source_query and conf_level?
| { | ||
| "model_id": model_id, | ||
| "table_name": self.table_name, | ||
| "content_column": content_col, |
There was a problem hiding this comment.
Should we add a null check for content_column and embedding_column?
| } | ||
| async with self.engine.connect() as conn: | ||
| result = await conn.execute(text(query), params) | ||
| result_map = result.mappings() |
There was a problem hiding this comment.
nit: We can just use
return [dict(row) for row in result.mappings()] instead of the next three lines.
This can be cleaned up at multiple places in this PR
| if self.pct_leaves_to_search is not None: | ||
| params.append(f"scann.pct_leaves_to_search = {self.pct_leaves_to_search}") | ||
| if self.num_leaves_to_search is not None: | ||
| params.append(f"scann.num_leaves_to_search = {self.num_leaves_to_search}") |
There was a problem hiding this comment.
What happens if pct_leaves_to_search and num_leaves_to_search both are set? Should we document this behaviour or link to somewhere?
| return [] | ||
|
|
||
| spec_id = specs[0].get("vector_spec_id") | ||
| if not spec_id: |
There was a problem hiding this comment.
Should spec_id=0 be considered valid? It would also be true for if not spec_id
| mock_conn = AsyncMock() | ||
| mock_connect.return_value.__aenter__.return_value = mock_conn | ||
| await vs.aenable_columnar_engine(["content"]) | ||
| assert mock_conn.execute.called |
There was a problem hiding this comment.
Ideally we should assert on some data not just that something is executed on the table? Maybe the generated sql and the results. Is there some limitations to doing that here?
Same goes for multiple places at which this assertion is present.
| async def test_live_vector_assist(self, vs): | ||
| """Test vector assist spec definition, application, and recommendations against live AlloyDB instance.""" | ||
| specs = await vs.adefine_vector_assist_spec() | ||
| assert isinstance(specs, list) |
There was a problem hiding this comment.
Should we have stronger assertions based on the content or structure instead? Same goes for other tests in this PR
|
/gcbrun |
3 participants
Part 1 of 3 for AlloyDB AI Features Epic. Contains Vector Optimizations including ScaNN and Columnar Engine.