Skip to content

Commit

Permalink
[core] Skip count query when not needed (#9291)
Browse files Browse the repository at this point in the history
  • Loading branch information
@tdraier
tdraier authored Dec 11, 2024
1 parent 1842a7a commit a74674e
Show file tree
Hide file tree
Showing 4 changed files with 62 additions and 38 deletions.
2 changes: 2 additions & 0 deletions core/bin/core_api.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1585,6 +1585,7 @@ async fn data_sources_documents_versions_list(
None => None,
},
&query.latest_hash,
true,
)
.await
{
Expand Down Expand Up @@ -1787,6 +1788,7 @@ async fn data_sources_documents_list(
&document_ids,
limit_offset,
true, // remove system tags
true,
)
.await
{
Expand Down
3 changes: 3 additions & 0 deletions core/src/data_sources/data_source.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1535,6 +1535,7 @@ impl DataSource {
// With top_k documents, we should be guaranteed to have at least top_k chunks, if
// we make the assumption that each document has at least one chunk.
Some((top_k, 0)),
false,
)
.await?;

Expand Down Expand Up @@ -1831,6 +1832,7 @@ impl DataSource {
None,
&None,
&None,
false,
)
.await?;

Expand Down Expand Up @@ -1887,6 +1889,7 @@ impl DataSource {
None,
&None,
&None,
false,
)
.await?;

Expand Down
92 changes: 54 additions & 38 deletions core/src/stores/postgres.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1573,6 +1573,7 @@ impl Store for PostgresStore {
limit_offset: Option<(usize, usize)>,
view_filter: &Option<SearchFilter>,
latest_hash: &Option<String>,
include_count: bool,
) -> Result<(Vec<DocumentVersion>, usize)> {
let project_id = project.project_id();
let data_source_id = data_source_id.to_string();
Expand Down Expand Up @@ -1697,23 +1698,27 @@ impl Store for PostgresStore {
});
}

let total = match limit_offset {
None => versions.len(),
Some(_) => {
let stmt = c
.prepare(
format!(
"SELECT COUNT(*) FROM data_sources_documents dsd \
INNER JOIN data_sources_nodes dsn ON dsn.document=dsd.id \
WHERE {}",
where_clauses.join(" AND ")
let total = if include_count {
match limit_offset {
None => versions.len(),
Some(_) => {
let stmt = c
.prepare(
format!(
"SELECT COUNT(*) FROM data_sources_documents dsd \
INNER JOIN data_sources_nodes dsn ON dsn.document=dsd.id \
WHERE {}",
where_clauses.join(" AND ")
)
.as_str(),
)
.as_str(),
)
.await?;
let t: i64 = c.query_one(&stmt, &params).await?.get(0);
t as usize
.await?;
let t: i64 = c.query_one(&stmt, &params).await?.get(0);
t as usize
}
}
} else {
0
};

Ok((versions, total))
Expand All @@ -1726,6 +1731,7 @@ impl Store for PostgresStore {
filter: &Option<SearchFilter>,
view_filter: &Option<SearchFilter>,
limit_offset: Option<(usize, usize)>,
include_count: bool,
) -> Result<(Vec<String>, usize)> {
let pool = self.pool.clone();
let c = pool.get().await?;
Expand Down Expand Up @@ -1772,14 +1778,19 @@ impl Store for PostgresStore {
params.extend(view_filter_params);

// compute the total count
let count_query = format!(
"SELECT COUNT(*) \
FROM data_sources_documents dsd \
INNER JOIN data_sources_nodes dsn ON dsn.document=dsd.id \
WHERE {}",
where_clauses.join(" AND ")
);
let count: i64 = c.query_one(&count_query, &params).await?.get(0);
let count = if include_count {
let count_query = format!(
"SELECT COUNT(*) \
FROM data_sources_documents dsd \
INNER JOIN data_sources_nodes dsn ON dsn.document=dsd.id \
WHERE {}",
where_clauses.join(" AND ")
);
let count: i64 = c.query_one(&count_query, &params).await?.get(0);
count as usize
} else {
0
};

let mut query = format!(
"SELECT document_id FROM data_sources_documents dsd \
Expand Down Expand Up @@ -1923,6 +1934,7 @@ impl Store for PostgresStore {
document_ids: &Option<Vec<String>>,
limit_offset: Option<(usize, usize)>,
remove_system_tags: bool,
include_count: bool,
) -> Result<(Vec<Document>, usize)> {
let project_id = project.project_id();
let data_source_id = data_source_id.to_string();
Expand Down Expand Up @@ -2053,23 +2065,27 @@ impl Store for PostgresStore {
})
.collect::<Result<Vec<_>>>()?;

let total = match limit_offset {
None => documents.len(),
Some(_) => {
let stmt = c
.prepare(
format!(
"SELECT COUNT(*) FROM data_sources_documents dsd \
INNER JOIN data_sources_nodes dsn ON dsn.document=dsd.id \
WHERE {}",
where_clauses.join(" AND ")
let total: usize = if include_count {
match limit_offset {
None => documents.len(),
Some(_) => {
let stmt = c
.prepare(
format!(
"SELECT COUNT(*) FROM data_sources_documents dsd \
INNER JOIN data_sources_nodes dsn ON dsn.document=dsd.id \
WHERE {}",
where_clauses.join(" AND ")
)
.as_str(),
)
.as_str(),
)
.await?;
let t: i64 = c.query_one(&stmt, &params).await?.get(0);
t as usize
.await?;
let t: i64 = c.query_one(&stmt, &params).await?.get(0);
t as usize
}
}
} else {
0
};

Ok((documents, total))
Expand Down
3 changes: 3 additions & 0 deletions core/src/stores/store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,6 +182,7 @@ pub trait Store {
filter: &Option<SearchFilter>,
view_filter: &Option<SearchFilter>,
limit_offset: Option<(usize, usize)>,
include_count: bool,
) -> Result<(Vec<String>, usize)>;
async fn create_data_source_document(
&self,
Expand Down Expand Up @@ -219,6 +220,7 @@ pub trait Store {
limit_offset: Option<(usize, usize)>,
view_filter: &Option<SearchFilter>,
latest_hash: &Option<String>,
include_count: bool,
) -> Result<(Vec<DocumentVersion>, usize)>;
async fn list_data_source_documents(
&self,
Expand All @@ -228,6 +230,7 @@ pub trait Store {
document_ids: &Option<Vec<String>>,
limit_offset: Option<(usize, usize)>,
remove_system_tags: bool,
include_count: bool,
) -> Result<(Vec<Document>, usize)>;
async fn delete_data_source_document(
&self,
Expand Down

0 comments on commit a74674e

Please sign in to comment.