refactor: move ann.proto compilation to lance crate, use VectorMetricType enum

Address PR review feedback:
- Move ann.proto from lance-datafusion to lance crate (package lance.pb)
  so it can reference types from lance-index without circular deps
- Use lance.index.pb.VectorMetricType enum instead of string for metric_type
- Rename key_arrow_ipc to query_vector_arrow_ipc
- Make minimum_nprobes and dist_q_c optional
- Add comment about prefilter_source handled by DataFusion child serialization
- Add comment explaining ANNIvfSubIndexExecProto purpose

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: LuQQiu

Cargo.lock

@@ -3,32 +3,39 @@
 
 syntax = "proto3";
 
-package lance.datafusion;
+package lance.pb;
 
 import "table_identifier.proto";
 import "table.proto";
+import "index.proto";
 
 // Serialized vector query parameters.
 message VectorQueryProto {
   // Query vector as Arrow IPC bytes (supports Float16, Float32, Float64, UInt8, etc.)
-  bytes key_arrow_ipc = 1;
+  bytes query_vector_arrow_ipc = 1;
   string column = 2;
   uint32 k = 3;
   optional float lower_bound = 4;
   optional float upper_bound = 5;
-  uint32 minimum_nprobes = 6;
+  optional uint32 minimum_nprobes = 6;
   optional uint32 maximum_nprobes = 7;
   optional uint32 ef = 8;
   optional uint32 refine_factor = 9;
-  // Distance metric type as string ("l2", "cosine", "dot", "hamming").
-  // Absent means None (use the index's default metric).
-  optional string metric_type = 10;
+  // Distance metric type. Absent means None (use the index's default metric).
+  optional lance.index.pb.VectorMetricType metric_type = 10;
   bool use_index = 11;
-  float dist_q_c = 12;
+  optional float dist_q_c = 12;
 }
 
+// Serializable form of ANNIvfSubIndexExec — the IVF sub-index search node.
+//
+// Note: ANNIvfSubIndexExec.prefilter_source (child ExecutionPlan) is NOT
+// serialized here. DataFusion's PhysicalExtensionCodec handles child plans
+// automatically via children() / with_new_children(). The codec receives
+// deserialized children in the `inputs` parameter of try_decode and
+// reconstructs the PreFilterSource from them.
 message ANNIvfSubIndexExecProto {
   VectorQueryProto query = 1;
-  TableIdentifier table = 2;
+  lance.datafusion.TableIdentifier table = 2;
   repeated lance.table.IndexMetadata indices = 3;
 }

protos/ann.proto

@@ -20,7 +20,6 @@ fn main() -> Result<()> {
         &[
             "./protos/table_identifier.proto",
             "./protos/filtered_read.proto",
-            "./protos/ann.proto",
         ],
         &["./protos"],
     )?;

python/Cargo.lock

@@ -84,6 +84,10 @@ semver.workspace = true
 tokio-stream = { workspace = true }
 tokio-util = { workspace = true }
 
+[build-dependencies]
+prost-build.workspace = true
+protobuf-src = { version = "2.1", optional = true }
+
 [target.'cfg(target_os = "linux")'.dev-dependencies]
 pprof.workspace = true
 # Need this so we can prevent dynamic linking in binaries (see cli feature)
