seaweedfs/seaweed-worker/crates/lance/tests/compaction.rs

//! Drives the compaction handler against a live namespace.
//!
//! Skipped unless WEED_LANCE_NAMESPACE names one, the way the Go integration
//! tests skip without Docker: compaction rewrites real files, and there is
//! nothing to learn from it against a fake.

use anyhow::Result;
use seaweed_worker_core::pb::{
    ExecuteJobRequest, JobSpec, RunDetectionRequest, config_value::Kind,
};
use seaweed_worker_core::{JobHandler, PreviewProvider};
use weed_lance_worker::jobs::cleanup::CleanupVersionsHandler;
use weed_lance_worker::jobs::compact::{CompactHandler, JOB_TYPE};
use weed_lance_worker::jobs::indices::OptimizeIndicesHandler;
use weed_lance_worker::preview::LancePreview;

mod common;
use common::{Recorder, client, int_config, namespace_url, signing_key, storage};

/// These tests drive one live gateway and one shared catalog: `list_all_tables`
/// sweeps everything, so a table another test is writing shows up in this test's
/// detection. Rust runs a binary's tests concurrently, so take a lock.
static GATEWAY: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());

/// Declares a table through the namespace and writes `fragments` one-row
/// appends into it, so a test brings its own state instead of depending on
/// whatever a previous run left behind.
async fn seed_fragmented_table(url: &str, name: &str, fragments: usize) -> Result<String> {
    seed_table(url, name, fragments, 1, false).await
}

/// Writes `batches` appends of `rows_each` into a freshly declared table, and
/// optionally builds a vector index after the first batch so the later ones are
/// rows no index covers.
async fn seed_table(
    url: &str,
    name: &str,
    batches: usize,
    rows_each: usize,
    with_index: bool,
) -> Result<String> {
    use arrow_array::{
        FixedSizeListArray, Float32Array, Int64Array, RecordBatch, RecordBatchIterator,
    };
    use arrow_schema::{DataType, Field, Schema};
    use lance::dataset::{Dataset, WriteMode, WriteParams};
    use std::sync::Arc;

    // Declaring is the namespace's job, not the worker's, so the test asks for
    // it directly rather than widening the client the worker uses. The bucket and
    // namespace come first: a table cannot be declared under a parent that does
    // not exist, and a test that assumes one is a test that only passes twice.
    let http = reqwest::Client::new();
    for parent in ["vec", "vec$ml"] {
        http.post(format!("{url}/v1/namespace/{parent}/create"))
            .json(&serde_json::json!({"mode": "EXIST_OK"}))
            .send()
            .await?
            .error_for_status()?;
    }
    let encoded = format!("vec$ml${name}");
    http.post(format!("{url}/v1/table/{encoded}/declare"))
        .json(&serde_json::json!({}))
        .send()
        .await?
        .error_for_status()?;

    let client = client(url);
    let id = vec!["vec".to_string(), "ml".to_string(), name.to_string()];
    let description = client.describe_table(&id).await?;

    let wrapper = storage().await.wrapper(&description.location).await?;
    const DIM: i32 = 16;
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new(
            "vec",
            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM),
            false,
        ),
    ]));

    for i in 0..batches {
        let ids: Vec<i64> = (0..rows_each).map(|r| (i * rows_each + r) as i64).collect();
        let values: Vec<f32> = ids
            .iter()
            .flat_map(|id| (0..DIM).map(move |d| (*id as f32) + d as f32))
            .collect();
        let vectors = FixedSizeListArray::new(
            Arc::new(Field::new("item", DataType::Float32, true)),
            DIM,
            Arc::new(Float32Array::from(values)),
            None,
        );
        let batch = RecordBatch::try_new(
            schema.clone(),
            vec![Arc::new(Int64Array::from(ids)), Arc::new(vectors)],
        )?;
        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
        let params = WriteParams {
            mode: if i == 0 {
                WriteMode::Overwrite
            } else {
                WriteMode::Append
            },
            store_params: Some(weed_lance_worker::dataset::store_params(wrapper.clone())),
            ..Default::default()
        };
        let dataset = Dataset::write(reader, description.location.as_str(), Some(params)).await?;

        // The index is built after the first batch, so everything appended
        // afterwards is a row it does not cover.
        if with_index && i == 0 {
            use lance::index::DatasetIndexExt;
            use lance::index::vector::VectorIndexParams;
            use lance_index::IndexType;
            use lance_index::vector::{ivf::IvfBuildParams, pq::PQBuildParams};

            let mut dataset = dataset;
            let params = VectorIndexParams::with_ivf_pq_params(
                lance_linalg::distance::MetricType::L2,
                IvfBuildParams::new(1),
                PQBuildParams::new(4, 8),
            );
            dataset
                .create_index(&["vec"], IndexType::Vector, None, &params, true)
                .await?;
        }
    }
    Ok(encoded)
}

