feat: support distributed analyze plan metrics in clients (#3675)

Adds client-side support for analyze_plan distributed metrics modes
across Rust, Python, and TypeScript clients. Defaults to aggregate for
backward compatibility and sends the remote distributed_metrics
parameter only when a non-default mode is requested.

python/python/lancedb/_lancedb.pyi

@@ -30,6 +30,7 @@ from .types import BaseTokenizerType
 IvfHnswPq: type[HnswPq] = HnswPq
 IvfHnswSq: type[HnswSq] = HnswSq
 IvfHnswFlat: type[HnswFlat] = HnswFlat
+AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
 
 class MetricPoint:
     name: str
@@ -393,7 +394,9 @@ class Query:
         self, max_batch_length: Optional[int], timeout: Optional[timedelta]
     ) -> RecordBatchStream: ...
     async def explain_plan(self, verbose: Optional[bool]) -> str: ...
-    async def analyze_plan(self) -> str: ...
+    async def analyze_plan(
+        self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
+    ) -> str: ...
     def to_query_request(self) -> PyQueryRequest: ...
 
 class TakeQuery:
@@ -401,6 +404,10 @@ class TakeQuery:
     def with_row_id(self): ...
     async def output_schema(self) -> pa.Schema: ...
     async def execute(self) -> RecordBatchStream: ...
+    async def explain_plan(self, verbose: Optional[bool]) -> str: ...
+    async def analyze_plan(
+        self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
+    ) -> str: ...
     def to_query_request(self) -> PyQueryRequest: ...
 
 class FTSQuery:
@@ -421,6 +428,10 @@ class FTSQuery:
     async def execute(
         self, max_batch_length: Optional[int], timeout: Optional[timedelta]
     ) -> RecordBatchStream: ...
+    async def explain_plan(self, verbose: Optional[bool]) -> str: ...
+    async def analyze_plan(
+        self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
+    ) -> str: ...
     def to_query_request(self) -> PyQueryRequest: ...
 
 class VectorQuery:
@@ -443,6 +454,10 @@ class VectorQuery:
     def bypass_vector_index(self): ...
     def nearest_to_text(self, query: dict) -> HybridQuery: ...
     def order_by(self, ordering: Optional[List[ColumnOrdering]]): ...
+    async def explain_plan(self, verbose: Optional[bool]) -> str: ...
+    async def analyze_plan(
+        self, distributed_metrics: Optional[AnalyzePlanDistributedMetrics] = None
+    ) -> str: ...
     def to_query_request(self) -> PyQueryRequest: ...
 
 class HybridQuery:

python/python/lancedb/query.py

@@ -79,6 +79,7 @@ if TYPE_CHECKING:
         from typing_extensions import Self
 
 T = TypeVar("T", bound="LanceModel")
+AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
 
 
 @runtime_checkable
@@ -1372,7 +1373,9 @@ class LanceQueryBuilder(ABC):
         self._order_by = ordering
         return self
 
-    def analyze_plan(self) -> str:
+    def analyze_plan(
+        self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
+    ) -> str:
         """
         Run the query and return its execution plan with runtime metrics.
 
@@ -1410,12 +1413,22 @@ class LanceQueryBuilder(ABC):
                       fragments_scanned=..., ranges_scanned=1, rows_scanned=1,
                       bytes_read=..., iops=..., requests=..., task_wait_time=...]
 
+        Parameters
+        ----------
+        distributed_metrics : Literal["aggregate", "per_worker", "full"]
+            Defaults to "aggregate".
+            How distributed worker metrics are displayed for remote query plans.
+            "aggregate" preserves the legacy summary, "per_worker" shows each
+            worker separately, and "full" includes both.
+
         Returns
         -------
         plan : str
             The physical query execution plan with runtime metrics.
         """
-        return self._table._analyze_plan(self.to_query_object())
+        return self._table._analyze_plan(
+            self.to_query_object(), distributed_metrics=distributed_metrics
+        )
 
     def vector(self, vector: Union[np.ndarray, list]) -> Self:
         """Set the vector to search for.
@@ -2581,9 +2594,17 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
         indented_fts = "\n".join("  " + line for line in fts_plan.splitlines())
         return f"{reranker_label}\n  {indented_vector}\n  {indented_fts}"
 
-    def analyze_plan(self):
+    def analyze_plan(
+        self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
+    ) -> str:
         """Execute the query and display with runtime metrics.
 
