Audit, Compact, and Optimize LanceDB Vector Tables with Stale Version Pruning

id: lancedb-vector-table-compaction-and-version-prune-guard
namespace: company.ai

description: |
  Audit LanceDB vector tables for fragment fragmentation, tombstoned rows, and stale manifest versions.
  Track fragment bloat across runs in Kestra KV, distinguish an urgent fragment explosion from routine growth,
  pause for operator approval, execute atomic fragment compaction, version pruning, and vector index optimization,
  and verify post-compaction storage and fragment reduction.

inputs:
  - id: db_uri
    type: STRING
    displayName: LanceDB Database URI
    description: Local directory path (e.g. /data/lancedb) or remote cloud object
      store URI (e.g. s3://company-rag-storage/lancedb) hosting the LanceDB
      database.
    defaults: "/data/lancedb"

  - id: table_name
    type: STRING
    displayName: Vector Table Name
    description: Name of the LanceDB vector table to audit, compact, and optimize.
    defaults: "enterprise_knowledge_base"

  - id: fragment_count_threshold
    type: INT
    displayName: Fragment Count Threshold
    description: Minimum fragment count on a table before it is tracked as a
      compaction candidate.
    defaults: 20

  - id: critical_fragment_threshold
    type: INT
    displayName: Critical Fragment Threshold
    description: Fragment count at or above which a table is escalated immediately
      for compaction, bypassing consecutive-run confirmation.
    defaults: 100

  - id: tombstone_percent_threshold
    type: INT
    displayName: Tombstone Percentage Threshold
    description: Minimum percentage of deleted or superseded tombstoned rows (e.g.
      15%) that flags a table for compaction even if fragment count is moderate.
    defaults: 15

  - id: stale_version_retention_hours
    type: INT
    displayName: Stale Version Retention (Hours)
    description: Number of hours after which obsolete manifest versions and
      unreferenced data files are permanently pruned.
    defaults: 24

  - id: target_rows_per_fragment
    type: INT
    displayName: Target Rows Per Fragment
    description: Target row count per compacted Lance data fragment to balance
      parallel search scan efficiency and file size.
    defaults: 100000

  - id: min_consecutive_runs
    type: INT
    displayName: Minimum Consecutive Bloated Runs
    description: Number of consecutive audit passes fragment bloat must persist
      before it is treated as genuinely stalled rather than a transient batch
      ingestion.
    defaults: 2

  - id: cursor_key
    type: STRING
    displayName: KV Cursor Key
    description: Kestra KV key used to persist per-table fragmentation history and
      trend state across runs.
    defaults: "lancedb_compaction_tracker"

  - id: dry_run
    type: BOOL
    displayName: Dry Run Mode
    description: When true, audits fragments and versions and logs the proposed
      compaction and pruning plan without mutating any files.
    defaults: true

  - id: auto_remediate
    type: BOOL
    displayName: Auto Remediate
    description: When true and dry_run is false, pauses for operator approval and
      executes compaction, cleanup_old_versions, and index optimization.
    defaults: false

triggers:
  - id: periodic_compaction_audit
    type: io.kestra.plugin.core.trigger.Schedule
    description: Audit LanceDB vector table fragmentation every 6 hours. Shipped
      disabled so the flow can be verified manually in safe mode first.
    cron: "0 */6 * * *"
    disabled: true

  - id: on_demand_compaction_webhook
    type: io.kestra.plugin.core.trigger.Webhook
    description: Authenticated HTTP endpoint invoked by RAG ingestion pipelines, CDC
      consumers, or data maintenance runbooks.
    key: "{{ secret('LANCEDB_GUARD_WEBHOOK_KEY') }}"

tasks:
  - id: validate_inputs
    type: io.kestra.plugin.core.execution.Assert
    description: Assert operational thresholds and table parameters satisfy sanity
      constraints before inspecting the database.
    conditions:
      - "{{ inputs.table_name | length > 0 }}"
      - "{{ inputs.db_uri | length > 0 }}"
      - "{{ inputs.fragment_count_threshold >= 1 }}"
      - "{{ inputs.critical_fragment_threshold >=
        inputs.fragment_count_threshold }}"
      - "{{ inputs.tombstone_percent_threshold >= 1 and
        inputs.tombstone_percent_threshold <= 90 }}"
      - "{{ inputs.stale_version_retention_hours >= 1 }}"
      - "{{ inputs.target_rows_per_fragment >= 1000 }}"
      - "{{ inputs.min_consecutive_runs >= 1 }}"
    errorMessage: "Invalid configuration: critical_fragment_threshold must be >=
      fragment_count_threshold, tombstone_percent_threshold between 1 and 90,
      and all parameters positive."

  - id: audit_table_fragments
    type: io.kestra.plugin.scripts.python.Script
    description: Inspect LanceDB dataset fragments, measure tombstone row ratios and
      manifest version count, and diff against KV state.
    containerImage: python:3.11-slim
    beforeCommands:
      - pip install lancedb pyarrow
    env:
      DB_URI: "{{ inputs.db_uri }}"
      TABLE_NAME: "{{ inputs.table_name }}"
      FRAGMENT_THRESHOLD: "{{ inputs.fragment_count_threshold }}"
      CRITICAL_FRAGMENT_THRESHOLD: "{{ inputs.critical_fragment_threshold }}"
      TOMBSTONE_THRESHOLD: "{{ inputs.tombstone_percent_threshold }}"
      RETENTION_HOURS: "{{ inputs.stale_version_retention_hours }}"
      MIN_CONSECUTIVE_RUNS: "{{ inputs.min_consecutive_runs }}"
      EXECUTION_ID: "{{ execution.id }}"
      PREVIOUS_STATE_JSON: "{{ kv(inputs.cursor_key, errorOnMissing=false) ?? '{}' }}"
    script: |
      import json
      import os
      from datetime import datetime, timezone

      try:
          from kestra import Kestra
      except ImportError:
          Kestra = None

      import lancedb

      now_iso = datetime.now(timezone.utc).isoformat()
      db_uri = os.environ.get("DB_URI", "/data/lancedb")
      table_name = os.environ.get("TABLE_NAME", "enterprise_knowledge_base")
      fragment_threshold = int(os.environ.get("FRAGMENT_THRESHOLD", "20"))
      critical_fragment_threshold = int(os.environ.get("CRITICAL_FRAGMENT_THRESHOLD", "100"))
      tombstone_threshold = int(os.environ.get("TOMBSTONE_THRESHOLD", "15"))
      min_consecutive_runs = int(os.environ.get("MIN_CONSECUTIVE_RUNS", "2"))
      execution_id = os.environ.get("EXECUTION_ID", "manual")

      prev_state_raw = os.environ.get("PREVIOUS_STATE_JSON", "{}")
      try:
          prev_state = json.loads(prev_state_raw) if prev_state_raw.strip() else {}
      except Exception:
          prev_state = {}

      db = lancedb.connect(db_uri)
      existing_tables = db.table_names()

      if table_name not in existing_tables:
          raise ValueError(f"Table '{table_name}' does not exist in LanceDB at '{db_uri}'. Existing tables: {existing_tables}")

      table = db.open_table(table_name)
      ds = table.to_lance()

      fragments = ds.fragments
      fragment_count = len(fragments)
      active_rows = table.count_rows()

      total_fragment_rows = sum(f.count_rows() for f in fragments) if fragments else active_rows
      tombstone_count = max(0, total_fragment_rows - active_rows)
      tombstone_pct = round((tombstone_count / total_fragment_rows * 100.0), 2) if total_fragment_rows > 0 else 0.0

      versions = table.list_versions()
      version_count = len(versions)
      current_version = table.version

      table_key = f"{db_uri}/{table_name}"
      prev_entry = prev_state.get(table_key, {})
      prev_fragments = int(prev_entry.get("fragment_count", 0))
      prev_consecutive = int(prev_entry.get("consecutive_bloated_runs", 0))

      is_bloated = (fragment_count >= fragment_threshold) or (tombstone_pct >= tombstone_threshold)

      if is_bloated and (fragment_count >= prev_fragments or tombstone_pct >= float(prev_entry.get("tombstone_pct", 0.0))):
          consecutive_runs = prev_consecutive + 1
      elif is_bloated:
          consecutive_runs = max(1, prev_consecutive)
      else:
          consecutive_runs = 0

      new_state_snapshot = dict(prev_state)
      new_state_snapshot[table_key] = {
          "fragment_count": fragment_count,
          "active_rows": active_rows,
          "tombstone_count": tombstone_count,
          "tombstone_pct": tombstone_pct,
          "version_count": version_count,
          "current_version": current_version,
          "consecutive_bloated_runs": consecutive_runs,
          "last_evaluated_at": now_iso,
      }

      is_critical = fragment_count >= critical_fragment_threshold
      is_stuck_bloat = is_bloated and consecutive_runs >= min_consecutive_runs
      high_tombstones = tombstone_pct >= tombstone_threshold

      needs_compaction = is_critical or is_stuck_bloat or high_tombstones

      reasons = []
      if is_critical:
          reasons.append(f"CRITICAL_FRAGMENT_EXPLOSION (fragments={fragment_count} >= {critical_fragment_threshold})")
      if is_stuck_bloat:
          reasons.append(f"PERSISTENT_FRAGMENT_BLOAT (consecutive_runs={consecutive_runs} >= {min_consecutive_runs})")
      if high_tombstones:
          reasons.append(f"HIGH_TOMBSTONE_RATIO (tombstone_pct={tombstone_pct}% >= {tombstone_threshold}%)")

      escalation_reason = "; ".join(reasons) if reasons else "NOMINAL"
      overall_status = "CRITICAL" if is_critical else ("WARNING" if needs_compaction else "HEALTHY")

      summary_text = (
          f"LanceDB Guard: {overall_status} | Table '{table_name}' has {fragment_count} fragment(s), "
          f"{tombstone_count} tombstone row(s) ({tombstone_pct}%), {version_count} manifest version(s). "
          f"Reason: {escalation_reason}."
      )
      print(summary_text)

      audit_manifest = {
          "overall_status": overall_status,
          "table_name": table_name,
          "db_uri": db_uri,
          "fragment_count": fragment_count,
          "active_rows": active_rows,
          "total_fragment_rows": total_fragment_rows,
          "tombstone_count": tombstone_count,
          "tombstone_pct": tombstone_pct,
          "version_count": version_count,
          "current_version": current_version,
          "consecutive_runs": consecutive_runs,
          "escalation_reason": escalation_reason,
          "needs_compaction": needs_compaction,
          "evaluated_at": now_iso,
          "execution_id": execution_id,
      }

      if Kestra:
          Kestra.outputs({
              "needs_compaction": needs_compaction,
              "overall_status": overall_status,
              "fragment_count": fragment_count,
              "active_rows": active_rows,
              "tombstone_count": tombstone_count,
              "tombstone_pct": tombstone_pct,
              "version_count": version_count,
              "escalation_reason": escalation_reason,
              "summary": summary_text,
              "audit_manifest": audit_manifest,
              "new_state_snapshot": new_state_snapshot,
          })

  - id: update_compaction_tracker_kv
    type: io.kestra.plugin.core.kv.Set
    description: Persist the latest LanceDB table fragmentation metrics and
      consecutive run history into Kestra KV.
    key: "{{ inputs.cursor_key }}"
    value: "{{ outputs.audit_table_fragments.vars.new_state_snapshot | toJson }}"
    kvType: JSON

  - id: check_compaction_needed
    type: io.kestra.plugin.core.flow.If
    description: Evaluate whether the table exhibits excessive fragments or
      tombstone dead rows requiring compaction.
    condition: "{{ outputs.audit_table_fragments.vars.needs_compaction == true }}"
    then:
      - id: check_execution_mode
        type: io.kestra.plugin.core.flow.If
        description: Check whether active auto-remediation is authorized or whether safe
          dry run logging should take precedence.
        condition: "{{ not inputs.dry_run and inputs.auto_remediate }}"
        then:
          - id: pause_for_operator_approval
            type: io.kestra.plugin.core.flow.Pause
            description: Pause workflow execution for human operator review before
              initiating atomic file compaction and version deletion.
            onResume:
              - id: approved
                type: BOOL
                defaults: false
                description: Set true to confirm table compaction, obsolete version cleanup, and
                  index optimization on the LanceDB dataset.

          - id: branch_on_approval
            type: io.kestra.plugin.core.flow.If
            description: Execute remediation only if explicit operator approval was provided
              during the pause gate.
            condition: "{{ outputs.pause_for_operator_approval.onResume.approved == true }}"
            then:
              - id: execute_compaction_and_cleanup
                type: io.kestra.plugin.scripts.python.Script
                description: Run compact_files, purge unreferenced tombstoned versions via
                  cleanup_old_versions, and optimize vector indices.
                containerImage: python:3.11-slim
                beforeCommands:
                  - pip install lancedb pyarrow
                env:
                  DB_URI: "{{ inputs.db_uri }}"
                  TABLE_NAME: "{{ inputs.table_name }}"
                  TARGET_ROWS: "{{ inputs.target_rows_per_fragment }}"
                  RETENTION_HOURS: "{{ inputs.stale_version_retention_hours }}"
                script: |
                  import os
                  from datetime import timedelta

                  try:
                      from kestra import Kestra
                  except ImportError:
                      Kestra = None

                  import lancedb

                  db_uri = os.environ.get("DB_URI", "/data/lancedb")
                  table_name = os.environ.get("TABLE_NAME", "enterprise_knowledge_base")
                  target_rows = int(os.environ.get("TARGET_ROWS", "100000"))
                  retention_hours = int(os.environ.get("RETENTION_HOURS", "24"))

                  db = lancedb.connect(db_uri)
                  table = db.open_table(table_name)

                  print(f"Starting atomic file compaction on table '{table_name}' with target {target_rows} rows/fragment...")
                  compaction_stats = table.compact_files(
                      target_rows_per_fragment=target_rows,
                      materialize_deletions=True,
                  )
                  print(f"Compaction completed: {compaction_stats}")

                  print(f"Pruning manifest versions older than {retention_hours} hour(s)...")
                  cleanup_stats = table.cleanup_old_versions(
                      older_than=timedelta(hours=retention_hours),
                      delete_unreferenced=True,
                  )
                  print(f"Cleanup completed: {cleanup_stats}")

                  print(f"Optimizing indices and file layout on table '{table_name}'...")
                  try:
                      table.optimize()
                      print("Table optimization finished successfully.")
                  except Exception as e:
                      print(f"Optimize warning (non-fatal): {e}")

                  if Kestra:
                      Kestra.outputs({
                          "compaction_stats": str(compaction_stats),
                          "cleanup_stats": str(cleanup_stats),
                          "remediation_status": "COMPLETED",
                      })

              - id: verify_compaction_health
                type: io.kestra.plugin.scripts.python.Script
                description: Re-inspect dataset fragments, verify zero tombstoned rows, and
                  benchmark canary vector search latency.
                containerImage: python:3.11-slim
                beforeCommands:
                  - pip install lancedb pyarrow
                env:
                  DB_URI: "{{ inputs.db_uri }}"
                  TABLE_NAME: "{{ inputs.table_name }}"
                  PRE_FRAGMENTS: "{{ outputs.audit_table_fragments.vars.fragment_count }}"
                  PRE_TOMBSTONES: "{{ outputs.audit_table_fragments.vars.tombstone_count }}"
                script: |
                  import os
                  import time
                  from datetime import datetime, timezone

                  try:
                      from kestra import Kestra
                  except ImportError:
                      Kestra = None

                  import lancedb

                  db_uri = os.environ.get("DB_URI", "/data/lancedb")
                  table_name = os.environ.get("TABLE_NAME", "enterprise_knowledge_base")
                  pre_fragments = int(os.environ.get("PRE_FRAGMENTS", "0"))
                  pre_tombstones = int(os.environ.get("PRE_TOMBSTONES", "0"))

                  db = lancedb.connect(db_uri)
                  table = db.open_table(table_name)
                  ds = table.to_lance()

                  post_fragments = len(ds.fragments)
                  active_rows = table.count_rows()
                  total_rows = sum(f.count_rows() for f in ds.fragments) if ds.fragments else active_rows
                  post_tombstones = max(0, total_rows - active_rows)
                  post_version = table.version
                  post_version_count = len(table.list_versions())

                  # Run canary vector search to verify table queryability
                  canary_passed = True
                  search_latency_ms = 0.0
                  try:
                      schema = table.schema
                      vector_col = None
                      for field in schema:
                          if "vector" in field.name.lower() or hasattr(field.type, "list_size"):
                              vector_col = field.name
                              break

                      t0 = time.perf_counter()
                      if vector_col:
                          dim = getattr(schema.field(vector_col).type, "list_size", 128) or 128
                          sample_vector = [0.0] * dim
                          table.search(sample_vector).limit(5).to_arrow()
                      else:
                          table.search().limit(5).to_arrow()
                      search_latency_ms = round((time.perf_counter() - t0) * 1000.0, 2)
                  except Exception as e:
                      print(f"Canary search warning: {e}")
                      canary_passed = False

                  fragments_reduced = (post_fragments < pre_fragments) or (pre_fragments <= 1 and post_fragments <= 1)
                  tombstones_cleared = post_tombstones <= pre_tombstones

                  verification_passed = fragments_reduced and tombstones_cleared and canary_passed

                  audit_record = {
                      "table_name": table_name,
                      "pre_fragments": pre_fragments,
                      "post_fragments": post_fragments,
                      "fragments_reduced": fragments_reduced,
                      "pre_tombstones": pre_tombstones,
                      "post_tombstones": post_tombstones,
                      "tombstones_cleared": tombstones_cleared,
                      "active_rows": active_rows,
                      "post_version": post_version,
                      "post_version_count": post_version_count,
                      "canary_search_passed": canary_passed,
                      "canary_latency_ms": search_latency_ms,
                      "verification_passed": verification_passed,
                      "verified_at": datetime.now(timezone.utc).isoformat(),
                  }

                  print(f"Compaction Verification: {'PASSED' if verification_passed else 'FAILED'}")
                  print(f"Fragments: {pre_fragments} -> {post_fragments} | Tombstones: {pre_tombstones} -> {post_tombstones} | Canary: {search_latency_ms}ms")

                  if Kestra:
                      Kestra.outputs({
                          "verification_passed": verification_passed,
                          "post_fragment_count": post_fragments,
                          "post_tombstone_count": post_tombstones,
                          "canary_latency_ms": search_latency_ms,
                          "audit_record": audit_record,
                      })

              - id: assert_compaction_success
                type: io.kestra.plugin.core.flow.If
                description: Abort with an error if post-compaction verification failed to
                  confirm fragment reduction or table queryability.
                condition: "{{ outputs.verify_compaction_health.vars.verification_passed != true
                  }}"
                then:
                  - id: fail_unresolved_bloat
                    type: io.kestra.plugin.core.execution.Fail
                    description: Abort execution with diagnostic details when table compaction did
                      not achieve expected file or tombstone reduction.
                    errorMessage: "LanceDB table compaction completed but verification failed.
                      Fragments were not reduced (was {{
                      outputs.audit_table_fragments.vars.fragment_count }}, now
                      {{
                      outputs.verify_compaction_health.vars.post_fragment_count
                      }}) or tombstones were not purged. Check disk space and
                      concurrent writers."
                else:
                  - id: log_compaction_success
                    type: io.kestra.plugin.core.log.Log
                    description: Record successful table compaction, fragment consolidation, and
                      latency benchmark.
                    message: "LanceDB table '{{ inputs.table_name }}' successfully compacted.
                      Fragments: {{
                      outputs.audit_table_fragments.vars.fragment_count }} -> {{
                      outputs.verify_compaction_health.vars.post_fragment_count
                      }}. Canary search latency: {{
                      outputs.verify_compaction_health.vars.canary_latency_ms
                      }}ms."

              - id: save_recovery_audit_kv
                type: io.kestra.plugin.core.kv.Set
                description: Persist the remediation audit record into Kestra KV for compliance
                  and operational tracking.
                key: "lancedb_compaction_audit_{{ execution.id }}"
                value: "{{ outputs.verify_compaction_health.vars.audit_record | toJson }}"
                kvType: JSON

              - id: notify_resolution_slack
                type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
                description: Dispatch resolution alert to Slack confirming successful table
                  compaction, version pruning, and vector index optimization.
                url: "{{ secret('SLACK_WEBHOOK_URL') }}"
                payload: |
                  {
                    "attachments": [
                      {
                        "color": "#2EB67D",
                        "title": ":white_check_mark: LanceDB Vector Table Compaction Completed",
                        "text": "*Database:* `{{ inputs.db_uri }}`\n*Table:* `{{ inputs.table_name }}`\n*Fragments Consolidated:* `{{ outputs.audit_table_fragments.vars.fragment_count }}` -> `{{ outputs.verify_compaction_health.vars.post_fragment_count }}`\n*Canary Search Latency:* `{{ outputs.verify_compaction_health.vars.canary_latency_ms }} ms`\n*Execution:* `{{ execution.id }}`"
                      }
                    ]
                  }
            else:
              - id: log_approval_declined
                type: io.kestra.plugin.core.log.Log
                description: Record that the human operator declined approval for compacting the
                  LanceDB table.
                message: "Operator declined approval for LanceDB table compaction. Table '{{
                  inputs.table_name }}' files left untouched."
        else:
          - id: log_dry_run_plan
            type: io.kestra.plugin.core.log.Log
            description: Log the proposed compaction and pruning actions without mutating
              any files.
            message: |
              Dry run active or auto_remediate disabled (dry_run={{ inputs.dry_run }}, auto_remediate={{ inputs.auto_remediate }}).
              Table '{{ inputs.table_name }}' flagged for compaction:
              - Fragments: {{ outputs.audit_table_fragments.vars.fragment_count }} (threshold={{ inputs.fragment_count_threshold }})
              - Tombstones: {{ outputs.audit_table_fragments.vars.tombstone_count }} ({{ outputs.audit_table_fragments.vars.tombstone_pct }}%)
              - Versions: {{ outputs.audit_table_fragments.vars.version_count }}
              - Reason: {{ outputs.audit_table_fragments.vars.escalation_reason }}
    else:
      - id: log_healthy_state
        type: io.kestra.plugin.core.log.Log
        description: Confirm nominal table state when fragments and tombstone ratios are
          within healthy limits.
        message: "LanceDB table '{{ inputs.table_name }}' is healthy at {{ inputs.db_uri
          }}. Fragment count ({{
          outputs.audit_table_fragments.vars.fragment_count }}) and tombstones
          ({{ outputs.audit_table_fragments.vars.tombstone_pct }}%) within
          thresholds."