/// A table with more fragments than the policy allows is proposed, and running
/// the proposal leaves it with fewer than it started with.
#[tokio::test]
async fn compacts_a_fragmented_table() {
    let _gateway = GATEWAY.lock().await;
    let Some(url) = namespace_url() else {
        eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
        return;
    };
    // Seeded here rather than by a script, so the test is repeatable: a previous
    // run compacts the table it depended on.
    let encoded = seed_fragmented_table(&url, "compactme", 12)
        .await
        .expect("seed a fragmented table");
    let handler = CompactHandler::new(url)
        .with_storage(storage().await)
        .with_signing_key(signing_key());
    let recorder = Recorder::default();

    let request = RunDetectionRequest {
        request_id: "detect-1".to_string(),
        job_type: JOB_TYPE.to_string(),
        worker_config_values: int_config("min_fragments", 4),
        ..Default::default()
    };
    handler
        .detect(&request, &recorder)
        .await
        .expect("detection failed");

    let proposals = recorder.proposals.lock().unwrap().clone();
    assert!(
        !proposals.is_empty(),
        "expected a proposal for the fragmented table"
    );
    let proposal = proposals
        .iter()
        .find(|p| p.summary.contains(encoded.as_str()))
        .expect("no proposal for the seeded table");

    // Detection opened the dataset to decide, so it reports what it saw. This is
    // the only description of a Lance table anything outside the format can give.
    let observations = recorder.observations.lock().unwrap().clone();
    let observed = observations
        .iter()
        .find(|o| o.object_id.last().map(String::as_str) == Some("compactme"))
        .expect("detection reported no observation for the seeded table");
    assert_eq!(observed.format, "LANCE");
    for attribute in ["fragments", "rows", "versions", "schema"] {
        assert!(
            observed.attributes.contains_key(attribute),
            "observation is missing {attribute}: {:?}",
            observed.attributes.keys().collect::<Vec<_>>()
        );
    }

    let execute = ExecuteJobRequest {
        request_id: "execute-1".to_string(),
        job: Some(JobSpec {
            job_id: "job-1".to_string(),
            job_type: JOB_TYPE.to_string(),
            parameters: proposal.parameters.clone(),
            ..Default::default()
        }),
        worker_config_values: int_config("target_rows_per_fragment", 1_048_576),
        ..Default::default()
    };
    handler
        .execute(&execute, &recorder)
        .await
        .expect("execution failed");

    let completed = recorder.completed.lock().unwrap().clone();
    let result = completed.first().expect("no completion reported");
    assert!(
        result.success,
        "compaction reported failure: {}",
        result.error_message
    );
    let summary = result
        .result
        .as_ref()
        .map(|r| r.summary.clone())
        .unwrap_or_default();
    assert!(
        summary.contains("became"),
        "completion carried no fragment counts: {summary}"
    );
    eprintln!("compaction result: {summary}");
}

/// A table with more versions than the floor is proposed, and running the job
/// reports what it removed. The compaction test above leaves one behind.
#[tokio::test]
async fn cleans_up_old_versions() {
    let _gateway = GATEWAY.lock().await;
    let Some(url) = namespace_url() else {
        eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
        return;
    };
    let encoded = seed_table(&url, "cleanme", 6, 4, false)
        .await
        .expect("seed a table with versions to clean");

    let handler = CleanupVersionsHandler::new(url)
        .with_storage(storage().await)
        .with_signing_key(signing_key());
    let recorder = Recorder::default();

    let request = RunDetectionRequest {
        request_id: "detect-cleanup".to_string(),
        job_type: "lance_cleanup_versions".to_string(),
        worker_config_values: int_config("min_versions_to_keep", 2),
        ..Default::default()
    };
    handler
        .detect(&request, &recorder)
        .await
        .expect("detection failed");
    let proposals = recorder.proposals.lock().unwrap().clone();
    let proposal = proposals
        .iter()
        .find(|p| p.summary.contains(encoded.as_str()))
        .cloned()
        .expect("no cleanup proposal for the seeded table");

    // Retain nothing, so every version outside the current one is fair game and
    // the job has something to report rather than a no-op.
    let execute = ExecuteJobRequest {
        request_id: "execute-cleanup".to_string(),
        job: Some(JobSpec {
            job_id: "job-cleanup".to_string(),
            job_type: "lance_cleanup_versions".to_string(),
            parameters: proposal.parameters.clone(),
            ..Default::default()
        }),
        worker_config_values: int_config("retain_hours", 0),
        ..Default::default()
    };
    handler
        .execute(&execute, &recorder)
        .await
        .expect("cleanup failed");

    let completed = recorder.completed.lock().unwrap().clone();
    let result = completed.last().expect("no completion reported");
    assert!(
        result.success,
        "cleanup reported failure: {}",
        result.error_message
    );
    eprintln!(
        "cleanup result: {}",
        result.result.as_ref().unwrap().summary
    );
}