+        Parameters
+        ----------
+        distributed_metrics : Literal["aggregate", "per_worker", "full"]
+            Defaults to "aggregate".
+            How distributed worker metrics are displayed for remote query plans.
+
         Returns
         -------
         plan : str
@@ -2591,9 +2612,19 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
         self._create_query_builders()
 
         results = ["Vector Search Plan:"]
-        results.append(self._table._analyze_plan(self._vector_query.to_query_object()))
+        results.append(
+            self._table._analyze_plan(
+                self._vector_query.to_query_object(),
+                distributed_metrics=distributed_metrics,
+            )
+        )
         results.append("FTS Search Plan:")
-        results.append(self._table._analyze_plan(self._fts_query.to_query_object()))
+        results.append(
+            self._table._analyze_plan(
+                self._fts_query.to_query_object(),
+                distributed_metrics=distributed_metrics,
+            )
+        )
         return "\n".join(results)
 
     def _create_query_builders(self):
@@ -3080,14 +3111,22 @@ class AsyncQueryBase(object):
         """  # noqa: E501
         return await self._inner.explain_plan(verbose)
 
-    async def analyze_plan(self):
+    async def analyze_plan(
+        self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
+    ) -> str:
         """Execute the query and display with runtime metrics.
 
+        Parameters
+        ----------
+        distributed_metrics : Literal["aggregate", "per_worker", "full"]
+            Defaults to "aggregate".
+            How distributed worker metrics are displayed for remote query plans.
+
         Returns
         -------
         plan : str
         """
-        return await self._inner.analyze_plan()
+        return await self._inner.analyze_plan(distributed_metrics)
 
 
 class AsyncStandardQuery(AsyncQueryBase):
@@ -3866,7 +3905,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         indented_fts = "\n".join("  " + line for line in fts_plan.splitlines())
         return f"{self._reranker}\n  {indented_vector}\n  {indented_fts}"
 
-    async def analyze_plan(self):
+    async def analyze_plan(
+        self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
+    ) -> str:
         """
         Execute the query and return the physical execution plan with runtime metrics.
 
@@ -3875,14 +3916,24 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         elapsed time, I/O stats, and more. It’s useful for debugging and
         performance analysis.
 
+        Parameters
+        ----------
+        distributed_metrics : Literal["aggregate", "per_worker", "full"]
+            Defaults to "aggregate".
+            How distributed worker metrics are displayed for remote query plans.
+
         Returns
         -------
         plan : str
         """
         results = ["Vector Search Query:"]
-        results.append(await self._inner.to_vector_query().analyze_plan())
+        results.append(
+            await self._inner.to_vector_query().analyze_plan(distributed_metrics)
+        )
         results.append("FTS Search Query:")
-        results.append(await self._inner.to_fts_query().analyze_plan())
+        results.append(
+            await self._inner.to_fts_query().analyze_plan(distributed_metrics)
+        )
 
         return "\n".join(results)
 
@@ -4166,14 +4217,22 @@ class BaseQueryBuilder(object):
         """  # noqa: E501
         return LOOP.run(self._inner.explain_plan(verbose))
 
-    def analyze_plan(self):
+    def analyze_plan(
+        self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
+    ) -> str:
         """Execute the query and display with runtime metrics.
 
+        Parameters
+        ----------
+        distributed_metrics : Literal["aggregate", "per_worker", "full"]
+            Defaults to "aggregate".
+            How distributed worker metrics are displayed for remote query plans.
+
         Returns
         -------
         plan : str
         """
-        return LOOP.run(self._inner.analyze_plan())
+        return LOOP.run(self._inner.analyze_plan(distributed_metrics))
 
 
 class LanceTakeQueryBuilder(BaseQueryBuilder):

python/python/lancedb/remote/table.py

@@ -56,7 +56,12 @@ from lancedb.merge import LanceMergeInsertBuilder
 from lancedb.embeddings import EmbeddingFunctionRegistry
 from lancedb.table import _normalize_progress
 
