use std::sync::Arc;
use std::time::Duration;
use uuid::Uuid;
use crate::ILResult;
use crate::catalog::{
CatalogDataType, CatalogDatabase, CatalogHelper, CatalogSchema, CatalogSchemaRef, Column,
DataFileRecord, FieldRecord, INTERNAL_ROW_ID_FIELD_NAME, IndexFileRecord, IndexRecord,
InlineIndexRecord, Row, RowStream, TableRecord, TransactionHelper, inline_row_table_name,
};
use crate::expr::Expr;
use crate::utils::timestamp_ms_from_now;
impl TransactionHelper {
pub(crate) async fn get_namespace_id(
&mut self,
namespace_name: &str,
) -> ILResult<Option<Uuid>> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"namespace_id",
CatalogDataType::Uuid,
false,
)]));
let row = self
.query_single(
&format!(
"SELECT namespace_id FROM indexlake_namespace WHERE namespace_name = '{namespace_name}'"
),
schema,
)
.await?;
match row {
Some(row) => row.uuid(0),
None => Ok(None),
}
}
pub(crate) async fn get_table_id(
&mut self,
namespace_id: &Uuid,
table_name: &str,
) -> ILResult<Option<Uuid>> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"table_id",
CatalogDataType::Uuid,
false,
)]));
let row = self.query_single(&format!("SELECT table_id FROM indexlake_table WHERE namespace_id = {} AND table_name = '{table_name}'", self.catalog.sql_uuid_literal(namespace_id)), schema).await?;
match row {
Some(row) => row.uuid(0),
None => Ok(None),
}
}
pub(crate) async fn get_table_field(
&mut self,
table_id: &Uuid,
field_name: &str,
) -> ILResult<Option<FieldRecord>> {
let catalog_schema = Arc::new(FieldRecord::catalog_schema());
let row = self
.query_single(
&format!(
"SELECT {} FROM indexlake_field WHERE table_id = {} and field_name = {}",
catalog_schema
.select_items(self.catalog.as_ref())
.join(", "),
self.catalog.sql_uuid_literal(table_id),
self.catalog.sql_string_literal(field_name),
),
catalog_schema,
)
.await?;
match row {
Some(row) => {
let field = FieldRecord::from_row(row)?;
Ok(Some(field))
}
None => Ok(None),
}
}
pub(crate) async fn scan_inline_rows(
&mut self,
table_id: &Uuid,
schema: &CatalogSchemaRef,
filters: &[Expr],
limit: Option<usize>,
order_by: Option<Expr>,
) -> ILResult<RowStream<'_>> {
let filter_strs = filters
.iter()
.map(|f| self.catalog.unparse_expr(f))
.collect::<Result<Vec<_>, _>>()?;
let where_clause = if filter_strs.is_empty() {
"".to_string()
} else {
format!(" WHERE {}", filter_strs.join(" AND "))
};
let limit_clause = limit
.map(|limit| format!(" LIMIT {limit}"))
.unwrap_or_default();
let order_by_clause = if let Some(expr) = order_by {
format!(" ORDER BY {}", self.catalog.unparse_expr(&expr)?)
} else {
"".to_string()
};
self.transaction
.query(
&format!(
"SELECT {} FROM {}{where_clause}{order_by_clause}{limit_clause}",
schema.select_items(self.catalog.as_ref()).join(", "),
inline_row_table_name(table_id),
),
Arc::clone(schema),
)
.await
}
pub(crate) async fn count_inline_rows(&mut self, table_id: &Uuid) -> ILResult<i64> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"count",
CatalogDataType::Int64,
false,
)]));
let rows = self
.query_rows(
&format!("SELECT COUNT(1) FROM {}", inline_row_table_name(table_id)),
schema,
)
.await?;
let count = rows[0].int64(0)?.expect("count is not null");
Ok(count)
}
pub(crate) async fn get_data_files(
&mut self,
table_id: &Uuid,
) -> ILResult<Vec<DataFileRecord>> {
let schema = Arc::new(DataFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_data_file WHERE table_id = {}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(table_id),
),
schema,
)
.await?;
let mut data_files = Vec::with_capacity(rows.len());
for row in rows {
data_files.push(DataFileRecord::from_row(row)?);
}
Ok(data_files)
}
pub(crate) async fn get_index_id(
&mut self,
table_id: &Uuid,
index_name: &str,
) -> ILResult<Option<Uuid>> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"index_id",
CatalogDataType::Uuid,
false,
)]));
let row = self.query_single(&format!("SELECT index_id FROM indexlake_index WHERE table_id = {} AND index_name = '{index_name}'", self.catalog.sql_uuid_literal(table_id)), schema).await?;
match row {
Some(row) => row.uuid(0),
None => Ok(None),
}
}
pub(crate) async fn get_table_index_files(
&mut self,
table_id: &Uuid,
) -> ILResult<Vec<IndexFileRecord>> {
let schema = Arc::new(IndexFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_index_file WHERE table_id = {}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(table_id)
),
schema,
)
.await?;
let mut index_files = Vec::with_capacity(rows.len());
for row in rows {
index_files.push(IndexFileRecord::from_row(row)?);
}
Ok(index_files)
}
pub(crate) async fn get_table_fields(&mut self, table_id: &Uuid) -> ILResult<Vec<FieldRecord>> {
let catalog_schema = Arc::new(FieldRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_field WHERE table_id = {} order by field_id asc",
catalog_schema
.select_items(self.catalog.as_ref())
.join(", "),
self.catalog.sql_uuid_literal(table_id)
),
catalog_schema,
)
.await?;
let mut field_records = Vec::with_capacity(rows.len());
for row in rows {
field_records.push(FieldRecord::from_row(row)?);
}
Ok(field_records)
}
}
impl CatalogHelper {
pub(crate) async fn get_namespace_id(&self, namespace_name: &str) -> ILResult<Option<Uuid>> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"namespace_id",
CatalogDataType::Uuid,
false,
)]));
let row = self
.query_single(
&format!(
"SELECT namespace_id FROM indexlake_namespace WHERE namespace_name = '{namespace_name}'"
),
schema,
)
.await?;
match row {
Some(row) => row.uuid(0),
None => Ok(None),
}
}
pub(crate) async fn get_table(
&self,
namespace_id: &Uuid,
table_name: &str,
) -> ILResult<Option<TableRecord>> {
let schema = Arc::new(TableRecord::catalog_schema());
let row = self
.query_single(
&format!(
"SELECT {} FROM indexlake_table WHERE namespace_id = {} AND table_name = '{table_name}'",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(namespace_id)
),
schema,
)
.await?;
match row {
Some(row) => Ok(Some(TableRecord::from_row(row)?)),
None => Ok(None),
}
}
pub(crate) async fn get_table_fields(&self, table_id: &Uuid) -> ILResult<Vec<FieldRecord>> {
let catalog_schema = Arc::new(FieldRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_field WHERE table_id = {} order by field_id asc",
catalog_schema
.select_items(self.catalog.as_ref())
.join(", "),
self.catalog.sql_uuid_literal(table_id)
),
catalog_schema,
)
.await?;
let mut field_records = Vec::with_capacity(rows.len());
for row in rows {
field_records.push(FieldRecord::from_row(row)?);
}
Ok(field_records)
}
pub(crate) async fn get_table_indexes(&self, table_id: &Uuid) -> ILResult<Vec<IndexRecord>> {
let catalog_schema = Arc::new(IndexRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_index WHERE table_id = {}",
catalog_schema
.select_items(self.catalog.as_ref())
.join(", "),
self.catalog.sql_uuid_literal(table_id)
),
catalog_schema,
)
.await?;
let mut indexes = Vec::with_capacity(rows.len());
for row in rows {
indexes.push(IndexRecord::from_row(row)?);
}
Ok(indexes)
}
pub(crate) async fn count_inline_rows(
&self,
table_id: &Uuid,
filters: &[Expr],
) -> ILResult<i64> {
let filter_strs = filters
.iter()
.map(|f| self.catalog.unparse_expr(f))
.collect::<Result<Vec<_>, _>>()?;
let where_clause = if filter_strs.is_empty() {
"".to_string()
} else {
format!(" WHERE {}", filter_strs.join(" AND "))
};
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"count",
CatalogDataType::Int64,
false,
)]));
let rows = self
.query_rows(
&format!(
"SELECT COUNT(1) FROM {}{where_clause}",
inline_row_table_name(table_id)
),
schema,
)
.await?;
let count = rows[0].int64(0)?.expect("count is not null");
Ok(count)
}
pub(crate) async fn scan_inline_rows(
&self,
table_id: &Uuid,
table_schema: &CatalogSchemaRef,
row_ids: Option<&[Uuid]>,
filters: &[Expr],
offset: Option<usize>,
) -> ILResult<RowStream<'static>> {
if let Some(row_ids) = row_ids
&& row_ids.is_empty()
{
return Ok(Box::pin(futures::stream::empty::<ILResult<Row>>()));
}
let mut filter_strs = filters
.iter()
.map(|f| self.catalog.unparse_expr(f))
.collect::<Result<Vec<_>, _>>()?;
if let Some(row_ids) = row_ids {
filter_strs.push(format!(
"{INTERNAL_ROW_ID_FIELD_NAME} IN ({})",
row_ids
.iter()
.map(|id| self.catalog.sql_uuid_literal(id))
.collect::<Vec<_>>()
.join(", ")
));
}
let where_clause = if filter_strs.is_empty() {
"".to_string()
} else {
format!(" WHERE {}", filter_strs.join(" AND "))
};
let offset_clause = if let Some(offset) = offset {
match self.catalog.database() {
CatalogDatabase::Sqlite => format!(" LIMIT {} OFFSET {}", i64::MAX, offset),
CatalogDatabase::Postgres => format!(" OFFSET {}", offset),
}
} else {
"".to_string()
};
self.catalog
.query(
&format!(
"SELECT {} FROM {}{where_clause}{offset_clause}",
table_schema.select_items(self.catalog.as_ref()).join(", "),
inline_row_table_name(table_id),
),
Arc::clone(table_schema),
)
.await
}
pub(crate) async fn count_data_files(&self, table_id: &Uuid) -> ILResult<i64> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"count",
CatalogDataType::Int64,
false,
)]));
let rows = self
.query_rows(
&format!(
"SELECT COUNT(1) FROM indexlake_data_file WHERE table_id = {}",
self.catalog.sql_uuid_literal(table_id)
),
schema,
)
.await?;
let count = rows[0].int64(0)?.expect("count is not null");
Ok(count)
}
pub(crate) async fn get_partitioned_data_files(
&self,
table_id: &Uuid,
offset: usize,
limit: usize,
) -> ILResult<Vec<DataFileRecord>> {
let schema = Arc::new(DataFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_data_file WHERE table_id = {} ORDER BY data_file_id ASC LIMIT {limit} OFFSET {offset}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(table_id),
),
schema,
)
.await?;
let mut data_files = Vec::with_capacity(rows.len());
for row in rows {
data_files.push(DataFileRecord::from_row(row)?);
}
Ok(data_files)
}
pub(crate) async fn get_data_files(&self, table_id: &Uuid) -> ILResult<Vec<DataFileRecord>> {
let schema = Arc::new(DataFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_data_file WHERE table_id = {}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(table_id)
),
schema,
)
.await?;
let mut data_files = Vec::with_capacity(rows.len());
for row in rows {
data_files.push(DataFileRecord::from_row(row)?);
}
Ok(data_files)
}
pub(crate) async fn get_data_files_size(&self, table_id: &Uuid) -> ILResult<i64> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"size",
CatalogDataType::Int64,
true,
)]));
let row = self
.query_single(
&format!(
"SELECT CAST(SUM(size) AS BIGINT) FROM indexlake_data_file WHERE table_id = {}",
self.catalog.sql_uuid_literal(table_id)
),
schema,
)
.await?;
match row {
Some(row) => {
let size = row.int64(0)?.unwrap_or(0);
Ok(size)
}
None => Ok(0),
}
}
pub(crate) async fn get_table_index_files(
&self,
table_id: &Uuid,
) -> ILResult<Vec<IndexFileRecord>> {
let schema = Arc::new(IndexFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_index_file WHERE table_id = {}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(table_id)
),
schema,
)
.await?;
let mut index_files = Vec::with_capacity(rows.len());
for row in rows {
index_files.push(IndexFileRecord::from_row(row)?);
}
Ok(index_files)
}
pub(crate) async fn get_index_files_by_index_id(
&self,
index_id: &Uuid,
) -> ILResult<Vec<IndexFileRecord>> {
let schema = Arc::new(IndexFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_index_file WHERE index_id = {}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(index_id)
),
schema,
)
.await?;
let mut index_files = Vec::with_capacity(rows.len());
for row in rows {
index_files.push(IndexFileRecord::from_row(row)?);
}
Ok(index_files)
}
pub(crate) async fn get_index_files_by_data_file_id(
&self,
data_file_id: &Uuid,
) -> ILResult<Vec<IndexFileRecord>> {
let schema = Arc::new(IndexFileRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_index_file WHERE data_file_id = {}",
schema.select_items(self.catalog.as_ref()).join(", "),
self.catalog.sql_uuid_literal(data_file_id)
),
schema,
)
.await?;
let mut index_files = Vec::with_capacity(rows.len());
for row in rows {
index_files.push(IndexFileRecord::from_row(row)?);
}
Ok(index_files)
}
pub(crate) async fn get_index_files_size(&self, table_id: &Uuid) -> ILResult<i64> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"size",
CatalogDataType::Int64,
true,
)]));
let row = self
.query_single(
&format!(
"SELECT CAST(SUM(size) AS BIGINT) FROM indexlake_index_file WHERE table_id = {}",
self.catalog.sql_uuid_literal(table_id)
),
schema,
)
.await?;
match row {
Some(row) => {
let size = row.int64(0)?.unwrap_or(0);
Ok(size)
}
None => Ok(0),
}
}
pub(crate) async fn task_exists(&self, task_id: &str) -> ILResult<bool> {
let schema = Arc::new(CatalogSchema::new(vec![Column::new(
"task_id",
CatalogDataType::Utf8,
false,
)]));
let current_timestamp_ms = timestamp_ms_from_now(Duration::ZERO);
let rows = self
.query_rows(
&format!(
"SELECT task_id FROM indexlake_task WHERE task_id = {} AND start_at + max_lifetime > {current_timestamp_ms}",
self.catalog.sql_string_literal(task_id)
),
schema,
)
.await?;
Ok(!rows.is_empty())
}
pub(crate) async fn get_inline_indexes(
&self,
index_ids: &[Uuid],
) -> ILResult<Vec<InlineIndexRecord>> {
if index_ids.is_empty() {
return Ok(Vec::new());
}
let schema = Arc::new(InlineIndexRecord::catalog_schema());
let rows = self
.query_rows(
&format!(
"SELECT {} FROM indexlake_inline_index WHERE index_id IN ({}) ORDER BY created_at ASC",
schema.select_items(self.catalog.as_ref()).join(", "),
index_ids
.iter()
.map(|id| self.catalog.sql_uuid_literal(id))
.collect::<Vec<_>>()
.join(", ")
),
schema,
)
.await?;
let mut inline_indexes = Vec::with_capacity(rows.len());
for row in rows {
inline_indexes.push(InlineIndexRecord::from_row(row)?);
}
Ok(inline_indexes)
}
}