/// A table with no indices has nothing to optimize, so detection proposes
/// nothing rather than queueing work that would do nothing.
#[tokio::test]
async fn skips_tables_without_indices() {
    let _gateway = GATEWAY.lock().await;
    let Some(url) = namespace_url() else {
        eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
        return;
    };
    let encoded = seed_table(&url, "noindex", 2, 8, false)
        .await
        .expect("seed a table without an index");

    let handler = OptimizeIndicesHandler::new(url)
        .with_storage(storage().await)
        .with_signing_key(signing_key());
    let recorder = Recorder::default();

    let request = RunDetectionRequest {
        request_id: "detect-indices".to_string(),
        job_type: "lance_optimize_indices".to_string(),
        worker_config_values: int_config("max_unindexed_rows", 1),
        ..Default::default()
    };
    handler
        .detect(&request, &recorder)
        .await
        .expect("detection failed");
    // Judge this table only: the catalog holds every other test's tables too,
    // and an indexed one with uncovered rows is supposed to be proposed.
    assert!(
        !recorder
            .proposals
            .lock()
            .unwrap()
            .iter()
            .any(|p| p.summary.contains(encoded.as_str())),
        "a table with no indices must not be proposed for reindexing"
    );
}

/// The job with no Iceberg equivalent: rows appended after an index was built
/// are invisible to a search of it until this runs. Needs a table with an index
/// and rows outside it, which `indexed.py` in the scratchpad seeds.
#[tokio::test]
async fn reindexes_rows_an_index_does_not_cover() {
    let _gateway = GATEWAY.lock().await;
    let Some(url) = namespace_url() else {
        eprintln!("WEED_LANCE_NAMESPACE is unset, skipping");
        return;
    };
    let encoded = seed_table(&url, "reindexme", 2, 512, true)
        .await
        .expect("seed an indexed table with uncovered rows");
    let handler = OptimizeIndicesHandler::new(url)
        .with_storage(storage().await)
        .with_signing_key(signing_key());
    let recorder = Recorder::default();

    let request = RunDetectionRequest {
        request_id: "detect-reindex".to_string(),
        job_type: "lance_optimize_indices".to_string(),
        worker_config_values: int_config("max_unindexed_rows", 100),
        ..Default::default()
    };
    handler
        .detect(&request, &recorder)
        .await
        .expect("detection failed");
    let proposals = recorder.proposals.lock().unwrap().clone();
    let proposal = proposals
        .iter()
        .find(|p| p.summary.contains(encoded.as_str()))
        .cloned()
        .expect("the seeded indexed table was not proposed");

    let execute = ExecuteJobRequest {
        request_id: "execute-reindex".to_string(),
        job: Some(JobSpec {
            job_id: "job-reindex".to_string(),
            job_type: "lance_optimize_indices".to_string(),
            parameters: proposal.parameters.clone(),
            ..Default::default()
        }),
        ..Default::default()
    };
    handler
        .execute(&execute, &recorder)
        .await
        .expect("reindex failed");

    let completed = recorder.completed.lock().unwrap().clone();
    let result = completed.last().expect("no completion reported");
    assert!(
        result.success,
        "reindex reported failure: {}",
        result.error_message
    );
    let output = &result.result.as_ref().unwrap().output_values;
    let after = match output
        .get("unindexed_rows_after")
        .and_then(|v| v.kind.as_ref())
    {
        Some(Kind::Int64Value(value)) => *value,
        other => panic!("no unindexed_rows_after in {other:?}"),
    };
    assert_eq!(
        after, 0,
        "rows are still outside the index after optimizing"
    );
    eprintln!(
        "reindex result: {}",
        result.result.as_ref().unwrap().summary
    );
}