-from ..query import LanceVectorQueryBuilder, LanceQueryBuilder, LanceTakeQueryBuilder
+from ..query import (
+    AnalyzePlanDistributedMetrics,
+    LanceQueryBuilder,
+    LanceTakeQueryBuilder,
+    LanceVectorQueryBuilder,
+)
 from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags
 from ..types import BaseTokenizerType
 
@@ -718,8 +723,15 @@ class RemoteTable(Table):
     def _explain_plan(self, query: Query, verbose: Optional[bool] = False) -> str:
         return LOOP.run(self._table._explain_plan(query, verbose))
 
-    def _analyze_plan(self, query: Query) -> str:
-        return LOOP.run(self._table._analyze_plan(query))
+    def _analyze_plan(
+        self,
+        query: Query,
+        *,
+        distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
+    ) -> str:
+        return LOOP.run(
+            self._table._analyze_plan(query, distributed_metrics=distributed_metrics)
+        )
 
     def _output_schema(self, query: Query) -> pa.Schema:
         return LOOP.run(self._table._output_schema(query))

python/python/lancedb/table.py

@@ -73,6 +73,7 @@ from .expr import Expr
 from .merge import LanceMergeInsertBuilder
 from .pydantic import LanceModel, model_to_dict
 from .query import (
+    AnalyzePlanDistributedMetrics,
     AsyncFTSQuery,
     AsyncHybridQuery,
     AsyncQuery,
@@ -1552,7 +1553,12 @@ class Table(ABC):
     def _explain_plan(self, query: Query, verbose: Optional[bool] = False) -> str: ...
 
     @abstractmethod
-    def _analyze_plan(self, query: Query) -> str: ...
+    def _analyze_plan(
+        self,
+        query: Query,
+        *,
+        distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
+    ) -> str: ...
 
     @abstractmethod
     def _output_schema(self, query: Query) -> pa.Schema: ...
@@ -3630,8 +3636,15 @@ class LanceTable(Table):
     def _explain_plan(self, query: Query, verbose: Optional[bool] = False) -> str:
         return LOOP.run(self._table._explain_plan(query, verbose))
 
-    def _analyze_plan(self, query: Query) -> str:
-        return LOOP.run(self._table._analyze_plan(query))
+    def _analyze_plan(
+        self,
+        query: Query,
+        *,
+        distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
+    ) -> str:
+        return LOOP.run(
+            self._table._analyze_plan(query, distributed_metrics=distributed_metrics)
+        )
 
     def _output_schema(self, query: Query) -> pa.Schema:
         return LOOP.run(self._table._output_schema(query))
@@ -5390,10 +5403,15 @@ class AsyncTable:
         async_query = self._sync_query_to_async(query)
         return await async_query.explain_plan(verbose)
 
-    async def _analyze_plan(self, query: Query) -> str:
+    async def _analyze_plan(
+        self,
+        query: Query,
+        *,
+        distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate",
+    ) -> str:
         # This method is used by the sync table
         async_query = self._sync_query_to_async(query)
-        return await async_query.analyze_plan()
+        return await async_query.analyze_plan(distributed_metrics)
 
     async def _output_schema(self, query: Query) -> pa.Schema:
         async_query = self._sync_query_to_async(query)

python/python/tests/test_hybrid_query.py

@@ -196,18 +196,26 @@ async def test_analyze_plan(table: AsyncTable):
 def test_hybrid_phrase_query_is_preserved_in_analyze_plan():
     table = mock.Mock()
     analyzed_queries = []
-    table._analyze_plan.side_effect = lambda query: analyzed_queries.append(query) or ""
+    distributed_metric_modes = []
+
+    def capture_query(query, *, distributed_metrics="aggregate"):
+        analyzed_queries.append(query)
+        distributed_metric_modes.append(distributed_metrics)
+        return ""
+
+    table._analyze_plan.side_effect = capture_query
 
     (
         LanceHybridQueryBuilder(table)
         .vector([0.1, 0.2])
         .text("puppy runs")
         .phrase_query()
-        .analyze_plan()
+        .analyze_plan(distributed_metrics="full")
     )
 
     assert len(analyzed_queries) == 2
     assert analyzed_queries[1].full_text_query.query == '"puppy runs"'
+    assert distributed_metric_modes == ["full", "full"]
 
 
 @pytest.fixture