use crate::{
RetryPolicy,
config::{self},
handlers::NamespaceIdent,
sessions::{self},
};
use anyhow::{Context, Result};
use lance::Dataset;
use lance::dataset::builder::DatasetBuilder;
use lance::dataset::optimize::{CompactionOptions, compact_files};
use lance::dataset::write::merge_insert::SourceDedupeBehavior;
use lance::dataset::{MergeInsertBuilder, WhenMatched, WhenNotMatched, WriteMode};
use lance::deps::arrow_array::{RecordBatch, RecordBatchIterator};
use lance::index::DatasetIndexExt;
use lance::index::DatasetIndexInternalExt;
use lance::index::vector::VectorIndexParams;
use lance::session::Session;
use lance_index::IndexType;
use lance_index::optimize::OptimizeOptions;
use lance_index::scalar::{BuiltinIndexType, InvertedIndexParams, ScalarIndexParams};
use lance_io::object_store::{
ObjectStore, ObjectStoreParams, ObjectStoreRegistry, StorageOptionsAccessor,
};
use lance_linalg::distance::MetricType;
use lance_namespace::LanceNamespace;
use lance_namespace::error::NamespaceError;
use lance_namespace::models::DescribeTableRequest;
use lance_namespace_impls::ConnectBuilder;
use std::{
collections::HashMap,
sync::Arc,
time::{Duration, Instant},
};
use tokio::sync::{Mutex, OnceCell};
use tokio_stream::StreamExt;
use url::Url;
pub const VECTOR_INDEX_ACTIVATION_ROWS: usize = 100_000;
pub const DEFAULT_INDEX_LAG_THRESHOLD: usize = 4;
static INDEX_LAG_THRESHOLD_RUNTIME: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
pub fn init_index_lag_threshold(value: usize) {
INDEX_LAG_THRESHOLD_RUNTIME.get_or_init(|| value);
}
pub fn index_lag_threshold() -> usize {
INDEX_LAG_THRESHOLD_RUNTIME
.get()
.copied()
.unwrap_or(DEFAULT_INDEX_LAG_THRESHOLD)
}
#[derive(Debug, Clone)]
pub struct IndexIntent {
pub name: &'static str,
pub column: &'static str,
pub trigger: IndexTrigger,
pub params: IndexParamsKind,
}
#[derive(Debug, Clone)]
pub enum IndexTrigger {
OnAnyRows,
OnNonNullCount {
column: &'static str,
threshold: usize,
},
}
#[derive(Debug, Clone)]
pub enum IndexParamsKind {
Scalar(BuiltinIndexType),
InvertedFtsNgram { min: u32, max: u32 },
IvfPqCosine {
sub_vectors: usize,
num_bits: u8,
max_iters: usize,
},
}
impl IndexTrigger {
async fn should_create(&self, dataset: &Dataset) -> Result<bool> {
match self {
Self::OnAnyRows => Ok(dataset.count_rows(None).await? > 0),
Self::OnNonNullCount { column, threshold } => {
let count = dataset
.count_rows(Some(format!("{column} IS NOT NULL")))
.await?;
Ok(count >= *threshold)
}
}
}
}
impl IndexParamsKind {
fn index_type(&self) -> IndexType {
match self {
Self::Scalar(BuiltinIndexType::Bitmap) => IndexType::Bitmap,
Self::Scalar(_) => IndexType::BTree,
Self::InvertedFtsNgram { .. } => IndexType::Inverted,
Self::IvfPqCosine { .. } => IndexType::Vector,
}
}
async fn build(&self, dataset: &Dataset) -> Result<Box<dyn lance::index::IndexParams>> {
match self {
Self::Scalar(kind) => Ok(Box::new(ScalarIndexParams::for_builtin(kind.clone()))),
Self::InvertedFtsNgram { min, max } => Ok(Box::new(
InvertedIndexParams::default()
.base_tokenizer("ngram".to_owned())
.ngram_min_length(*min)
.ngram_max_length(*max)
.stem(false)
.remove_stop_words(false),
)),
Self::IvfPqCosine {
sub_vectors,
num_bits,
max_iters,
} => {
let count = dataset
.count_rows(Some("vector IS NOT NULL".to_owned()))
.await?;
let partitions = count.checked_div(4096).unwrap_or(0).max(1);
Ok(Box::new(VectorIndexParams::ivf_pq(
partitions,
*num_bits,
*sub_vectors,
MetricType::Cosine,
*max_iters,
)))
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IndexStatus {
pub table: Table,
pub intent_name: String,
pub fragments_covered: usize,
pub unindexed_rows: usize,
pub exists: bool,
}
#[derive(Debug, Clone, Copy)]
pub struct ConflictExhausted {
pub attempts: u8,
}
impl std::fmt::Display for ConflictExhausted {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
formatter,
"commit conflict exhausted after {} attempt(s)",
self.attempts
)
}
}
impl std::error::Error for ConflictExhausted {}
#[derive(Debug)]
pub enum PhaseOutcome {
Ok,
Noop,
SkippedConflict,
Failed(anyhow::Error),
NotAttempted,
}
impl PhaseOutcome {
pub fn is_failed(&self) -> bool {
matches!(self, Self::Failed(_))
}
}
#[derive(Debug)]
pub struct TableOptimizeOutcome {
pub table: Table,
pub indices: PhaseOutcome,
pub compaction: PhaseOutcome,
}
#[derive(Debug, Clone)]
pub enum OptimizeEvent {
PhaseStart {
table: Table,
phase: OptimizePhase,
detail: Option<String>,
},
PhaseDone {
table: Table,
phase: OptimizePhase,
elapsed_ms: u64,
},
}
#[derive(Debug, Clone, Copy)]
pub enum OptimizePhase {
Compact,
Cleanup,
IndexCreate,
IndexRebuild,
IndexAppend,
}
impl OptimizePhase {
pub fn label(self) -> &'static str {
match self {
Self::Compact => "compact",
Self::Cleanup => "cleanup",
Self::IndexCreate => "index-create",
Self::IndexRebuild => "index-rebuild",
Self::IndexAppend => "index-append",
}
}
}
pub type OptimizeProgressFn = Box<dyn Fn(OptimizeEvent) + Send + Sync>;
fn emit(progress: Option<&OptimizeProgressFn>, event: OptimizeEvent) {
if let Some(callback) = progress {
callback(event);
}
}
pub fn is_commit_conflict(error: &anyhow::Error) -> bool {
error.downcast_ref::<lance::Error>().is_some_and(|err| {
matches!(
err,
lance::Error::CommitConflict { .. }
| lance::Error::RetryableCommitConflict { .. }
| lance::Error::TooMuchWriteContention { .. }
)
})
}
fn is_conflict_exhausted(error: &anyhow::Error) -> bool {
error.chain().any(|cause| cause.is::<ConflictExhausted>())
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct TableSizes {
pub sessions: u64,
pub messages: u64,
pub parts: u64,
pub other: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ScalarValue {
String(String),
Int32(i32),
Raw(String),
}
impl From<&str> for ScalarValue {
fn from(value: &str) -> Self {
Self::String(value.to_owned())
}
}
impl From<String> for ScalarValue {
fn from(value: String) -> Self {
Self::String(value)
}
}
impl From<i32> for ScalarValue {
fn from(value: i32) -> Self {
Self::Int32(value)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Predicate {
Eq(&'static str, ScalarValue),
Ne(&'static str, ScalarValue),
IsNull(&'static str),
IsNotNull(&'static str),
In(&'static str, Vec<ScalarValue>),
LikeContains(&'static str, String),
Regex(&'static str, String),
Gte(&'static str, ScalarValue),
Lte(&'static str, ScalarValue),
And(Vec<Predicate>),
Or(Vec<Predicate>),
}
impl Predicate {
pub fn to_lance(&self) -> String {
match self {
Self::Eq(column, value) => format!("{column} = {}", value.to_lance()),
Self::Ne(column, value) => format!("{column} <> {}", value.to_lance()),
Self::IsNull(column) => format!("{column} IS NULL"),
Self::IsNotNull(column) => format!("{column} IS NOT NULL"),
Self::In(column, values) => {
let values = values
.iter()
.map(ScalarValue::to_lance)
.collect::<Vec<_>>()
.join(", ");
format!("{column} IN ({values})")
}
Self::LikeContains(column, value) => {
format!("{column} LIKE {} ESCAPE '\\'", like_contains(value))
}
Self::Regex(column, pattern) => {
format!("regexp_like({column}, {})", quoted_string(pattern))
}
Self::Gte(column, value) => format!("{column} >= {}", value.to_lance()),
Self::Lte(column, value) => format!("{column} <= {}", value.to_lance()),
Self::And(predicates) => predicates
.iter()
.map(Self::to_lance)
.filter(|predicate| !predicate.is_empty())
.collect::<Vec<_>>()
.join(" AND "),
Self::Or(predicates) => {
let body = predicates
.iter()
.map(Self::to_lance)
.filter(|predicate| !predicate.is_empty())
.collect::<Vec<_>>()
.join(" OR ");
if body.is_empty() {
String::new()
} else {
format!("({body})")
}
}
}
}
}
#[derive(Default)]
pub struct ScanOpts<'a> {
pub predicate: Option<&'a Predicate>,
pub projection: Option<&'a [&'a str]>,
}
impl<'a> ScanOpts<'a> {
pub fn project_only(projection: &'a [&'a str]) -> Self {
Self {
predicate: None,
projection: Some(projection),
}
}
pub fn with_predicate_and_projection(
predicate: &'a Predicate,
projection: &'a [&'a str],
) -> Self {
Self {
predicate: Some(predicate),
projection: Some(projection),
}
}
}
impl ScalarValue {
fn to_lance(&self) -> String {
match self {
Self::String(value) => quoted_string(value),
Self::Int32(value) => value.to_string(),
Self::Raw(value) => value.clone(),
}
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct RuntimeCaps {
pub index_cache_bytes: Option<usize>,
pub metadata_cache_bytes: Option<usize>,
}
impl RuntimeCaps {
pub fn from_config(config: &crate::config::RuntimeConfig) -> Self {
Self {
index_cache_bytes: config.index_cache_bytes,
metadata_cache_bytes: config.metadata_cache_bytes,
}
}
}
const LOCAL_INDEX_CACHE_BYTES: usize = 256 * 1024 * 1024;
const LOCAL_METADATA_CACHE_BYTES: usize = 128 * 1024 * 1024;
const REMOTE_INDEX_CACHE_BYTES: usize = 2 * 1024 * 1024 * 1024;
const REMOTE_METADATA_CACHE_BYTES: usize = 512 * 1024 * 1024;
fn resolve_cache_caps(location: &Url, caps: RuntimeCaps) -> (usize, usize) {
let (index_default, metadata_default) = if config::is_local(location) {
(LOCAL_INDEX_CACHE_BYTES, LOCAL_METADATA_CACHE_BYTES)
} else {
(REMOTE_INDEX_CACHE_BYTES, REMOTE_METADATA_CACHE_BYTES)
};
(
caps.index_cache_bytes.unwrap_or(index_default),
caps.metadata_cache_bytes.unwrap_or(metadata_default),
)
}
pub struct Handle {
datasets: DatasetSet,
retry: RetryPolicy,
#[allow(dead_code)]
session: Arc<Session>,
nm: Arc<dyn LanceNamespace>,
nm_ident: NamespaceIdent,
storage_options: HashMap<String, String>,
location: Url,
parts_refresh_after: Duration,
}
impl std::fmt::Debug for Handle {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("Handle")
.field("datasets", &self.datasets)
.field("retry", &self.retry)
.field("nm_ident", &self.nm_ident)
.field("storage_options", &self.storage_options)
.field("location", &self.location)
.finish()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Table {
Sessions,
Messages,
Parts,
}
impl Table {
pub fn as_str(self) -> &'static str {
self.label()
}
fn label(self) -> &'static str {
match self {
Self::Sessions => "sessions",
Self::Messages => "messages",
Self::Parts => "parts",
}
}
}
#[derive(Debug)]
struct DatasetSet {
sessions: Mutex<CachedDataset>,
messages: Mutex<CachedDataset>,
parts: OnceCell<Mutex<CachedDataset>>,
}
#[derive(Debug)]
struct CachedDataset {
dataset: Dataset,
last_refresh: Instant,
refresh_after: Duration,
}
impl CachedDataset {
async fn latest(&mut self) -> Result<Dataset> {
if self.last_refresh.elapsed() >= self.refresh_after {
self.dataset.checkout_latest().await?;
self.last_refresh = Instant::now();
}
Ok(self.dataset.clone())
}
fn replace(&mut self, dataset: Dataset) {
self.dataset = dataset;
self.last_refresh = Instant::now();
}
}
impl Handle {
pub async fn open(location: &Url) -> Result<Self> {
Self::open_with_options(location, HashMap::new(), RuntimeCaps::default()).await
}
pub async fn open_with_options(
location: &Url,
mut storage_options: HashMap<String, String>,
caps: RuntimeCaps,
) -> Result<Self> {
if let Some(path) = config::local_path(location) {
tokio::fs::create_dir_all(&path)
.await
.with_context(|| format!("failed to create data dir {}", path.display()))?;
} else {
apply_remote_storage_defaults(&mut storage_options);
}
let (index_cache_bytes, metadata_cache_bytes) = resolve_cache_caps(location, caps);
let session = Arc::new(Session::new(
index_cache_bytes,
metadata_cache_bytes,
Arc::new(ObjectStoreRegistry::default()),
));
let root = location.as_str().trim_end_matches('/').to_string();
let mut connect = ConnectBuilder::new("dir")
.property("root", root)
.session(session.clone());
for (key, value) in &storage_options {
connect = connect.property(format!("storage.{key}"), value.clone());
}
let nm: Arc<dyn LanceNamespace> = connect
.connect()
.await
.context("failed to connect lance Directory namespace")?;
let nm_ident = NamespaceIdent::root();
let refresh_after = if config::is_local(location) {
Duration::ZERO
} else {
Duration::from_secs(5)
};
let handle = Self {
datasets: DatasetSet {
sessions: Mutex::new(CachedDataset {
dataset: open_or_create_via_ns(
&nm,
&nm_ident,
sessions::SESSIONS,
sessions::session_schema(),
&session,
&storage_options,
)
.await?,
last_refresh: Instant::now(),
refresh_after,
}),
messages: Mutex::new(CachedDataset {
dataset: open_or_create_via_ns(
&nm,
&nm_ident,
sessions::MESSAGES,
sessions::message_schema(),
&session,
&storage_options,
)
.await?,
last_refresh: Instant::now(),
refresh_after,
}),
parts: OnceCell::new(),
},
retry: RetryPolicy::default(),
session,
nm,
nm_ident,
storage_options,
location: location.clone(),
parts_refresh_after: refresh_after,
};
Ok(handle)
}
pub fn location(&self) -> &Url {
&self.location
}
pub fn storage_options(&self) -> &HashMap<String, String> {
&self.storage_options
}
pub async fn row_counts(&self) -> Result<(usize, usize, usize)> {
Ok((
self.count_rows(Table::Sessions).await?,
self.count_rows(Table::Messages).await?,
self.count_rows(Table::Parts).await?,
))
}
pub(crate) async fn merge_insert(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
) -> Result<u64> {
self.merge(
table,
batch,
row_count,
"merge_insert",
WhenMatched::DoNothing,
WhenNotMatched::InsertAll,
)
.await
}
pub(crate) async fn merge_update(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
) -> Result<u64> {
self.merge(
table,
batch,
row_count,
"merge_update",
WhenMatched::UpdateAll,
WhenNotMatched::DoNothing,
)
.await
}
async fn merge(
&self,
table: Table,
batch: RecordBatch,
row_count: usize,
op: &'static str,
when_matched: WhenMatched,
when_not_matched: WhenNotMatched,
) -> Result<u64> {
if row_count == 0 {
return Ok(0);
}
let started = Instant::now();
let result = self
.retry_lance(table.label(), || async {
let mut cached = self.cached(table).await?.lock().await;
let existing = cached.latest().await?;
let reader = RecordBatchIterator::new([Ok(batch.clone())], batch.schema());
let mut builder = MergeInsertBuilder::try_new(Arc::new(existing), Vec::new())?;
builder.when_matched(when_matched.clone());
builder.when_not_matched(when_not_matched.clone());
builder.source_dedupe_behavior(SourceDedupeBehavior::FirstSeen);
builder.skip_auto_cleanup(true);
let (dataset, stats) = builder
.try_build()?
.execute_reader(Box::new(reader))
.await?;
cached.replace(dataset.as_ref().clone());
Ok((
stats.num_inserted_rows + stats.num_updated_rows,
stats.num_skipped_duplicates,
))
})
.await;
let skipped = result.as_ref().map(|(_, s)| *s).unwrap_or(0);
tracing::info!(
target: "pond::perf",
op,
table = %table.label(),
rows = row_count,
elapsed_ms = started.elapsed().as_millis() as u64,
skipped,
"merge",
);
result.map(|(affected, _)| affected)
}
pub async fn optimize_table(
&self,
table: Table,
intents: &[IndexIntent],
progress: Option<&OptimizeProgressFn>,
cleanup: crate::sessions::CleanupConfig,
) -> TableOptimizeOutcome {
let compaction = self
.run_optimize_compact_phase(table, progress, cleanup)
.await;
let indices = self
.run_optimize_indices_phase(table, intents, progress)
.await;
TableOptimizeOutcome {
table,
indices,
compaction,
}
}
pub async fn optimize_table_indices_only(
&self,
table: Table,
intents: &[IndexIntent],
progress: Option<&OptimizeProgressFn>,
) -> PhaseOutcome {
self.run_optimize_indices_phase(table, intents, progress)
.await
}
async fn run_optimize_indices_phase(
&self,
table: Table,
intents: &[IndexIntent],
progress: Option<&OptimizeProgressFn>,
) -> PhaseOutcome {
if intents.is_empty() {
return PhaseOutcome::Noop;
}
let result = self
.retry_lance(table.label(), || async {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
let did_work =
optimize_table_indices(&mut dataset, intents, table, progress).await?;
guard.replace(dataset);
Ok::<_, anyhow::Error>(did_work)
})
.await;
match result {
Ok(true) => PhaseOutcome::Ok,
Ok(false) => PhaseOutcome::Noop,
Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
Err(error) => PhaseOutcome::Failed(error),
}
}
async fn run_optimize_compact_phase(
&self,
table: Table,
progress: Option<&OptimizeProgressFn>,
cleanup: crate::sessions::CleanupConfig,
) -> PhaseOutcome {
let result = self
.retry_lance(table.label(), || async {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
optimize_table_compact(&mut dataset, table, progress, cleanup).await?;
guard.replace(dataset);
Ok::<_, anyhow::Error>(())
})
.await;
match result {
Ok(()) => PhaseOutcome::Ok,
Err(error) if is_conflict_exhausted(&error) => PhaseOutcome::SkippedConflict,
Err(error) => PhaseOutcome::Failed(error),
}
}
pub async fn rebuild_index(&self, table: Table, intent: &IndexIntent) -> Result<()> {
self.retry_lance(table.label(), || async {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
rebuild_index(&mut dataset, intent).await?;
guard.replace(dataset);
Ok(())
})
.await
}
pub async fn index_status(
&self,
table: Table,
intents: &[IndexIntent],
) -> Result<Vec<IndexStatus>> {
let dataset = self.dataset(table).await?;
index_status(table, &dataset, intents).await
}
pub(crate) async fn dataset(&self, table: Table) -> Result<Dataset> {
let mut cached = self.cached(table).await?.lock().await;
cached.latest().await
}
pub(crate) async fn scanner(
&self,
table: Table,
predicate: Option<&Predicate>,
) -> Result<lance::dataset::scanner::Scanner> {
let dataset = self.dataset(table).await?;
scanner_with_prefilter(&dataset, predicate)
}
pub async fn scan(
&self,
table: Table,
opts: ScanOpts<'_>,
) -> Result<lance::dataset::scanner::Scanner> {
let mut scanner = self.scanner(table, opts.predicate).await?;
if let Some(projection) = opts.projection {
scanner.project(projection)?;
}
Ok(scanner)
}
pub(crate) async fn scan_batch(
&self,
table: Table,
predicate: Option<&Predicate>,
projection: &[&str],
) -> Result<RecordBatch> {
let opts = ScanOpts {
predicate,
projection: (!projection.is_empty()).then_some(projection),
};
self.scan(table, opts)
.await?
.try_into_batch()
.await
.context("scan failed")
}
pub async fn count_rows(&self, table: Table) -> Result<usize> {
self.dataset(table)
.await?
.count_rows(None)
.await
.map_err(Into::into)
}
#[cfg(test)]
pub(crate) async fn messages_index_names(&self) -> Result<Vec<String>> {
let dataset = self.dataset(Table::Messages).await?;
let indices = dataset.load_indices().await?;
Ok(indices.iter().map(|index| index.name.clone()).collect())
}
pub(crate) async fn unindexed_row_count(
&self,
table: Table,
index_name: &str,
) -> Result<usize> {
let dataset = self.dataset(table).await?;
let fragments = dataset
.unindexed_fragments(index_name)
.await
.with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
Ok(fragments
.iter()
.map(|fragment| fragment.num_rows().unwrap_or(0))
.sum())
}
pub(crate) async fn drop_index(&self, table: Table, name: &str) -> Result<()> {
let mut guard = self.cached(table).await?.lock().await;
let mut dataset = guard.latest().await?;
dataset
.drop_index(name)
.await
.with_context(|| format!("drop_index({name}) failed for {}", table.label()))?;
guard.replace(dataset);
Ok(())
}
async fn table_location(&self, table_name: &str) -> Result<String> {
let request = DescribeTableRequest {
id: Some(self.nm_ident.as_table_id(table_name)),
..Default::default()
};
let response = self
.nm
.describe_table(request)
.await
.with_context(|| format!("failed to describe table {table_name}"))?;
response
.location
.with_context(|| format!("namespace returned no location for table {table_name}"))
}
pub async fn table_sizes(&self) -> Result<TableSizes> {
let registry = Arc::new(ObjectStoreRegistry::default());
let params = ObjectStoreParams {
storage_options_accessor: (!self.storage_options.is_empty()).then(|| {
Arc::new(StorageOptionsAccessor::with_static_options(
self.storage_options.clone(),
))
}),
..Default::default()
};
let sessions = self
.listed_size(
®istry,
¶ms,
&self.table_location(sessions::SESSIONS).await?,
)
.await?;
let messages = self
.listed_size(
®istry,
¶ms,
&self.table_location(sessions::MESSAGES).await?,
)
.await?;
let parts = self
.listed_size(
®istry,
¶ms,
&self.table_location(sessions::PARTS).await?,
)
.await?;
let root_total = self
.listed_size(®istry, ¶ms, self.location.as_str())
.await?;
let other = root_total.saturating_sub(sessions + messages + parts);
Ok(TableSizes {
sessions,
messages,
parts,
other,
})
}
async fn listed_size(
&self,
registry: &Arc<ObjectStoreRegistry>,
params: &ObjectStoreParams,
uri: &str,
) -> Result<u64> {
let (store, base) = ObjectStore::from_uri_and_params(registry.clone(), uri, params)
.await
.with_context(|| format!("failed to open object store for {uri}"))?;
let mut listing = store.list(Some(base));
let mut total = 0u64;
while let Some(meta) = listing.next().await {
let meta = meta.with_context(|| format!("listing {uri} failed"))?;
total += meta.size;
}
Ok(total)
}
async fn cached(&self, table: Table) -> Result<&Mutex<CachedDataset>> {
match table {
Table::Sessions => Ok(&self.datasets.sessions),
Table::Messages => Ok(&self.datasets.messages),
Table::Parts => self.parts_cached().await,
}
}
async fn parts_cached(&self) -> Result<&Mutex<CachedDataset>> {
self.datasets
.parts
.get_or_try_init(|| async {
let dataset = open_or_create_via_ns(
&self.nm,
&self.nm_ident,
sessions::PARTS,
sessions::part_schema(),
&self.session,
&self.storage_options,
)
.await?;
Ok::<_, anyhow::Error>(Mutex::new(CachedDataset {
dataset,
last_refresh: Instant::now(),
refresh_after: self.parts_refresh_after,
}))
})
.await
}
async fn retry_lance<T, Fut, Op>(&self, label: &str, mut operation: Op) -> Result<T>
where
Fut: std::future::Future<Output = Result<T>>,
Op: FnMut() -> Fut,
{
let mut attempt = 0u8;
loop {
attempt = attempt.saturating_add(1);
match operation().await {
Ok(value) => return Ok(value),
Err(error) if attempt < self.retry.attempts => {
let backoff = self.backoff(attempt);
let error_chain = format!("{error:#}");
tracing::warn!(
label,
attempt,
?backoff,
error = %error_chain,
"retrying Lance operation"
);
tokio::time::sleep(backoff).await;
}
Err(error) => {
let error_chain = format!("{error:#}");
tracing::warn!(
label,
attempt,
error = %error_chain,
"Lance operation exhausted retries"
);
if is_commit_conflict(&error) {
return Err(error.context(ConflictExhausted { attempts: attempt }));
}
return Err(error);
}
}
}
}
fn backoff(&self, attempt: u8) -> Duration {
let shift = u32::from(attempt.saturating_sub(1));
let multiplier = 1u32.checked_shl(shift).unwrap_or(u32::MAX);
let base = self.retry.initial_backoff.saturating_mul(multiplier);
let factor = (1.0 + self.retry.jitter * (fastrand::f64() * 2.0 - 1.0)).max(0.0);
base.mul_f64(factor).min(self.retry.max_backoff)
}
}
async fn optimize_table_compact(
dataset: &mut Dataset,
table: Table,
progress: Option<&OptimizeProgressFn>,
cleanup: crate::sessions::CleanupConfig,
) -> Result<()> {
let compaction = CompactionOptions {
defer_index_remap: false,
..CompactionOptions::default()
};
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::Compact,
detail: None,
},
);
let started = Instant::now();
compact_files(dataset, compaction, None).await?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::Compact,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::Cleanup,
detail: None,
},
);
let started = Instant::now();
dataset
.cleanup_old_versions(
cleanup.older_than,
Some(cleanup.delete_unverified),
Some(false),
)
.await
.context("cleanup_old_versions failed during index optimize")?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::Cleanup,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
Ok(())
}
async fn optimize_table_indices(
dataset: &mut Dataset,
intents: &[IndexIntent],
table: Table,
progress: Option<&OptimizeProgressFn>,
) -> Result<bool> {
let existing = dataset.load_indices().await?;
let existing_names: std::collections::HashSet<String> =
existing.iter().map(|index| index.name.clone()).collect();
let mut append_indices: Vec<String> = Vec::new();
let mut did_work = false;
for intent in intents {
let exists = existing_names.contains(intent.name);
if !exists {
if !intent.trigger.should_create(dataset).await? {
continue;
}
let params = intent.params.build(dataset).await?;
let index_type = intent.params.index_type();
tracing::info!(
index = intent.name,
column = intent.column,
"creating Lance index (trigger fired)",
);
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexCreate,
detail: Some(intent.name.to_owned()),
},
);
let started = Instant::now();
dataset
.create_index(
&[intent.column],
index_type,
Some(intent.name.to_owned()),
params.as_ref(),
false,
)
.await
.with_context(|| format!("failed to create index {}", intent.name))?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexCreate,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
did_work = true;
continue;
}
let unindexed = dataset.unindexed_fragments(intent.name).await?;
if unindexed.is_empty() {
continue;
}
if unindexed.len() < index_lag_threshold() {
continue;
}
match intent.params {
IndexParamsKind::Scalar(BuiltinIndexType::BTree) => {
let params = intent.params.build(dataset).await?;
let index_type = intent.params.index_type();
tracing::debug!(
target: "pond::perf",
index = intent.name,
column = intent.column,
"rebuilding Lance BTree index",
);
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexRebuild,
detail: Some(intent.name.to_owned()),
},
);
let started = Instant::now();
dataset
.create_index(
&[intent.column],
index_type,
Some(intent.name.to_owned()),
params.as_ref(),
true,
)
.await
.with_context(|| format!("failed to rebuild index {}", intent.name))?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexRebuild,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
did_work = true;
}
IndexParamsKind::Scalar(BuiltinIndexType::Bitmap)
| IndexParamsKind::InvertedFtsNgram { .. }
| IndexParamsKind::IvfPqCosine { .. } => {
append_indices.push(intent.name.to_owned());
}
IndexParamsKind::Scalar(_) => {
let params = intent.params.build(dataset).await?;
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexRebuild,
detail: Some(intent.name.to_owned()),
},
);
let started = Instant::now();
dataset
.create_index(
&[intent.column],
intent.params.index_type(),
Some(intent.name.to_owned()),
params.as_ref(),
true,
)
.await
.with_context(|| format!("failed to rebuild index {}", intent.name))?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexRebuild,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
did_work = true;
}
}
}
if !append_indices.is_empty() {
let to_append = append_indices.clone();
emit(
progress,
OptimizeEvent::PhaseStart {
table,
phase: OptimizePhase::IndexAppend,
detail: Some(append_indices.join(", ")),
},
);
let started = Instant::now();
dataset
.optimize_indices(&OptimizeOptions::append().index_names(to_append))
.await
.context("optimize_indices(append) failed during index optimize")?;
emit(
progress,
OptimizeEvent::PhaseDone {
table,
phase: OptimizePhase::IndexAppend,
elapsed_ms: started.elapsed().as_millis() as u64,
},
);
tracing::debug!(
target: "pond::perf",
indices = ?append_indices,
"appended trailing fragments into indices",
);
did_work = true;
}
Ok(did_work)
}
async fn rebuild_index(dataset: &mut Dataset, intent: &IndexIntent) -> Result<()> {
if !intent.trigger.should_create(dataset).await? {
return Ok(());
}
let params = intent.params.build(dataset).await?;
dataset
.create_index(
&[intent.column],
intent.params.index_type(),
Some(intent.name.to_owned()),
params.as_ref(),
true,
)
.await
.with_context(|| format!("failed to rebuild index {}", intent.name))?;
Ok(())
}
async fn index_status(
table: Table,
dataset: &Dataset,
intents: &[IndexIntent],
) -> Result<Vec<IndexStatus>> {
let existing = dataset.load_indices().await?;
let existing_names: std::collections::HashSet<String> =
existing.iter().map(|index| index.name.clone()).collect();
let total_fragments = dataset.get_fragments().len();
let total_rows = dataset.count_rows(None).await?;
let mut statuses = Vec::with_capacity(intents.len());
for intent in intents {
let exists = existing_names.contains(intent.name);
if !exists {
statuses.push(IndexStatus {
table,
intent_name: intent.name.to_owned(),
fragments_covered: 0,
unindexed_rows: total_rows,
exists,
});
continue;
}
let unindexed = dataset
.unindexed_fragments(intent.name)
.await
.with_context(|| format!("unindexed_fragments failed for {}", table.label()))?;
let unindexed_rows = unindexed
.iter()
.map(|fragment| fragment.num_rows().unwrap_or(0))
.sum();
statuses.push(IndexStatus {
table,
intent_name: intent.name.to_owned(),
fragments_covered: total_fragments.saturating_sub(unindexed.len()),
unindexed_rows,
exists,
});
}
Ok(statuses)
}
async fn open_or_create_via_ns(
nm: &Arc<dyn LanceNamespace>,
nm_ident: &NamespaceIdent,
table_name: &str,
schema: lance::deps::arrow_schema::SchemaRef,
session: &Arc<Session>,
storage_options: &HashMap<String, String>,
) -> Result<Dataset> {
let table_id = nm_ident.as_table_id(table_name);
let request = DescribeTableRequest {
id: Some(table_id.clone()),
..Default::default()
};
match nm.describe_table(request).await {
Ok(response) => {
let location = response.location.with_context(|| {
format!("namespace returned no location for table {table_name}")
})?;
let mut builder = DatasetBuilder::from_uri(&location).with_session(session.clone());
if !storage_options.is_empty() {
builder = builder.with_storage_options(storage_options.clone());
}
let dataset = builder
.load()
.await
.with_context(|| format!("failed to open table {table_name}"))?;
ensure_schema_matches(&dataset, schema.as_ref(), table_name)?;
return Ok(dataset);
}
Err(error) => match &error {
lance::Error::Namespace { source, .. }
if matches!(
source.downcast_ref::<NamespaceError>(),
Some(NamespaceError::TableNotFound { .. })
) =>
{
}
_ => {
return Err(anyhow::Error::from(error))
.with_context(|| format!("failed to describe table {table_name}"));
}
},
}
let mut write_params = sessions::write_params_for_create();
write_params.session = Some(session.clone());
write_params.mode = WriteMode::Create;
if !storage_options.is_empty() {
write_params.store_params = Some(ObjectStoreParams {
storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
storage_options.clone(),
))),
..Default::default()
});
}
let reader = sessions::empty_reader(schema)?;
Dataset::write_into_namespace(reader, nm.clone(), table_id, Some(write_params))
.await
.with_context(|| format!("failed to create table {table_name}"))
}
fn scanner_with_prefilter(
dataset: &Dataset,
predicate: Option<&Predicate>,
) -> Result<lance::dataset::scanner::Scanner> {
let mut scanner = dataset.scan();
scanner.prefilter(true);
if let Some(predicate) = predicate {
let filter = predicate.to_lance();
if !filter.is_empty() {
scanner.filter(&filter)?;
}
}
Ok(scanner)
}
fn ensure_schema_matches(
dataset: &Dataset,
expected: &lance::deps::arrow_schema::Schema,
table_name: &str,
) -> Result<()> {
use lance::deps::arrow_schema::DataType;
use std::collections::BTreeSet;
let actual = lance::deps::arrow_schema::Schema::from(dataset.schema());
let actual_names: BTreeSet<&str> = actual.fields().iter().map(|f| f.name().as_str()).collect();
let expected_names: BTreeSet<&str> = expected
.fields()
.iter()
.map(|f| f.name().as_str())
.collect();
if actual_names != expected_names {
anyhow::bail!(
"table {table_name} has columns {actual_names:?} but this pond build expects \
{expected_names:?} - the on-disk store predates a schema change; delete the \
data directory and re-run `pond ingest`",
);
}
for actual_field in actual.fields() {
let Some(expected_field) = expected.field_with_name(actual_field.name()).ok() else {
continue;
};
if let (DataType::FixedSizeList(_, actual_dim), DataType::FixedSizeList(_, expected_dim)) =
(actual_field.data_type(), expected_field.data_type())
&& actual_dim != expected_dim
{
tracing::warn!(
table = table_name,
column = actual_field.name(),
actual_dim,
expected_dim,
"embedding dimension differs from config; open proceeds because model swaps are operator-driven",
);
}
}
Ok(())
}
fn apply_remote_storage_defaults(options: &mut HashMap<String, String>) {
fn set_default(options: &mut HashMap<String, String>, aliases: &[&str], value: &str) {
if aliases
.iter()
.any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)))
{
return;
}
options.insert(aliases[0].to_owned(), value.to_owned());
}
set_default(options, &["pool_idle_timeout"], "300 seconds");
set_default(options, &["connect_timeout"], "10 seconds");
let has_custom_endpoint = ["aws_endpoint", "endpoint"]
.iter()
.any(|alias| options.keys().any(|k| k.eq_ignore_ascii_case(alias)));
if has_custom_endpoint {
set_default(
options,
&["aws_unsigned_payload", "unsigned_payload"],
"true",
);
}
}
fn quoted_string(value: &str) -> String {
format!("'{}'", value.replace('\'', "''"))
}
fn like_contains(value: &str) -> String {
let escaped = value
.replace('\\', "\\\\")
.replace('%', "\\%")
.replace('_', "\\_")
.replace('\'', "''");
format!("'%{escaped}%'")
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[tokio::test]
async fn store_opens_via_namespace_and_scan_works() -> Result<()> {
let temp = TempDir::new()?;
let url = Url::from_directory_path(temp.path())
.map_err(|()| anyhow::anyhow!("temp path is not absolute"))?;
let handle = Handle::open(&url).await?;
let cases: [(Table, &[&str]); 3] = [
(Table::Sessions, &["id"]),
(Table::Messages, &["id"]),
(Table::Parts, &["id"]),
];
for (table, projection) in cases {
let scanner = handle
.scan(table, ScanOpts::project_only(projection))
.await?;
let batch = scanner.try_into_batch().await?;
assert_eq!(batch.num_rows(), 0, "fresh table should be empty");
}
Ok(())
}
}