/// The UI's whole reason for asking a worker: admin cannot read a Lance table,
/// so the rows have to come back already rendered.
#[tokio::test]
async fn previews_rows_of_a_table() {
    let _gateway = GATEWAY.lock().await;
    let Some(url) = namespace_url() else {
        eprintln!("set WEED_LANCE_NAMESPACE to run this test");
        return;
    };
    seed_table(&url, "previewme", 2, 3, false)
        .await
        .expect("seed a table to preview");

    let provider = LancePreview::new(url, storage().await);
    let id = vec!["vec".to_string(), "ml".to_string(), "previewme".to_string()];
    let preview = provider.preview(&id, 4).await.expect("preview the table");

    assert_eq!(preview.columns, vec!["id".to_string(), "vec".to_string()]);
    assert_eq!(preview.total_rows, 6, "total is the table, not the sample");
    assert_eq!(preview.rows.len(), 4, "row_limit bounds the sample");
    assert!(
        preview.rows[0][1].starts_with('['),
        "a vector column should render as a list, got {:?}",
        preview.rows[0][1]
    );
}

/// The claim that removed managed versioning is not "one writer wins the
/// conditional PUT" - that is only the mechanism. It is that concurrent writers
/// lose nothing: the loser sees the conflict, rebases, and commits again. Eight
/// writers appending at once must leave all eight batches in the table.
#[tokio::test]
async fn concurrent_writers_keep_every_commit() {
    let _gateway = GATEWAY.lock().await;
    let Some(url) = namespace_url() else {
        eprintln!("set WEED_LANCE_NAMESPACE to run this test");
        return;
    };
    const WRITERS: i64 = 8;
    const ROWS_EACH: i64 = 4;

    seed_table(&url, "racers", 1, ROWS_EACH as usize, false)
        .await
        .expect("seed the table the writers will append to");

    let client = client(&url);
    let id = vec!["vec".to_string(), "ml".to_string(), "racers".to_string()];
    let description = client.describe_table(&id).await.expect("describe");
    let wrapper = storage()
        .await
        .wrapper(&description.location)
        .await
        .expect("storage for the table");

    let writes: Vec<_> = (0..WRITERS)
        .map(|writer| {
            let location = description.location.clone();
            let wrapper = wrapper.clone();
            tokio::spawn(
                async move { append_rows(&location, wrapper, writer * 1000, ROWS_EACH).await },
            )
        })
        .collect();
    for (writer, handle) in writes.into_iter().enumerate() {
        handle
            .await
            .expect("writer panicked")
            .unwrap_or_else(|err| panic!("writer {writer} failed to commit: {err:#}"));
    }

    let table = weed_lance_worker::dataset::open(&client, &id, &storage().await)
        .await
        .expect("reopen the table");
    let rows = table.dataset.count_rows(None).await.expect("count rows");
    let expected = (ROWS_EACH + WRITERS * ROWS_EACH) as usize;
    assert_eq!(
        rows, expected,
        "concurrent commits lost data: {rows} rows, want {expected}"
    );
}

/// Appends one batch to an existing dataset, the way an independent writer would.
async fn append_rows(
    location: &str,
    wrapper: Option<std::sync::Arc<dyn lance::io::WrappingObjectStore>>,
    first_id: i64,
    rows: i64,
) -> Result<()> {
    use arrow_array::{
        FixedSizeListArray, Float32Array, Int64Array, RecordBatch, RecordBatchIterator,
    };
    use arrow_schema::{DataType, Field, Schema};
    use lance::dataset::{Dataset, WriteMode, WriteParams};
    use std::sync::Arc;

    const DIM: i32 = 16;
    let schema = Arc::new(Schema::new(vec![
        Field::new("id", DataType::Int64, false),
        Field::new(
            "vec",
            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), DIM),
            false,
        ),
    ]));
    let ids: Vec<i64> = (0..rows).map(|r| first_id + r).collect();
    let values: Vec<f32> = ids
        .iter()
        .flat_map(|id| (0..DIM).map(move |d| (*id as f32) + d as f32))
        .collect();
    let vectors = FixedSizeListArray::new(
        Arc::new(Field::new("item", DataType::Float32, true)),
        DIM,
        Arc::new(Float32Array::from(values)),
        None,
    );
    let batch = RecordBatch::try_new(
        schema.clone(),
        vec![Arc::new(Int64Array::from(ids)), Arc::new(vectors)],
    )?;
    let params = WriteParams {
        mode: WriteMode::Append,
        store_params: Some(weed_lance_worker::dataset::store_params(wrapper)),
        ..Default::default()
    };
    Dataset::write(
        RecordBatchIterator::new(vec![Ok(batch)], schema.clone()),
        location,
        Some(params),
    )
    .await?;
    Ok(())
}