feat(query): Add custom SQL query generation for pgvector search

[shamb0]

Hi @timonv,

I’d like to align quickly on this.

This PR addresses our recent discussion, particularly the idea:

"Or we could go lazy (not so lazy, it's not trivial) and add a custom strategy that takes an SQLx query, or a closure that returns an SQLx query with the nodes yielded. That would enable a lot of doors in fully controlling the search process."

As part of this PR, I’ve implemented a Dynamic Search Strategy to enable dynamic SQL query generation for pgvector search.

Key benefits:

• Combines vector similarity with metadata filtering
• Allows users to fully customize query construction
• Provides precise control over filtering and ordering

Looking forward to your feedback on how this can be refined further!

[timonv]

Hey @shamb0, the concern I have with this, is that it only works with Sql and I suspect it's an uphill battle to implement every conceivable query method for every supported provider. I was thinking of a broader solution such that Swiftide makes several big wins:

1. More complex solutions might require custom queries
2. Can't provide defaults for every use case
3. Still allowing for the underlying library to be used is, I believe, a very good practice.

It also removes the need / pressure to implement every search strategy for every storage, as there's always the manual option. Might be a bit harder, but huge relieve when it works.

For instance, say we have a CustomSearchStrategy<Q>, where Q is the internal query type of the storage provider. Then only that has to be implemented, everything is possible, and we can add it to providers on our own pace.

For instance, for what we have now, Q could be:

sqlx -> QueryBuilder https://docs.rs/sqlx/latest/sqlx/struct.QueryBuilder.html
qdrant -> SearchPointsBuilder https://docs.rs/sqlx/latest/sqlx/struct.QueryBuilder.html
lancedb -> Query https://docs.rs/sqlx/latest/sqlx/struct.QueryBuilder.html

In pseudo code, something like this might work:

struct CustomQuery<Q> {
  query: Q
}

impl<Q: Send + Sync + Clone> CustomQuery<Q> {
    pub fn from_query(
        query: impl Fn(&CustomQuery<Q>, &Query<states::Pending>) -> Q + Send + Sync + 'static,
    ) -> Self {
        CustomQuery {
            query: Arc::new(query),
        }
    }

    pub fn build_query(&self, query_node: &Query<states::Pending>) -> Q {
        (*self.query)(self, query_node)
    }
}

  

impl Retrieve<CustomQuery<QueryPoints>> for PgVector {
    async fn retrieve(
        &self,
        search_strategy: &CustomQuery<sqlx::QueryBuilder>,
        query: Query<states::Pending>,
    ) -> Result<Query<states::Retrieved>> {
       ...
       let builder = search_strategy.build_query(&query);

       ... // anything else and execute
    }
}

// A user can then use it as follows:

let strat = CustomQuery::from_query(|strat, query_node| sqls::QueryBuilder::new ... })

This way, there's no uphill battle in trying to implement every conceivable query method, and users can do whatever they want. Other query strategies are there for sane and quick defaults. We could also let the provider (i.e. PGVector) provide a build_custom_query to make it a bit easier. What do you think?

[timonv]

Just to be clear, I'd be happy to merge an implementation that only supports SqlxBuilder at the moment, then I or someone else can expand on it in the future 👍

[shamb0]

Just to be clear, I'd be happy to merge an implementation that only supports SqlxBuilder at the moment, then I or someone else can expand on it in the future 👍

Hi @timonv,

Apologies for the delay :), and thank you for sharing the detailed design guidelines! I'll aim to complete the SqlxBuilder implementation this week.

Next week, I'll start working on the search strategy for the remaining storage. 👍🏾

[shamb0]

Hi @timonv,

I’ve implemented the custom search strategy for sqlx::QueryBuilder based on our recent discussion. The key updates include:

• Dynamic SQL query generation for vector similarity search
• Integration with pgvector
• Comprehensive integration test coverage

Please review and let me know if further adjustments are needed.

Thank you! :)

[timonv]

This looks really nice! Why did you opt for wrapping the CustomQuery with a concrete type? If it's possible to use the CustomQuery instead (let's call it CustomStrategy, to avoid confusion with the Query pipeline), then it removes one layer of abstraction.

[timonv]

You have two files with the same name, one with a typo

[shamb0]

Hi @timonv,

Thank you for your review comments :). I have reworked the code and cleaned it up as suggested.

• Initially, I wrapped CustomQuery<Q> inside a new type PgVecCustomStrategy as I thought it would make the code more concise and Rust-idiomatic. However, as part of this patch, I have removed that additional layer.

  pub struct PgVecCustomStrategy(pub CustomQuery<sqlx::QueryBuilder<'static, sqlx::Postgres>>);  

• I believe the PR is now ready for intake. Please let me know if any further enhancements are needed. :)

[timonv]

Excellent, looks really good. Two more small comments, otherwise this looks very useful and looks ready to merge. Nice work (again) 🚀

[timonv]

The correct name here is retriever instead of provider

[timonv]

Looks like top_k and vector_field aren't used in the retrieval implementation. Since it builds from a closure anyway, and any data can be moved there, how about simplifying by removing the fields, having only query and _marker here?

[timonv]

Two tests here are ignored

[shamb0]

Thank you @timonv for your thorough and constructive feedback :). I appreciate your patience in guiding me through these improvements.

I have addressed both key points from your review:

1. Updated the terminology from "provider" to "retriever" throughout the documentation to better reflect the architectural intent
2. Simplified the CustomStrategy implementation by moving configuration into the query generator closure, reducing structural complexity

The documentation has been enhanced with:

• Clearer examples demonstrating the new configuration pattern
• Detailed explanation of the query generator's responsibilities
• Updated implementation notes for users of the custom search strategy API

I believe these changes have made the code more maintainable and the API more intuitive to use. Please let me know if you'd like to see any further improvements or have additional suggestions.

[timonv]

Looks good! I want to go over the docs one more time before merging, I'll do the changes. I will merge this asap.