use serde_json::{json, Value};

use sqlx::{Executor, Pool, Postgres};

use tokio::task::JoinSet;

use crate::pipeline::FieldAction;

let mut args = args.unwrap_or_default();

let args = args.as_object_mut().context("args must be a JSON object")?;

let mut parsed_schemas = vec![];

let pipelines: Vec<(Pipeline, HashMap<String, FieldAction>)> =

pipelines.into_iter().zip(parsed_schemas).collect();

let batch_size = args

.remove("batch_size")

.map(|x| x.try_to_u64())

.unwrap_or(Ok(100))?;

let parallel_batches = args

.get("parallel_batches")

.map(|x| x.try_to_u64())

.unwrap_or(Ok(1))? as usize;

let mut set = JoinSet::new();

for batch in documents.chunks(batch_size as usize) {

if set.len() < parallel_batches {

let local_self = self.clone();

let local_batch = batch.to_owned();

let local_args = args.clone();

let local_pipelines = pipelines.clone();

let local_pool = pool.clone();

set.spawn(async move {

local_self

._upsert_documents(local_batch, local_args, local_pipelines, local_pool)

.await

});

} else {

if let Some(res) = set.join_next().await {

res??;

progress_bar.inc(batch_size);

}

}

}

while let Some(res) = set.join_next().await {

res??;

progress_bar.inc(batch_size);

}

progress_bar.println("Done Upserting Documents\n");

progress_bar.finish();

Ok(())

}

async fn _upsert_documents(

self,

batch: Vec<Json>,

args: serde_json::Map<String, Value>,

mut pipelines: Vec<(Pipeline, HashMap<String, FieldAction>)>,

pool: Pool<Postgres>,

) -> anyhow::Result<()> {

let project_info = &self.database_data.as_ref().unwrap().project_info;

let mut transaction = pool.begin().await?;

let mut query_values = String::new();

let mut binding_parameter_counter = 1;

for _ in 0..batch.len() {

query_values = format!(

"{query_values}, (${}, ${}, ${})",

binding_parameter_counter,

binding_parameter_counter + 1,

binding_parameter_counter + 2

binding_parameter_counter += 3;

}

let query = query.replace(

"{values_parameters}",

&query_values.chars().skip(1).collect::<String>(),

);

let query = query.replace(

"{binding_parameter}",

&format!("${binding_parameter_counter}"),

);

let mut query = sqlx::query_as(&query);

let mut source_uuids = vec![];

for document in &batch {

let id = document

.get("id")

.context("`id` must be a key in document")?

.to_string();

let md5_digest = md5::compute(id.as_bytes());

let source_uuid = uuid::Uuid::from_slice(&md5_digest.0)?;

source_uuids.push(source_uuid);

let start = SystemTime::now();

let timestamp = start

.duration_since(UNIX_EPOCH)

.expect("Time went backwards")

.as_millis();

let versions: HashMap<String, serde_json::Value> = document

.as_object()

.context("document must be an object")?

.iter()

.try_fold(HashMap::new(), |mut acc, (key, value)| {

let md5_digest = md5::compute(serde_json::to_string(value)?.as_bytes());

let md5_digest = format!("{md5_digest:x}");

acc.insert(

key.to_owned(),

serde_json::json!({

"last_updated": timestamp,

"md5": md5_digest

}),

);

anyhow::Ok(acc)

})?;

let versions = serde_json::to_value(versions)?;

query = query.bind(source_uuid).bind(document).bind(versions);

}

let results: Vec<(i64, Option<Json>)> = query

.bind(source_uuids)

.fetch_all(&mut *transaction)

.await?;

let dp: Vec<(i64, Json, Option<Json>)> = results

.into_iter()

.zip(batch)

.map(|((id, previous_document), document)| (id, document.to_owned(), previous_document))

.collect();

for (pipeline, parsed_schema) in &mut pipelines {

let ids_to_run_on: Vec<i64> = dp

.iter()

.filter(|(_, document, previous_document)| match previous_document {

Some(previous_document) => parsed_schema

.iter()

.any(|(key, _)| document[key] != previous_document[key]),

None => true,

.map(|(document_id, _, _)| *document_id)

if !ids_to_run_on.is_empty() {

pipeline

.sync_documents(ids_to_run_on, project_info, &mut transaction)

.await

.expect("Failed to execute pipeline");

transaction.commit().await?;

#[tokio::test]

async fn can_add_pipeline_and_upsert_documents_with_parallel_batches() -> anyhow::Result<()> {

internal_init_logger(None, None).ok();

let collection_name = "test_r_c_capaud_107";

let pipeline_name = "test_r_p_capaud_6";

let mut pipeline = Pipeline::new(

pipeline_name,

Some(

json!({

"title": {

"semantic_search": {

"model": "intfloat/e5-small"

}

},

"body": {

"splitter": {

"model": "recursive_character",

"parameters": {

"chunk_size": 1000,

"chunk_overlap": 40

}

},

"semantic_search": {

"model": "hkunlp/instructor-base",

"parameters": {

"instruction": "Represent the Wikipedia document for retrieval"

}

},

"full_text_search": {

"configuration": "english"

}

}

})

.into(),

),

)?;

let mut collection = Collection::new(collection_name, None)?;

collection.add_pipeline(&mut pipeline).await?;

let documents = generate_dummy_documents(20);

collection

.upsert_documents(

documents.clone(),

Some(

json!({

"batch_size": 4,

"parallel_batches": 5

})

.into(),

),

)

.await?;

let pool = get_or_initialize_pool(&None).await?;

let documents_table = format!("{}.documents", collection_name);

let queried_documents: Vec<models::Document> =

sqlx::query_as(&query_builder!("SELECT * FROM %s", documents_table))

.fetch_all(&pool)

.await?;

assert!(queried_documents.len() == 20);

let chunks_table = format!("{}_{}.title_chunks", collection_name, pipeline_name);

let title_chunks: Vec<models::Chunk> =

sqlx::query_as(&query_builder!("SELECT * FROM %s", chunks_table))

.fetch_all(&pool)

.await?;

assert!(title_chunks.len() == 20);

let chunks_table = format!("{}_{}.body_chunks", collection_name, pipeline_name);

let body_chunks: Vec<models::Chunk> =

sqlx::query_as(&query_builder!("SELECT * FROM %s", chunks_table))

.fetch_all(&pool)

.await?;

assert!(body_chunks.len() == 120);

let tsvectors_table = format!("{}_{}.body_tsvectors", collection_name, pipeline_name);

let tsvectors: Vec<models::TSVector> =

sqlx::query_as(&query_builder!("SELECT * FROM %s", tsvectors_table))

.fetch_all(&pool)

.await?;

assert!(tsvectors.len() == 120);

collection.archive().await?;

Ok(())

}