errors:
  - id: alert_on_workflow_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Dispatch emergency alert to Slack if the guard encounters an
      unhandled execution error.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "attachments": [
          {
            "color": "danger",
            "title": ":rotating_light: LanceDB Compaction Guard Execution Failed",
            "text": "Failed execution `{{ execution.id }}` while auditing or compacting table `{{ inputs.table_name }}` at `{{ inputs.db_uri }}`. Check storage permissions, worker logs, and table availability."
          }
        ]
      }

outputs:
  - id: table_name
    type: STRING
    value: "{{ inputs.table_name }}"

  - id: initial_fragments
    type: INT
    value: "{{ outputs.audit_table_fragments.vars.fragment_count ?? 0 }}"

  - id: post_fragments
    type: INT
    value: "{{ outputs.verify_compaction_health.vars.post_fragment_count ??
      outputs.audit_table_fragments.vars.fragment_count ?? 0 }}"

  - id: status
    type: STRING
    value: "{{ (outputs.audit_table_fragments.vars.needs_compaction ?? false) ?
      ((inputs.auto_remediate and not inputs.dry_run and
      (outputs.pause_for_operator_approval.onResume.approved ?? false)) ?
      'COMPACTED' : 'COMPACTION_REQUIRED') : 'HEALTHY' }}"