feat: Index markdown in pgvector

[shamb0]

Related to: #158

PR Overview (Part-01)

This PR introduces the first part of the indexing pipeline, which handles the process of loading, chunking, enhancing, embedding, and storing markdown content into the pgvector store.

How to Run the Example

The provided example, swiftide/examples/index_md_into_pgvector.rs, uses the Ollama LLM provider with the llama3.2:latest model for generating QA metadata.

Steps to Run:

1. Start the PostgreSQL services:

   docker compose -f ./scripts/docker/docker-compose-db-pg.yml up

2. Run the example with the default configuration:

   RUST_LOG=debug cargo run --example index-md-pgvector

[timonv]

Nice job, this looks already very solid! I have a couple of comments and change requests. Notably the N+1 is important.

Additionally, I'm also missing tests. Could you add tests for Persist, batch and non batch? For lance and qdrant I used testcontainers, see theirs for an example, it works pretty well. They're in their retrieve implementation iirc.

For the example test dataset, could you remove that and use the root readme instead? For the docker compose yml, I'm not sure, either we should have one at the root, leave it out, or use testcontainers in examples as well.

Really nice job on this, it's looking good!

[timonv]

Instead of having this here, you could create a default_connection_pool method instead, and have that as the default build_fn for the builder (in the attr macro) instead. Then this method could build and connect. Added benefit is that PgVector would then does not need an Option anymore, just a PgPool.

[shamb0]

Thanks for the suggestion! I’ve refactored the code accordingly, but please let me know if I’ve misunderstood anything.

I’ve wrapped PgPool inside the PgDBConnectionPool tuple struct and made connection_pool non-optional. The default is now handled via the builder macro with PgDBConnectionPool::default(), and the connection to the Postgres server is established through the try_connect_to_pool() method.

#[builder(default = "PgDBConnectionPool::default()")]
connection_pool: PgDBConnectionPool,

Let me know if you think further improvements are needed.

[timonv]

Yes exactly, almost! I've left some comments

[shamb0]

@timonv, Thanks a lot! I really appreciate your detailed feedback. I’ve addressed most of the points, but I’d like some clarification on your comment about the PgVector::connection_pool field. Could you confirm if I’ve understood it correctly, or if further changes are needed?

