-
Notifications
You must be signed in to change notification settings - Fork 332
pgml Python SDK with vector search support #636
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
run_create_or_insert_statement(conn, create_schema_statement) | ||
create_table_statement = ( | ||
"CREATE TABLE IF NOT EXISTS %s (\ | ||
id serial8 PRIMARY KEY,\ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
id serial8 PRIMARY KEY,\ | |
id bigserial PRIMARY KEY,\ |
More idiomatic, but doesn't matter.
document uuid NOT NULL,\ | ||
metadata jsonb NOT NULL DEFAULT '{}',\ | ||
text text NOT NULL,\ | ||
UNIQUE (document)\ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Duplicate primary key technically, since this is unique. Curious why we can't use the id
field as the document identifier?
run_create_or_insert_statement(conn, create_index_statement, autocommit=True) | ||
|
||
create_index_statement = ( | ||
"CREATE INDEX CONCURRENTLY IF NOT EXISTS \ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Don't need to create an index here on document
, UNIQUE
does it already automatically. So I think you end up with two indexes on the same field.
) | ||
run_create_or_insert_statement(conn, create_statement) | ||
|
||
index_statement = ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
UNIQUE (task, splitter, model)
in the table definition creates a compound index on those three columns. Having 3 additional individual indexes on the same columns may not be necessary.
created_at timestamptz NOT NULL DEFAULT now(), \ | ||
task text NOT NULL, \ | ||
splitter int8 NOT NULL REFERENCES %s\ | ||
ON DELETE CASCADE\ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Are you sure you want to CASCADE
? This has the effect of automatically deleting rows from this table when the row they are referencing in another table is deleted. This can delete a lot of data accidentally. The preferred way for me generally is to ON DELETE RESTRICT
which is the default. That way, you'll get an error when attempting to delete a splitter that's referenced by this table. This is default behavior also, so you can just remove ON DELETE CASCADE
.
self.transforms_table = self.name + ".transforms" | ||
create_statement = ( | ||
"CREATE TABLE IF NOT EXISTS %s (\ | ||
oid regclass PRIMARY KEY,\ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Odd way of referencing another table, I've never seen that done before. Is this compatible with logical replication? I.e. if we want to move this data to another database, the oids probably won't match anymore, since they are specific to a Postgres installation.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This should probably be called table_name
. An oid
is something very Postgres-specific and does not mean a table reference at all. In fact, it can reference a row in a TOAST table or a type in pg_type
.
"CREATE TABLE IF NOT EXISTS %s ( \ | ||
id serial8 PRIMARY KEY,\ | ||
created_at timestamptz NOT NULL DEFAULT now(),\ | ||
chunk int8 NOT NULL REFERENCES %s\ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The convention is sometimes to use field_id
to explain to the reader that this refers to the primary key (usually id
) column in another table. The reviewer would look in this case if the column being referenced has an index on it, which is often required for foreign keys: a foreign key validation needs to be an index scan, the quickest way to validate that this value exists in another table.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
_id
is not just a railsism, it's a popular and significant readability improvement to easily distinguish local fields from foreign keys, which isn't really Hungarian notation. It's part of what-this-field-represents, rather than what-data-type-is-this-field.
It's more similar to the convention many programming languages use by adding *
or &
to distinguish references and pointers from local objects on the stack.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A foreign key can be polymorphic, in which case, the naming convention suddenly becomes important. Sure, one can't really enforce naming conventions programmatically, but there is no reason to make something harder to understand.
For example, created_at
can be renamed to random_date
and only by looking at the definition of the column one would realize that this is actually a creation timestamp, but we don't do that, we call it created_at
because we want to be kind to our future selves and to reviewers and other users of our code. An argument can be made for calling variables abcd
and only by reading the code fully, one will truly understand what they do, but we don't do that either.
Rails can be many things that we don't agree with, but the way ActiveRecord handles data is fantastic, and we would be wise to learn from their experience and do something most engineers are familiar and happy to work with.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
And you will commonly see code that is forced to name two variables in a context things like document_ptr
and a document
, to disambiguate the one that has been dereferenced safely. No one would assume that the compiler checks _ptr
correctness, but if they ever saw a variable with that extension, next to a variable without the extension, and the truth of the references disagreed with the conventions we've established match expectations, people would say the variable was misnamed and likely to cause bugs or confusion.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We write code for others to read. Code should consistently follow the conventions established in the project. The convention has already been established, and there is no need to revisit that convention, when the suggestion obviously confused multiple team members in multiple places due to not just inconsistency with the greater project scope, but inconsistency within the necessary layers to achieve a solution, and also inconsistency within a single table.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Natural joins with USING
is generally useless, because it uses all columns in common, not just the key. In this case created_at
would be shared but would not be the same, so natural joins won't work.
Note
USING is reasonably safe from column changes in the joined relations since only the listed columns are combined. NATURAL is considerably more risky since any schema changes to either relation that cause a new matching column name to be present will cause the join to combine that new column as well.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If you're still not convinced, there is the user_uuid
in your example, which is yet another instance.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
But similarly, if you really want to use USING
, and not natural joins, it works just as well to call them both document_id
, and you can keep that feature of SQL.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Calling them both document_id
would also be a break with Rails conventions, though...
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You seem confused about the importance of Rails conventions, and to be missing the point that this schema is not an island unto itself, nor is this SDK. The important aspect is that a database and schema in this context are useless without an application layer of logic on top. That application layer will have concerns with ORM that will be simpler to maintain if the database follows a convention of using _id
prefixes. The reason you see the pattern replicated so widely, is because it's so frequently eases the mental burden of application developers. DBA's or people working only inside the database may not understand these concepts or concerns, but that doesn't mean they don't exist.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The tables feel a bit overindexed. In Postgres, each index carries a write penalty: for each row that's inserted into the table, each index needs to be updated accordingly. When all (or most) columns are updated, an update to any column requires an update to all indexes (i.e. Postgres can't do HOT, heap-only, updates).
I think I understand the schema overall, although it would be helpful to define it in a separate .sql
file which can then be executed when the SDK is used for the first time. Although, we do generate a lot of tables in the fly, so understandably that's not possible for all use cases.
Overall, I think this is great to start with, and might require some optimizations as it's deployed at scale.
Experience. Removing an index is dangerous, adding an index is safe. |
log.info("id key is not present.. hashing") | ||
document_id = hashlib.md5(text.encode("utf-8")).hexdigest() | ||
metadata = document | ||
delete_statement = "DELETE FROM %s WHERE document = %s" % ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This would be more clear to me if the field name was document_id
to match the variable name. I was half expecting document
to be a text value.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Right, so using one convention for names at one layer, and a different convention for names in the next layer down seems confusing, there is good reason to add _id
in both layers.
chunks = text_splitter.create_documents([text]) | ||
for chunk_id, chunk in enumerate(chunks): | ||
insert_statement = ( | ||
"INSERT INTO %s (document,splitter,chunk_id, chunk) VALUES (%s, %s, %s, %s);" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Missing _id
suffixes make this hard for me to read.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, chunk_id
is the chunk index for the given document text. For example: if the text "hello world" is split into two chunks then chunk_id 0 will map to "hello" and chunk_id 1 will map to "world". id
is the global id of the chunk across all documents and splitters.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yes, this is a perfect example of needing both a local and reference to the same "object" in a single context, and being forced to disambiguate one of them with an _id
suffix. The confusing ones in this line are document
, and splitter
. chunk
and chunk_id
are clear and easily distinguished, even though chunk_id
does not in fact refer to a foreign key in a chunks
table, it's clearly marked as a reference.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
chunk_id
is the ordering of the chunk within the document. chunk_index
sounds good.
"CREATE TABLE IF NOT EXISTS %s ( \ | ||
id serial8 PRIMARY KEY,\ | ||
created_at timestamptz NOT NULL DEFAULT now(),\ | ||
chunk int8 NOT NULL REFERENCES %s\ |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
_id
is not just a railsism, it's a popular and significant readability improvement to easily distinguish local fields from foreign keys, which isn't really Hungarian notation. It's part of what-this-field-represents, rather than what-data-type-is-this-field.
It's more similar to the convention many programming languages use by adding *
or &
to distinguish references and pointers from local objects on the stack.
model_params = results[0]["parameters"] | ||
|
||
# get all chunks that don't have embeddings | ||
embeddings_statement = ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This needs to be refactored into the insert statement to avoid round tripping the vector.
embeddings_table = self._create_or_get_embeddings_table( | ||
conn, model_id=model_id, splitter_id=splitter_id | ||
) | ||
select_statement = "SELECT name, parameters FROM %s WHERE id = %d;" % ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I would include this with a CTE in the embeddings statement as well to avoid the round trip.
results = run_select_statement(conn, select_statement) | ||
|
||
model = results[0]["name"] | ||
query_embeddings = self._get_embeddings( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A query builder can be used for the CTE here.
query_embeddings = self._get_embeddings( | ||
conn, query, model_name=model, parameters=query_parameters | ||
) | ||
embeddings_table = self._create_or_get_embeddings_table( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This needs to be cached some how to avoid multiple round trips to the db for this function.
) | ||
|
||
select_statement = ( | ||
"SELECT chunk, 1 - (%s.embedding <=> %s::float8[]::vector) AS score FROM %s ORDER BY score DESC LIMIT %d;" |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This needs to be embedded with all the results queries to avoid n+1 queries
for result in results: | ||
_out = {} | ||
_out["score"] = result["score"] | ||
select_statement = "SELECT chunk, document FROM %s WHERE id = %d" % ( |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Prefer joins in SQL to round trips from the app.
self.pool.putconn(conn) | ||
return Collection(self.pool, name) | ||
|
||
def delete_collection(self, name: str) -> None: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hard delete is a bit of a footgun.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, we could have some archiving mechanism, which simply alters the schema to name
+ '_archive', and then enable delete or restore on archives.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Adding more storage space is cheap. Recovering lost data is expensive and scary.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Also, we should support "undo".
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
that's what I meant by "restore"
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That does not guarantee people aren't still using the table, if they are ignoring the flag by using poorly written homegrown queries, which is a potentially desirable outcome for projects that outgrow simple SDK designs, and what to leverage the full expressive power of SQL.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🥳
Objective of this SDK is to provide an easy interface for PostgresML generative AI capabilities. This version supports vector search using multiple models and text splitters.
Quick start instructions are here