QuickwitDatafusion

[alexanderbianchi]

THIS IS A DRAFT, IT'S STILL VERY CLAUDE-Y AND I WILL NEED TO DO A DEEP REVIEW MYSELF BEFORE ASKING FOR ONE FROM OTHERS

Distributed DataFusion Execution for Quickwit

Adds a new quickwit-datafusion crate that enables distributed SQL query execution over Quickwit splits using DataFusion + datafusion-distributed + tantivy-datafusion.

Architecture

graph TD
    Client["Client (pyarrow / SQL / REST)"]
    
    Client --> Coordinator

    subgraph Coordinator["Coordinator (any searcher node)"]
        direction TB
        SB["QuickwitSessionBuilder"]
        SB --> CAT["QuickwitSchemaProvider<br/>(lazy catalog — resolves index names<br/>from metastore automatically)"]
        SB --> WR["QuickwitWorkerResolver<br/>(reads SearcherPool / Chitchat)"]
        SB --> TC["TantivyCodec (stateless)"]
        SB --> DPO["DistributedPhysicalOptimizerRule"]
        SB --> BSP["build_search_plan()"]
        BSP --> MS["metastore.list_splits()"]
        BSP --> JOIN["per-split: inv ⋈ f join"]
        BSP --> UNION["UNION ALL across splits<br/>(schema-aligned, NULL-filled)"]
        BSP --> SORT["sort / limit / agg on top"]
    end

    Coordinator -->|"Arrow Flight gRPC<br/>(same port as SearchService)"| W1
    Coordinator -->|"Arrow Flight gRPC"| W2
    Coordinator -->|"Arrow Flight gRPC"| W3

    subgraph W1["Searcher Node 1"]
        direction TB
        TC1["TantivyCodec → decode plan<br/>+ restore pushed filters"]
        OF1["OpenerFactory →<br/>StorageSplitOpener<br/>→ open_index_with_caches()<br/>+ TokenizerManager"]
        LN1["tantivy-df leaf nodes execute"]
        TC1 --> OF1 --> LN1
    end

    subgraph W2["Searcher Node 2"]
        direction TB
        TC2["decode + filters"] --> OF2["open split"] --> LN2["execute"]
    end

    subgraph W3["Searcher Node 3"]
        direction TB
        TC3["decode + filters"] --> OF3["open split"] --> LN3["execute"]
    end

    W1 -->|"RecordBatch stream"| Merge
    W2 -->|"RecordBatch stream"| Merge
    W3 -->|"RecordBatch stream"| Merge

    subgraph Merge["Coordinator merges"]
        SPM["SortPreservingMergeExec"]
        AGG["AggregateExec(mode=Final)"]
        SPM --> AGG
    end

    Merge --> Result["RecordBatches → Client"]

Loading

Per-split plan decomposition

graph TD
    subgraph "Single split (one UNION ALL child)"
        HJ["HashJoinExec<br/>on (_doc_id, _segment_ord)"]
        INV["InvertedIndexDataSource<br/>(query=true, full-text filter)"]
        FF["FastFieldDataSource<br/>(partitions=N, pushed_filters serialized)"]
        HJ --> INV
        HJ --> FF
    end

Loading

For queries without full-text search, the inverted index join is skipped — just a FastFieldDataSource scan with filter pushdown.

Splits are isolated by UNION ALL boundaries — (_doc_id, _segment_ord) can collide across splits, but never across union children. Schema differences across splits are handled by aligning to the canonical schema (from the doc mapper): missing columns → NULL, type mismatches → error with CAST suggestion.

How it works

Query path (SQL → RecordBatches):

1. Client sends SQL to any searcher node (via Arrow Flight do_get or REST)
2. QuickwitSchemaProvider lazily resolves index names from the metastore — no manual table registration
3. build_search_plan() queries metastore for splits, builds per-split join plans using DataFrame API, aligns schemas, unions across splits
4. DistributedPhysicalOptimizerRule decomposes into stages — leaf scans to workers, merge/sort/agg on coordinator
5. Workers receive plan fragments via Arrow Flight, TantivyCodec decodes nodes + pushed filters, OpenerFactory (from SessionConfig) creates StorageSplitOpener which calls open_index_with_caches() with tokenizer manager, executes tantivy-df leaf nodes, streams RecordBatches back
6. Coordinator merges via SortPreservingMergeExec / AggregateExec(Final)