I’ve also added unit tests, following the primary use cases in LanceDB, and included tests for retrieve. Everything seems to be working fine so far (#406). The PR #406 is still in draft, and I plan to rebase it after the current PR intake.

[timonv]

Hey @shamb0,

Getting there! I've provided some more feedback. I still need to run the example to see how it performs, will do that this week.

[timonv]

You should use pgvector/pgvector, ankane's hasn't been updated for a year and is unofficial.

[shamb0]

Sounds good—I’ve migrated to the latest image, pgvector/pgvector:pg17:

let postgres = testcontainers::GenericImage::new("pgvector/pgvector", "pg17")

Let me know if there’s anything else you’d suggest here!

[timonv]

You might not need the bindmount, or do you want to persist data across tests? We're currently having issues with github runner limitations, disk space is one, so we need to be a bit careful

[shamb0]

Given the resource constraints in the CI/CD build context, Mount::tmpfs_mount() seems like a promising choice since the volume resource is released once the container is stopped.

let postgres = testcontainers::GenericImage::new("pgvector/pgvector", "pg17")
    .with_wait_for(WaitFor::message_on_stdout(
        "database system is ready to accept connections",
    ))
    .with_mapped_port(host_port, 5432.tcp())
    .with_env_var("POSTGRES_USER", "myuser")
    .with_env_var("POSTGRES_PASSWORD", "mypassword")
    .with_env_var("POSTGRES_DB", "mydatabase")
    .with_mount(Mount::tmpfs_mount("/var/lib/postgresql/data"))
    .start()
    .await
    .expect("Failed to start Postgres container");

Does this approach fit the context, or would you suggest any adjustments?

[timonv]

Almost exactly right! I prefer it if builders do not do IO if they can avoid it, for multiple reasons. In this case, that also has the benefit of being able to connect lazilly and hiding the details of the connection pool.

i.e. the builder api like:

 let pgv_storage = PgVector::builder()
        .database_url(pgv_db_url)
        .pool_size(10) // With a sane default if ommitted 
        .vector_size(384)
        .with_vector(EmbeddedField::Combined)
        .with_metadata(METADATA_QA_TEXT_NAME)
        .table_name("swiftide_pgvector_test".to_string())
        .build()
        .unwrap();

And then in PgVector::setup (which is only called once):

async fn setup(&self) -> Result<()> {
  self.try_connect_to_pool(self.database_url, self.pool_size).await?;
  ...

[shamb0]

@timonv, I'm looking for your input on a design choice here.

If we decide to handle the database connection pool setup within fn setup(&self) instead of PgVectorBuilder, we'll need to mutate PgVector within fn setup(). This change would mean updating the function signature in trait Persist to:

async fn setup(&mut self) -> Result<()>

For example:

async fn setup(&mut self) -> Result<()> {
  self.connection_pool = self.try_connect_to_pool(self.database_url, self.pool_size).await?;
  ...
}

This adjustment would introduce breaking changes across the stack, particularly impacting:

• swiftide-indexing/src/persist/memory_storage.rs
• swiftide-integrations/src/lancedb/persist.rs
• swiftide-integrations/src/qdrant/persist.rs
• swiftide-integrations/src/redis/persist.rs

Would you prefer moving the IO operations into Persist::setup() for these components? If so, we could handle this as a separate PR to streamline the updates.

Looking forward to your thoughts!

[timonv]

This value isn't optional, is it? What happens if it is None?

[shamb0]

Just to clarify: this parameter can’t be None because if it is, the user won’t be able to build the PgVector::fields parameter, and PgVectorBuilder::with_vector() would fail. We’re ensuring that users configure this parameter correctly before launching the indexing pipeline.

I’d love to hear your thoughts and any suggestions for enhancing this approach.

[timonv]

Is it intentional that this value is optional and can be None?

[shamb0]

For the batch_size configuration, I decided to keep it optional, drawing inspiration from the reference Qdrant implementation. By default, the batch size for PGVector is set to 50—though we still plan to fine-tune this for optimal performance.

To answer your question, "Is it intentional that this value is optional and can be None?"—yes, it’s designed that way. In the indexing pipeline, Pipeline::then_store_with() functions as an adapter, routing Node Streams to backend storage depending on the batch setting.

Here’s the approach:

• When Persist::batch_size() returns None, each node is processed individually, with Persist::store() sending each chunk to the backend.
• If Persist::batch_size() returns Some(), batch processing is enabled. The stream of nodes is grouped into chunks based on the batch size, and Persist::batch_store() sends these batches to storage.

Does this implementation align with requirements? Let me know if there’s anything specific you’d like adjusted.

[timonv]

Does this do a query? Debug is called extensively with tracing. Which would mean (I think?) a query on every log / trace statement.

[timonv]

See other comment. Changing that and moving it to the main struct will also remove the need to clone, unwrap and option.

[timonv]

Could you add tests for storing nodes and verifying with a query that it got persisted as expected? For all the different embedding modes and with/without metadata.

[timonv]

It's good practice not to log errors in libraries if the error is bubbled up, as the caller is expected to handle it. Could you remove it?

[shamb0]

Sounds good! I've cleaned up the extra logs and traces.

[timonv]

I have a hunch this query can be made simpler, reducing the code and at the same time speeding things up. There's a lot of loops / allocations in the sql generation that might be avoided. I'm not not a 100% sure, haven't fully dived into it. Additionally, what happens if the node is already present? In swiftide we upsert on the id of the node.

Postgres supports bulk 'dynamic' bulk insertion with UNNEST, see https://github.com/launchbadge/sqlx/blob/main/FAQ.md#how-can-i-bind-an-array-to-a-values-clause-how-can-i-do-bulk-inserts.

[shamb0]

• Thanks @timnov for suggesting UNNEST; that sounds promising! I'll give it a try and keep you posted.

• Regarding the question, “What happens if the node is already present?”

  • In the current setup, if a node already exists, a duplicate entry is created. I considered using a caching layer, like Redis, to manage duplicate source files after the FileLoader stage, rather than handling it at the storage level.

• On the suggestion to “upsert on the ID of the node”

  • Upserting on the node’s ID could lead to data loss. Here’s why:

    Let’s take a look at the indexing pipeline structure:

    indexing::Pipeline::from_loader(FileLoader::new(test_dataset_path).with_extensions(&["md"]))
        .then_chunk(ChunkMarkdown::from_chunk_range(10..2048))
        .then(MetadataQAText::new(llm_client.clone()))
        .then_in_batch(Embed::new(fastembed.clone()).with_batch_size(100))
        .then_store_with(pgv_storage.clone())
        .run()
        .await?;

    • Explanation: In this example, the pipeline starts with FileLoader, which processes a single file and creates a node, let’s call it N1.
    • In the next stage, ChunkMarkdown, N1 is split into multiple nodes (one-to-X relationship), each inheriting the same ID as N1.
    • These nodes, sharing the same ID, are then passed to the storage stage (pgv_storage).

  If we use upsert based on the ID, data from previously processed nodes may be overwritten, leading to possible data loss. I’d be glad to hear your thoughts on this or correct any misunderstandings if there’s a different intended approach.

[timonv]

It shouldn't be an issue, the id is lazilly created based on a hash of the chunk (usually at index time). Do you think there is a bug here? You might be on to something, I recently redid id generation such that it's lazy, but it indeed seems that it copies over the id when chunking, while it shouldn't. Perhaps the id shouldn't be lazy at all.

[timonv]

@shamb0 I think you identified a pretty hefty potential bug. Nodes should never share the same ID unless their chunk and path is exactly the same. As long as that holds true it should be safe to upsert. What could happen (maybe with a node cache?), that the id gets set earlier and gets copied over during chunking, never updating it again. That's not intentional and is a bug. I've fixed it in #414, thanks!

[shamb0]

Thanks, @timnov. I’m glad our discussion led to identifying this bug. Apologies for the delay—I’ve been managing some personal matters but am back on this task now. I’ll start working with patch #414 today and keep you updated on any progress.

[timonv]

No worries, no scrum gods to please in open source. Thank you for identifying the bug.

[shamb0]

HaHa, Thanks! :-) Happy to help improve the project. Bugs are part of our journey—glad we caught it early!

[shamb0]

• Thanks, @timnov, for the second round of review comments. I believe around 90% of the points have been addressed.
• The Retrieve functionality integration is complete.
• For unit test coverage, I've included critical scenarios and kept the scope focused to facilitate the review process. Let me know if you'd like to see any additional use cases covered—I’d appreciate your feedback.
• I’m currently troubleshooting a failure with EmbedMode::PerField. I’ll keep you posted on any progress there.

[timonv]

• Thanks, @timnov, for the second round of review comments. I believe around 90% of the points have been addressed.

  • The Retrieve functionality integration is complete.

  • For unit test coverage, I've included critical scenarios and kept the scope focused to facilitate the review process. Let me know if you'd like to see any additional use cases covered—I’d appreciate your feedback.

  • I’m currently troubleshooting a failure with EmbedMode::PerField. I’ll keep you posted on any progress there.

Perfect, let me know if I can help / continue the review

[shamb0]

• Thanks, @timnov, for the second round of review comments. I believe around 90% of the points have been addressed.

  • The Retrieve functionality integration is complete.
  • For unit test coverage, I've included critical scenarios and kept the scope focused to facilitate the review process. Let me know if you'd like to see any additional use cases covered—I’d appreciate your feedback.
  • I’m currently troubleshooting a failure with EmbedMode::PerField. I’ll keep you posted on any progress there.

Perfect, let me know if I can help / continue the review

Yes @timonv, I need your help regarding the discussion thread on (#392 (comment)).

Could you review my recent comments and share any updated thoughts when convenient? Once I have your input, I can proceed with the implementation, and then we can move forward to the intake review process.