From 37032151d33c57217713239276421a3b6aee092d Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Wed, 15 Jul 2026 21:21:40 -0700 Subject: [PATCH] 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. --- Cargo.lock | 1 + docs/src/js/classes/Query.md | 8 +- docs/src/js/classes/QueryBase.md | 8 +- docs/src/js/classes/TakeQuery.md | 8 +- docs/src/js/classes/VectorQuery.md | 8 +- docs/src/js/globals.md | 1 + .../AnalyzePlanDistributedMetrics.md | 11 +++ nodejs/__test__/table.test.ts | 7 +- nodejs/lancedb/index.ts | 1 + nodejs/lancedb/query.ts | 15 +++- nodejs/src/query.rs | 77 ++++++++++++----- python/pyproject.toml | 5 +- python/python/lancedb/_lancedb.pyi | 17 +++- python/python/lancedb/query.py | 83 ++++++++++++++++--- python/python/lancedb/remote/table.py | 18 +++- python/python/lancedb/table.py | 28 +++++-- python/python/tests/test_hybrid_query.py | 12 ++- python/src/query.rs | 56 +++++++++++-- python/uv.lock | 20 +++-- rust/lancedb/Cargo.toml | 3 + rust/lancedb/src/query.rs | 30 +++++++ rust/lancedb/src/remote/table.rs | 54 +++++++++++- 22 files changed, 396 insertions(+), 75 deletions(-) create mode 100644 docs/src/js/type-aliases/AnalyzePlanDistributedMetrics.md diff --git a/Cargo.lock b/Cargo.lock index 6bcc03dc5..e759cbccb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5391,6 +5391,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-sql", "futures", + "goosefs-sdk", "half", "hf-hub", "http 1.4.2", diff --git a/docs/src/js/classes/Query.md b/docs/src/js/classes/Query.md index 15d78036b..c4bc3ff23 100644 --- a/docs/src/js/classes/Query.md +++ b/docs/src/js/classes/Query.md @@ -33,7 +33,7 @@ protected inner: Query | Promise; ### analyzePlan() ```ts -analyzePlan(): Promise +analyzePlan(distributedMetrics?): Promise ``` Executes the query and returns the physical query plan annotated with runtime metrics. @@ -41,6 +41,12 @@ Executes the query and returns the physical query plan annotated with runtime me This is useful for debugging and performance analysis, as it shows how the query was executed and includes metrics such as elapsed time, rows processed, and I/O statistics. +#### Parameters + +* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md) + How distributed worker metrics are displayed for remote query plans. + Defaults to `"aggregate"`. + #### Returns `Promise`<`string`> diff --git a/docs/src/js/classes/QueryBase.md b/docs/src/js/classes/QueryBase.md index 1dd81ff98..35c071525 100644 --- a/docs/src/js/classes/QueryBase.md +++ b/docs/src/js/classes/QueryBase.md @@ -38,7 +38,7 @@ protected inner: NativeQueryType | Promise; ### analyzePlan() ```ts -analyzePlan(): Promise +analyzePlan(distributedMetrics?): Promise ``` Executes the query and returns the physical query plan annotated with runtime metrics. @@ -46,6 +46,12 @@ Executes the query and returns the physical query plan annotated with runtime me This is useful for debugging and performance analysis, as it shows how the query was executed and includes metrics such as elapsed time, rows processed, and I/O statistics. +#### Parameters + +* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md) + How distributed worker metrics are displayed for remote query plans. + Defaults to `"aggregate"`. + #### Returns `Promise`<`string`> diff --git a/docs/src/js/classes/TakeQuery.md b/docs/src/js/classes/TakeQuery.md index 1329b99dc..0d658b32c 100644 --- a/docs/src/js/classes/TakeQuery.md +++ b/docs/src/js/classes/TakeQuery.md @@ -29,7 +29,7 @@ protected inner: TakeQuery | Promise; ### analyzePlan() ```ts -analyzePlan(): Promise +analyzePlan(distributedMetrics?): Promise ``` Executes the query and returns the physical query plan annotated with runtime metrics. @@ -37,6 +37,12 @@ Executes the query and returns the physical query plan annotated with runtime me This is useful for debugging and performance analysis, as it shows how the query was executed and includes metrics such as elapsed time, rows processed, and I/O statistics. +#### Parameters + +* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md) + How distributed worker metrics are displayed for remote query plans. + Defaults to `"aggregate"`. + #### Returns `Promise`<`string`> diff --git a/docs/src/js/classes/VectorQuery.md b/docs/src/js/classes/VectorQuery.md index 5ff8beec7..af3317029 100644 --- a/docs/src/js/classes/VectorQuery.md +++ b/docs/src/js/classes/VectorQuery.md @@ -51,7 +51,7 @@ addQueryVector(vector): VectorQuery ### analyzePlan() ```ts -analyzePlan(): Promise +analyzePlan(distributedMetrics?): Promise ``` Executes the query and returns the physical query plan annotated with runtime metrics. @@ -59,6 +59,12 @@ Executes the query and returns the physical query plan annotated with runtime me This is useful for debugging and performance analysis, as it shows how the query was executed and includes metrics such as elapsed time, rows processed, and I/O statistics. +#### Parameters + +* **distributedMetrics?**: [`AnalyzePlanDistributedMetrics`](../type-aliases/AnalyzePlanDistributedMetrics.md) + How distributed worker metrics are displayed for remote query plans. + Defaults to `"aggregate"`. + #### Returns `Promise`<`string`> diff --git a/docs/src/js/globals.md b/docs/src/js/globals.md index 907c89acc..e664bda29 100644 --- a/docs/src/js/globals.md +++ b/docs/src/js/globals.md @@ -118,6 +118,7 @@ ## Type Aliases +- [AnalyzePlanDistributedMetrics](type-aliases/AnalyzePlanDistributedMetrics.md) - [BaseTokenizer](type-aliases/BaseTokenizer.md) - [Data](type-aliases/Data.md) - [DataLike](type-aliases/DataLike.md) diff --git a/docs/src/js/type-aliases/AnalyzePlanDistributedMetrics.md b/docs/src/js/type-aliases/AnalyzePlanDistributedMetrics.md new file mode 100644 index 000000000..642931449 --- /dev/null +++ b/docs/src/js/type-aliases/AnalyzePlanDistributedMetrics.md @@ -0,0 +1,11 @@ +[**@lancedb/lancedb**](../README.md) • **Docs** + +*** + +[@lancedb/lancedb](../globals.md) / AnalyzePlanDistributedMetrics + +# Type Alias: AnalyzePlanDistributedMetrics + +```ts +type AnalyzePlanDistributedMetrics: "aggregate" | "per_worker" | "full"; +``` diff --git a/nodejs/__test__/table.test.ts b/nodejs/__test__/table.test.ts index dfe0af8c8..39345996a 100644 --- a/nodejs/__test__/table.test.ts +++ b/nodejs/__test__/table.test.ts @@ -2775,8 +2775,13 @@ describe("when calling analyzePlan", () => { .fill(1) .map(() => Math.random()); const plan = await table.query().nearestTo(queryVec).analyzePlan(); - console.log("Query Plan:\n", plan); // <--- Print the plan expect(plan).toMatch("AnalyzeExec"); + + const fullPlan = await table + .query() + .nearestTo(queryVec) + .analyzePlan("full"); + expect(fullPlan).toMatch("AnalyzeExec"); }); }); diff --git a/nodejs/lancedb/index.ts b/nodejs/lancedb/index.ts index 527f5f4a2..a54529153 100644 --- a/nodejs/lancedb/index.ts +++ b/nodejs/lancedb/index.ts @@ -93,6 +93,7 @@ export { QueryBase, VectorQuery, TakeQuery, + AnalyzePlanDistributedMetrics, QueryExecutionOptions, ColumnOrdering, FullTextSearchOptions, diff --git a/nodejs/lancedb/query.ts b/nodejs/lancedb/query.ts index 04a8a69d8..7c9d96750 100644 --- a/nodejs/lancedb/query.ts +++ b/nodejs/lancedb/query.ts @@ -79,6 +79,8 @@ export interface QueryExecutionOptions { timeoutMs?: number; } +export type AnalyzePlanDistributedMetrics = "aggregate" | "per_worker" | "full"; + export interface ColumnOrdering { columnName: string; ascending?: boolean; @@ -311,13 +313,20 @@ export class QueryBase< * KNNVectorDistance: metric=l2, metrics=[output_rows=1, elapsed_compute=114.333µs, output_batches=1] * LanceScan: uri=/path/to/data, projection=[vector], row_id=true, row_addr=false, ordered=false, metrics=[output_rows=1, elapsed_compute=103.626µs, bytes_read=549, iops=2, requests=2] * + * @param distributedMetrics - How distributed worker metrics are displayed for remote query plans. + * Defaults to `"aggregate"`. * @returns A query execution plan with runtime metrics for each step. */ - async analyzePlan(): Promise { + async analyzePlan( + distributedMetrics?: AnalyzePlanDistributedMetrics, + ): Promise { + const distributedMetricsMode = distributedMetrics ?? "aggregate"; if (this.inner instanceof Promise) { - return this.inner.then((inner) => inner.analyzePlan()); + return this.inner.then((inner) => + inner.analyzePlan(distributedMetricsMode), + ); } else { - return this.inner.analyzePlan(); + return this.inner.analyzePlan(distributedMetricsMode); } } diff --git a/nodejs/src/query.rs b/nodejs/src/query.rs index 6032619e3..7266e05fc 100644 --- a/nodejs/src/query.rs +++ b/nodejs/src/query.rs @@ -19,6 +19,7 @@ use lancedb::index::scalar::{ BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery, }; +use lancedb::query::AnalyzePlanDistributedMetrics; use lancedb::query::ExecutableQuery; use lancedb::query::Query as LanceDbQuery; use lancedb::query::QueryBase; @@ -47,6 +48,28 @@ impl From for LanceDbColumnOrdering { } } +fn analyze_plan_options( + distributed_metrics: Option, +) -> napi::Result { + let analyze_plan_distributed_metrics = + match distributed_metrics.as_deref().unwrap_or("aggregate") { + "aggregate" => AnalyzePlanDistributedMetrics::Aggregate, + "per_worker" => AnalyzePlanDistributedMetrics::PerWorker, + "full" => AnalyzePlanDistributedMetrics::Full, + mode => { + return Err(napi::Error::from_reason(format!( + "Invalid distributedMetrics value '{}'. Expected one of: \ + 'aggregate', 'per_worker', 'full'", + mode + ))); + } + }; + + let mut options = QueryExecutionOptions::default(); + options.analyze_plan_distributed_metrics = analyze_plan_distributed_metrics; + Ok(options) +} + fn bytes_to_arrow_array(data: Uint8Array, dtype: String) -> napi::Result> { let buf = arrow_buffer::Buffer::from(data.to_vec()); let num_bytes = buf.len(); @@ -200,13 +223,17 @@ impl Query { } #[napi(catch_unwind)] - pub async fn analyze_plan(&self) -> napi::Result { - self.inner.analyze_plan().await.map_err(|e| { - napi::Error::from_reason(format!( - "Failed to execute analyze plan: {}", - convert_error(&e) - )) - }) + pub async fn analyze_plan(&self, distributed_metrics: Option) -> napi::Result { + let options = analyze_plan_options(distributed_metrics)?; + self.inner + .analyze_plan_with_options(options) + .await + .map_err(|e| { + napi::Error::from_reason(format!( + "Failed to execute analyze plan: {}", + convert_error(&e) + )) + }) } } @@ -412,13 +439,17 @@ impl VectorQuery { } #[napi(catch_unwind)] - pub async fn analyze_plan(&self) -> napi::Result { - self.inner.analyze_plan().await.map_err(|e| { - napi::Error::from_reason(format!( - "Failed to execute analyze plan: {}", - convert_error(&e) - )) - }) + pub async fn analyze_plan(&self, distributed_metrics: Option) -> napi::Result { + let options = analyze_plan_options(distributed_metrics)?; + self.inner + .analyze_plan_with_options(options) + .await + .map_err(|e| { + napi::Error::from_reason(format!( + "Failed to execute analyze plan: {}", + convert_error(&e) + )) + }) } } @@ -491,13 +522,17 @@ impl TakeQuery { } #[napi(catch_unwind)] - pub async fn analyze_plan(&self) -> napi::Result { - self.inner.analyze_plan().await.map_err(|e| { - napi::Error::from_reason(format!( - "Failed to execute analyze plan: {}", - convert_error(&e) - )) - }) + pub async fn analyze_plan(&self, distributed_metrics: Option) -> napi::Result { + let options = analyze_plan_options(distributed_metrics)?; + self.inner + .analyze_plan_with_options(options) + .await + .map_err(|e| { + napi::Error::from_reason(format!( + "Failed to execute analyze plan: {}", + convert_error(&e) + )) + }) } } diff --git a/python/pyproject.toml b/python/pyproject.toml index 7806daed1..4e8391787 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -61,10 +61,11 @@ tests = [ "duckdb>=0.9.0", "pytz>=2023.3", "polars>=0.19, <=1.3.0", + "pyarrow<25", "pyarrow-stubs>=16.0", - "pylance>=5.0.0b5", + "pylance>=7,<8", "requests>=2.31.0", - "datafusion>=52,<53", + "datafusion>=53,<54", "opentelemetry-sdk>=1.30.0", ] dev = [ diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 368bdd069..d6e7f3407 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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: diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 1a4588808..49e3a171e 100644 --- a/python/python/lancedb/query.py +++ b/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): diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 75ad59b50..1006942ea 100644 --- a/python/python/lancedb/remote/table.py +++ b/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)) diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 05004eafa..4570c9788 100644 --- a/python/python/lancedb/table.py +++ b/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) diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 4ba61a4b8..aa3e3f5fd 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/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 diff --git a/python/src/query.rs b/python/src/query.rs index 2c682e38a..70ef3b281 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -19,6 +19,7 @@ use lancedb::index::scalar::{ BooleanQuery, BoostQuery, FtsQuery, FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery, }; +use lancedb::query::AnalyzePlanDistributedMetrics; use lancedb::query::QueryBase; use lancedb::query::QueryExecutionOptions; use lancedb::query::QueryFilter; @@ -42,6 +43,25 @@ use pyo3::{Borrowed, FromPyObject, exceptions::PyRuntimeError}; use pyo3::{PyErr, pyclass}; use pyo3::{exceptions::PyValueError, intern}; +fn analyze_plan_options(distributed_metrics: Option<&str>) -> PyResult { + let analyze_plan_distributed_metrics = match distributed_metrics.unwrap_or("aggregate") { + "aggregate" => AnalyzePlanDistributedMetrics::Aggregate, + "per_worker" => AnalyzePlanDistributedMetrics::PerWorker, + "full" => AnalyzePlanDistributedMetrics::Full, + mode => { + return Err(PyValueError::new_err(format!( + "Invalid distributed_metrics value '{}'. Expected one of: \ + 'aggregate', 'per_worker', 'full'", + mode + ))); + } + }; + + let mut options = QueryExecutionOptions::default(); + options.analyze_plan_distributed_metrics = analyze_plan_distributed_metrics; + Ok(options) +} + impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB { type Error = PyErr; @@ -571,11 +591,16 @@ impl Query { }) } - pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult> { + #[pyo3(signature = (distributed_metrics=None))] + pub fn analyze_plan( + self_: PyRef<'_, Self>, + distributed_metrics: Option, + ) -> PyResult> { let inner = self_.inner.clone(); + let options = analyze_plan_options(distributed_metrics.as_deref())?; future_into_py(self_.py(), async move { inner - .analyze_plan() + .analyze_plan_with_options(options) .await .map_err(|e| PyRuntimeError::new_err(e.to_string())) }) @@ -650,11 +675,16 @@ impl TakeQuery { }) } - pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult> { + #[pyo3(signature = (distributed_metrics=None))] + pub fn analyze_plan( + self_: PyRef<'_, Self>, + distributed_metrics: Option, + ) -> PyResult> { let inner = self_.inner.clone(); + let options = analyze_plan_options(distributed_metrics.as_deref())?; future_into_py(self_.py(), async move { inner - .analyze_plan() + .analyze_plan_with_options(options) .await .map_err(|e| PyRuntimeError::new_err(e.to_string())) }) @@ -777,14 +807,19 @@ impl FTSQuery { }) } - pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult> { + #[pyo3(signature = (distributed_metrics=None))] + pub fn analyze_plan( + self_: PyRef<'_, Self>, + distributed_metrics: Option, + ) -> PyResult> { let inner = self_ .inner .clone() .full_text_search(self_.fts_query.clone()); + let options = analyze_plan_options(distributed_metrics.as_deref())?; future_into_py(self_.py(), async move { inner - .analyze_plan() + .analyze_plan_with_options(options) .await .map_err(|e| PyRuntimeError::new_err(e.to_string())) }) @@ -958,11 +993,16 @@ impl VectorQuery { }) } - pub fn analyze_plan(self_: PyRef<'_, Self>) -> PyResult> { + #[pyo3(signature = (distributed_metrics=None))] + pub fn analyze_plan( + self_: PyRef<'_, Self>, + distributed_metrics: Option, + ) -> PyResult> { let inner = self_.inner.clone(); + let options = analyze_plan_options(distributed_metrics.as_deref())?; future_into_py(self_.py(), async move { inner - .analyze_plan() + .analyze_plan_with_options(options) .await .map_err(|e| PyRuntimeError::new_err(e.to_string())) }) diff --git a/python/uv.lock b/python/uv.lock index c45b3b5ba..7e6c9b5e6 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -850,19 +850,19 @@ nvtx = [ [[package]] name = "datafusion" -version = "52.3.0" +version = "53.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "pyarrow" }, { name = "typing-extensions", marker = "python_full_version < '3.13'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/db/d4/a5ad7b665a80008901892fde61dc667318db0652a955d706ddca3a224b5a/datafusion-52.3.0.tar.gz", hash = "sha256:2e8b02ad142b1a0d673f035d96a0944a640ac78275003d7e453cee4afe4a20a4", size = 205026, upload-time = "2026-03-16T10:54:07.739Z" } +sdist = { url = "https://files.pythonhosted.org/packages/58/2b/0f96f12b70839c93930c4e17d767fc32b6c77d548c78784128049e944701/datafusion-53.0.0.tar.gz", hash = "sha256:ba9a5ec06b5453fbd8710d6aeeb515a8bcac4b6c140e254409bb53a5f322ef22", size = 224267, upload-time = "2026-04-13T00:45:02.686Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/55/63/1bb0737988cefa77274b459d64fa4b57ba4cf755639a39733e9581b5d599/datafusion-52.3.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a73f02406b2985b9145dd97f8221a929c9ef3289a8ba64c6b52043e240938528", size = 31503230, upload-time = "2026-03-16T10:53:50.312Z" }, - { url = "https://files.pythonhosted.org/packages/d6/e3/ea3b79239953c3044d19d8e9581015da025b6640796db03799e435b17910/datafusion-52.3.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:118a1f0add6a3f91fcbc90c71819fe08750e2981637d5e7b346e099e94a20b8b", size = 28159497, upload-time = "2026-03-16T10:53:54.032Z" }, - { url = "https://files.pythonhosted.org/packages/24/c8/7d325feb4b7509ae03857fd7e164e95ec72e8c9f3dfd3178ec7f80d53977/datafusion-52.3.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:253ce7aee5fe84bd6ee290c20608114114bdb5115852617f97d3855d36ad9341", size = 30769154, upload-time = "2026-03-16T10:53:57.835Z" }, - { url = "https://files.pythonhosted.org/packages/37/ee/478689c69b3cb1ccabb2d52feac0c181f6cdf20b51a81df35344b1dab9a6/datafusion-52.3.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2af3469d2f06959bec88579ab107a72f965de18b32e607069bbdd0b859ed8dbb", size = 33060335, upload-time = "2026-03-16T10:54:01.715Z" }, - { url = "https://files.pythonhosted.org/packages/1c/48/01906ab5c1a70373c6874ac5192d03646fa7b94d9ff06e3f676cb6b0f43f/datafusion-52.3.0-cp310-abi3-win_amd64.whl", hash = "sha256:9fb35738cf4dbff672dbcfffc7332813024cb0ad2ab8cda1fb90b9054277ab0c", size = 33765807, upload-time = "2026-03-16T10:54:05.728Z" }, + { url = "https://files.pythonhosted.org/packages/af/4c/60e052813d81f1ffe3123ead013dbdd2cf961daa576cb9056cbb80228e6b/datafusion-53.0.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a0bd1a98d736571321416dc4ed361a9d1225da1ec9f6c5fad818d75f547697a7", size = 35774913, upload-time = "2026-04-13T00:44:46.235Z" }, + { url = "https://files.pythonhosted.org/packages/6e/59/beabe5301df3338d8206446cd624079e43bdad46e20377a6336017fb6ccf/datafusion-53.0.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:ce186a8d2405afd67e11e2fb75715019f16b00d070b8d0da89d8aa61cc74c8b5", size = 32667118, upload-time = "2026-04-13T00:44:50.269Z" }, + { url = "https://files.pythonhosted.org/packages/ae/94/636ab61ade98395daea6e733e225e9c7beef111c7c5b575ac851513e203c/datafusion-53.0.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:288a00a7ef03e2807a4667683f7560efd80d60ed1d41696ac15ca9ded14c8251", size = 35585824, upload-time = "2026-04-13T00:44:53.683Z" }, + { url = "https://files.pythonhosted.org/packages/34/80/b9f4889209af02f8d14bccb0e6f0519c329b072bc4d2595025a1303f144c/datafusion-53.0.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:8fef0004f0161fcfc556c025a7201f9cc3169aa3adb97a86419ebb34182d9efb", size = 38083690, upload-time = "2026-04-13T00:44:57.188Z" }, + { url = "https://files.pythonhosted.org/packages/4b/1a/ea4831fc6aeefedbcf186c9f6a273d507b1787c03cbb905bded7e1149a6a/datafusion-53.0.0-cp310-abi3-win_amd64.whl", hash = "sha256:4c8410f5f659b926677be6c7d443bbc05d825c078c970b7d8cf977ebcf948314", size = 38120687, upload-time = "2026-04-13T00:45:00.633Z" }, ] [[package]] @@ -1931,6 +1931,7 @@ tests = [ { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and python_full_version < '3.14'" }, { name = "pandas", version = "3.0.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.14'" }, { name = "polars" }, + { name = "pyarrow" }, { name = "pyarrow-stubs" }, { name = "pylance" }, { name = "pytest" }, @@ -1950,7 +1951,7 @@ requires-dist = [ { name = "botocore", marker = "extra == 'embeddings'", specifier = ">=1.31.57" }, { name = "cohere", marker = "extra == 'embeddings'", specifier = ">=4.0" }, { name = "colpali-engine", marker = "extra == 'embeddings'", specifier = ">=0.3.10" }, - { name = "datafusion", marker = "extra == 'tests'", specifier = ">=52,<53" }, + { name = "datafusion", marker = "extra == 'tests'", specifier = ">=53,<54" }, { name = "deprecation", specifier = ">=2.1.0" }, { name = "duckdb", marker = "extra == 'tests'", specifier = ">=0.9.0" }, { name = "google-genai", marker = "extra == 'embeddings'", specifier = ">=1.0.0" }, @@ -1978,10 +1979,11 @@ requires-dist = [ { name = "polars", marker = "extra == 'tests'", specifier = ">=0.19,<=1.3.0" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.5.0" }, { name = "pyarrow", specifier = ">=16" }, + { name = "pyarrow", marker = "extra == 'tests'", specifier = "<25" }, { name = "pyarrow-stubs", marker = "extra == 'tests'", specifier = ">=16.0" }, { name = "pydantic", specifier = ">=1.10" }, { name = "pylance", marker = "extra == 'pylance'", specifier = ">=5.0.0b5" }, - { name = "pylance", marker = "extra == 'tests'", specifier = ">=5.0.0b5" }, + { name = "pylance", marker = "extra == 'tests'", specifier = ">=7,<8" }, { name = "pyright", marker = "extra == 'dev'", specifier = ">=1.1.350" }, { name = "pytest", marker = "extra == 'tests'", specifier = ">=7.0" }, { name = "pytest-asyncio", marker = "extra == 'tests'", specifier = ">=0.21" }, diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index f71e4193a..791ee3d64 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -50,6 +50,8 @@ lance-namespace = { workspace = true } lance-namespace-impls = { workspace = true } metrics = { workspace = true, optional = true } metrics-util = { workspace = true, optional = true } +# Pin the transitive GooseFS SDK until the 0.1.6 compile break is fixed upstream. +goosefs-sdk = { version = "=0.1.5", optional = true } moka = { workspace = true } pin-project = { workspace = true } tokio = { version = "1.23", features = ["rt-multi-thread", "sync"] } @@ -132,6 +134,7 @@ azure = [ ] cos = ["lance/tencent", "lance-io/tencent"] goosefs = [ + "dep:goosefs-sdk", "lance/goosefs", "lance-io/goosefs", "lance-namespace-impls/dir-goosefs", diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index f3dbca1ab..727d87c38 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -614,6 +614,12 @@ pub struct QueryExecutionOptions { pub max_batch_length: u32, /// Max duration to wait for the query to execute before timing out. pub timeout: Option, + /// How distributed worker metrics should be displayed by + /// [`ExecutableQuery::analyze_plan`]. + /// + /// This only affects remote distributed query plans. Local query execution + /// ignores this option. + pub analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics, } impl Default for QueryExecutionOptions { @@ -621,6 +627,7 @@ impl Default for QueryExecutionOptions { Self { max_batch_length: 1024, timeout: None, + analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::Aggregate, } } } @@ -633,6 +640,29 @@ impl QueryExecutionOptions { } } +/// How distributed worker metrics are displayed in analyzed query plans. +#[non_exhaustive] +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum AnalyzePlanDistributedMetrics { + /// Preserve the legacy output: aggregate worker metrics into one synthetic tree. + #[default] + Aggregate, + /// Render one raw worker-side tree per distributed worker. + PerWorker, + /// Render the aggregate tree followed by the raw per-worker trees. + Full, +} + +impl AnalyzePlanDistributedMetrics { + pub(crate) fn as_query_param(self) -> &'static str { + match self { + Self::Aggregate => "aggregate", + Self::PerWorker => "per_worker", + Self::Full => "full", + } + } +} + /// A trait for a query object that can be executed to get results /// /// There are various kinds of queries but they all return results diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 3aa4246e9..8684d9070 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -36,7 +36,7 @@ use crate::{DistanceType, Error}; use crate::{ error::Result, index::{IndexBuilder, IndexConfig}, - query::QueryExecutionOptions, + query::{AnalyzePlanDistributedMetrics, QueryExecutionOptions}, table::{ AddDataBuilder, BaseTable, OptimizeAction, OptimizeStats, TableDefinition, UpdateBuilder, merge::MergeInsertBuilder, @@ -1993,9 +1993,16 @@ impl BaseTable for RemoteTable { async fn analyze_plan( &self, query: &AnyQuery, - _options: QueryExecutionOptions, + options: QueryExecutionOptions, ) -> Result { - let request = self.post_read(&format!("/v1/table/{}/analyze_plan/", self.identifier)); + let mut request = self.post_read(&format!("/v1/table/{}/analyze_plan/", self.identifier)); + + if options.analyze_plan_distributed_metrics != AnalyzePlanDistributedMetrics::Aggregate { + request = request.query(&[( + "distributed_metrics", + options.analyze_plan_distributed_metrics.as_query_param(), + )]); + } let query_bodies = self.prepare_query_bodies(query).await?; let requests: Vec = query_bodies @@ -2840,7 +2847,10 @@ mod tests { use crate::{ DistanceType, Error, Table, index::{Index, IndexStatistics, IndexType, vector::IvfPqIndexBuilder}, - query::{ColumnOrdering, ExecutableQuery, QueryBase}, + query::{ + AnalyzePlanDistributedMetrics, ColumnOrdering, ExecutableQuery, QueryBase, + QueryExecutionOptions, + }, remote::ARROW_FILE_CONTENT_TYPE, }; @@ -4048,6 +4058,42 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn test_analyze_plan_distributed_metrics_query_param() { + let table = Table::new_with_handler("my_table", |request| { + assert_eq!(request.method(), "POST"); + assert_eq!(request.url().path(), "/v1/table/my_table/analyze_plan/"); + assert_eq!( + request + .url() + .query_pairs() + .find(|(k, _)| k == "distributed_metrics"), + Some(("distributed_metrics".into(), "per_worker".into())) + ); + + let body = request.body().unwrap().as_bytes().unwrap(); + let body: serde_json::Value = serde_json::from_slice(body).unwrap(); + assert_eq!(body["k"], serde_json::json!(1)); + + http::Response::builder() + .status(200) + .body(r#""analyzed plan""#) + .unwrap() + }); + + let result = table + .query() + .limit(1) + .analyze_plan_with_options(QueryExecutionOptions { + analyze_plan_distributed_metrics: AnalyzePlanDistributedMetrics::PerWorker, + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(result, "analyzed plan"); + } + #[tokio::test] async fn test_query_structured_fts() { let table =