Split openers

Opener	Usage	How it opens
SplitIndexOpener	Tests	DashMap lookup (in-memory)
StorageSplitOpener	Production	open_index_with_caches(): S3/GCS/local download, footer + fast field caches, tokenizer manager from doc mapper

Crate structure

tantivy-datafusion (separate repo):

• IndexOpener trait with identifier() for serialization
• TantivyCodec — stateless codec, serializes all 3 DataSource types + pushed filters via PhysicalExprNode
• OpenerFactoryExt — registers opener factory on SessionConfig

quickwit-datafusion (this PR):

File	Purpose
catalog.rs	QuickwitSchemaProvider — lazy catalog, resolves index names from metastore, builds doc mapper, creates StorageSplitOpener factory with tokenizer manager
split_opener.rs	SplitIndexOpener (test) + StorageSplitOpener (production: open_index_with_caches + tokenizer)
table_provider.rs	QuickwitTableProvider — queries metastore for splits at scan time
query_translator.rs	SearchRequest → DataFrame: QueryAst→Expr, per-split join composition, UNION ALL, schema alignment
resolver.rs	QuickwitWorkerResolver — reads SearcherPool (Chitchat)
session.rs	QuickwitSessionBuilder — wires codec + resolver + catalog + optimizer rule
flight.rs	QuickwitFlightService — handles df-distributed worker traffic AND external SQL queries on same port
worker.rs	Test helper for df-distributed's start_localhost_context

quickwit-serve changes:

• grpc.rs — Flight service on same tonic gRPC port as SearchService
• lib.rs — Flight + session builder startup alongside SearcherPool
• datafusion_api.rs — REST endpoint behind enable_datafusion_endpoint flag
• quickwit-search/leaf.rs — open_index_with_caches made public

Client-facing endpoints

Arrow Flight (same gRPC port, always on):

import pyarrow.flight as flight
client = flight.connect("grpc://quickwit-node:7281")
ticket = flight.Ticket(b"SELECT id, price FROM my_index WHERE category = 'electronics'")
table = client.do_get(ticket).read_all()

Dispatch: UTF-8 ticket → SQL query. Binary protobuf ticket → df-distributed worker fragment.

REST (opt-in):

searcher:
  enable_datafusion_endpoint: true

POST /api/v1/{index_id}/datafusion
{"sql": "SELECT id, price FROM my_index ORDER BY id LIMIT 10"}

Distributed plan example

┌───── DistributedExec (coordinator) ──
│ SortPreservingMergeExec: [id ASC]
│   NetworkCoalesceExec: 3 tasks → 3 workers
└──────────────────────────────────────
  ┌───── Stage 1 (workers) ──
  │ DistributedUnionExec: worker0:[split-1] worker1:[split-2] worker2:[split-3]
  │   SortExec → HashJoinExec on (_doc_id, _segment_ord)
  │     InvertedIndexDataSource(query=true)
  │     FastFieldDataSource(pushed_filters=[...])  ← filters survive the wire
  │   ... ×3 splits
  └──────────────────────────────────────

Tests (15 passing)

Test file	Tests	What it covers
single_node.rs	3	Mock metastore → QuickwitTableProvider, full-text joins
join_plan.rs	6	Single-split and cross-split join plans (inv ⋈ f ⋈ d)
distributed_join_plan.rs	2	Distributed joins via df-distributed test infra
serve_integration.rs	1	Real TCP Flight servers + SearcherPool + distributed execution
schema_evolution.rs	3	NULL-fill, type mismatch error, SQL CAST

Remaining work

• Unnecessary RepartitionExec: Co-partitioned joins shouldn't shuffle — need to align target_partitions with segment count
• Cache affinity: Workers assigned round-robin; needs custom TaskEstimator for rendezvous-hash split affinity (team member working on plan-aware worker assignment in df-distributed)

[alexanderbianchi]

hacky, needs to be better

[alexanderbianchi]

of course would not be merged like this, going to need to move around dependencies