@@ -128,6 +132,7 @@ dynamodb = ["lance-table/dynamodb", "dep:aws-sdk-dynamodb"]
 dynamodb_tests = ["dynamodb"]
 substrait = ["lance-datafusion/substrait"]
 protoc = [
+    "dep:protobuf-src",
     "lance-encoding/protoc",
     "lance-file/protoc",
     "lance-index/protoc",

rust/lance-datafusion/build.rs

@@ -0,0 +1,24 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::io::Result;
+
+fn main() -> Result<()> {
+    println!("cargo:rerun-if-changed=protos");
+
+    #[cfg(feature = "protoc")]
+    // Use vendored protobuf compiler if requested.
+    unsafe {
+        std::env::set_var("PROTOC", protobuf_src::protoc());
+    }
+
+    let mut prost_build = prost_build::Config::new();
+    prost_build.extern_path(".lance.table", "::lance_table::format::pb");
+    prost_build.extern_path(".lance.index.pb", "::lance_index::pb");
+    prost_build.extern_path(".lance.datafusion", "::lance_datafusion::pb");
+    prost_build.protoc_arg("--experimental_allow_proto3_optional");
+    prost_build.enable_type_names();
+    prost_build.compile_protos(&["./protos/ann.proto"], &["./protos"])?;
+
+    Ok(())
+}

rust/lance/Cargo.toml

@@ -3,7 +3,7 @@
 
 //! Protobuf serialization for [`ANNIvfSubIndexExec`].
 //!
-//! Proto message definitions live in `lance-datafusion` (see `pb`).
+//! Proto message definitions live in `crate::pb` (compiled from `ann.proto`).
 //! Conversion functions live here because they need access to `ANNIvfSubIndexExec`
 //! and `Dataset`, which are defined in this crate.
 //!
@@ -15,13 +15,14 @@ use std::sync::Arc;
 use arrow_array::RecordBatch;
 use arrow_schema::{Field, Schema as ArrowSchema};
 use lance_core::{Error, Result};
-use lance_datafusion::pb;
+use lance_index::pb as index_pb;
 use lance_index::vector::Query;
 use lance_linalg::distance::DistanceType;
 use lance_table::format::IndexMetadata;
 use lance_table::format::pb as table_pb;
 
 use crate::Dataset;
+use crate::pb;
 
 use super::knn::ANNIvfSubIndexExec;
 use super::table_identifier::{resolve_dataset, table_identifier_from_dataset};
@@ -75,49 +76,53 @@ fn query_key_from_ipc_bytes(bytes: &[u8]) -> Result<arrow_array::ArrayRef> {
 }
 
 pub fn query_to_proto(query: &Query) -> Result<pb::VectorQueryProto> {
-    let key_arrow_ipc = query_key_to_ipc_bytes(query.key.as_ref())?;
+    let query_vector_arrow_ipc = query_key_to_ipc_bytes(query.key.as_ref())?;
 
-    let metric_type = query.metric_type.map(|dt| dt.to_string());
+    let metric_type = query
+        .metric_type
+        .map(|dt| index_pb::VectorMetricType::from(dt) as i32);
 
     Ok(pb::VectorQueryProto {
-        key_arrow_ipc,
+        query_vector_arrow_ipc,
         column: query.column.clone(),
         k: query.k as u32,
         lower_bound: query.lower_bound,
         upper_bound: query.upper_bound,
-        minimum_nprobes: query.minimum_nprobes as u32,
+        minimum_nprobes: Some(query.minimum_nprobes as u32),
         maximum_nprobes: query.maximum_nprobes.map(|n| n as u32),
         ef: query.ef.map(|n| n as u32),
         refine_factor: query.refine_factor,
         metric_type,
         use_index: query.use_index,
-        dist_q_c: query.dist_q_c,
+        dist_q_c: Some(query.dist_q_c),
     })
 }
 
 pub fn query_from_proto(proto: pb::VectorQueryProto) -> Result<Query> {
-    let key = query_key_from_ipc_bytes(&proto.key_arrow_ipc)?;
+    let key = query_key_from_ipc_bytes(&proto.query_vector_arrow_ipc)?;
 
     let metric_type = proto
         .metric_type
-        .as_deref()
-        .map(DistanceType::try_from)
-        .transpose()
-        .map_err(|e| Error::internal(format!("Invalid distance type: {e}")))?;
+        .map(|v| {
+            index_pb::VectorMetricType::try_from(v)
+                .map(DistanceType::from)
+                .map_err(|_| Error::internal(format!("Invalid VectorMetricType value: {v}")))
+        })
+        .transpose()?;
 
     Ok(Query {
         column: proto.column,
         key,
         k: proto.k as usize,
         lower_bound: proto.lower_bound,
         upper_bound: proto.upper_bound,
-        minimum_nprobes: proto.minimum_nprobes as usize,
+        minimum_nprobes: proto.minimum_nprobes.unwrap_or(0) as usize,
         maximum_nprobes: proto.maximum_nprobes.map(|n| n as usize),
         ef: proto.ef.map(|n| n as usize),
         refine_factor: proto.refine_factor,
         metric_type,
         use_index: proto.use_index,
-        dist_q_c: proto.dist_q_c,
+        dist_q_c: proto.dist_q_c.unwrap_or(0.0),
     })
 }
 

rust/lance/build.rs

@@ -85,6 +85,11 @@ pub mod session;
 pub mod table;
 pub mod utils;
 
+pub mod pb {
+    #![allow(clippy::use_self)]
+    include!(concat!(env!("OUT_DIR"), "/lance.pb.rs"));
+}
+
 pub use blob::{BlobArrayBuilder, blob_field};
 pub use dataset::Dataset;
 use lance_index::vector::DIST_COL;