feat(index): support float16 and float64 in IVF_FLAT (#6476)

## Feature

### What is the new feature?
This PR adds native `float16` and `float64` support for `IVF_FLAT` and
`IVF_HNSW_FLAT`.

### Why do we need this feature?
Flat IVF indexing previously only worked end to end for `float32`. That
meant users could not build, merge, reload, and query flat IVF indexes
on `float16` or `float64` vectors without running into
`Float32`-specific assumptions in flat storage, writer initialization,
and merge/query paths.

### How does it work?
The implementation makes flat IVF paths dispatch on the actual Arrow
element type from stored flat data instead of assuming `Float32`.

- `FlatFloatStorage` now dispatches distance calculators for `float16`,
`float32`, and `float64`.
- Query/training helpers that previously special-cased `Float32` now
accept the native float dtype where needed.
- Tests now cover flat storage distance, partition serde roundtrip, IVF
create/query/remap, and distributed merge behavior for `float16` /
`float64`.

## Validation

- `cargo fmt --all`
- `cargo check -p lance-index --lib`
- `cargo check -p lance --lib`
- `cargo test -p lance-index test_flat_float_storage_distance_f16 --
--nocapture`
- `cargo test -p lance-index
test_merge_ivf_flat_preserves_float64_schema -- --nocapture`
- `cargo test -p lance test_build_ivf_flat -- --nocapture`
- `cargo test -p lance test_create_ivf_hnsw_flat -- --nocapture`
- `cargo test -p lance test_create_ivf_flat_f16 -- --nocapture`

## Benchmark Note

I also benchmarked float32 `IVF_FLAT` before vs after, no obvious
performance diffs

Author: BubbleCal

rust/lance-index/src/vector/distributed/index_merger.rs

@@ -176,6 +176,7 @@ pub async fn init_writer_for_flat(
     object_store: &lance_io::object_store::ObjectStore,
     aux_out: &object_store::path::Path,
     d0: usize,
+    item_type: &DataType,
     dt: DistanceType,
     format_version: LanceFileVersion,
 ) -> Result<FileWriter> {
@@ -184,7 +185,7 @@ pub async fn init_writer_for_flat(
         Field::new(
             crate::vector::flat::storage::FLAT_COLUMN,
             DataType::FixedSizeList(
-                Arc::new(Field::new("item", DataType::Float32, true)),
+                Arc::new(Field::new("item", item_type.clone(), true)),
                 d0 as i32,
             ),
             true,
@@ -1129,9 +1130,13 @@ pub async fn merge_partial_vector_auxiliary_files(
                     .iter()
                     .find(|f| f.name() == crate::vector::flat::storage::FLAT_COLUMN)
                     .ok_or_else(|| Error::index("FLAT column missing".to_string()))?;
-                let d0 = match flat_field.data_type() {
-                    DataType::FixedSizeList(_, sz) => *sz as usize,
-                    _ => 0,
+                let (d0, item_type) = match flat_field.data_type() {
+                    DataType::FixedSizeList(item, sz) => (*sz as usize, item.data_type().clone()),
+                    _ => {
+                        return Err(Error::index(
+                            "FLAT column is not a FixedSizeList in shard schema".to_string(),
+                        ));
+                    }
                 };
                 dim.get_or_insert(d0);
                 if let Some(dprev) = dim
@@ -1140,53 +1145,29 @@ pub async fn merge_partial_vector_auxiliary_files(
                     return Err(Error::index("Dimension mismatch across shards".to_string()));
                 }
                 if v2w_opt.is_none() {
-                    let w = init_writer_for_flat(object_store, &aux_out, d0, dt, fv).await?;
+                    let w = init_writer_for_flat(object_store, &aux_out, d0, &item_type, dt, fv)
+                        .await?;
                     v2w_opt = Some(w);
                 }
             }
             SupportedIvfIndexType::IvfHnswFlat => {
-                // Treat HNSW_FLAT storage the same as FLAT: create schema with ROW_ID + flat vectors
-                // Determine dimension from shard schema (flat column) or fallback to STORAGE_METADATA_KEY
+                // Treat HNSW_FLAT storage the same as FLAT and preserve the actual flat item dtype.
                 let schema_arrow: ArrowSchema = reader.schema().as_ref().into();
-                // Try to find flat column and derive dim
-                let d0 = if let Some(flat_field) = schema_arrow
+                let Some(flat_field) = schema_arrow
                     .fields
                     .iter()
                     .find(|f| f.name() == crate::vector::flat::storage::FLAT_COLUMN)
-                {
-                    match flat_field.data_type() {
-                        DataType::FixedSizeList(_, sz) => *sz as usize,
-                        _ => 0,
-                    }
-                } else {
-                    // Fallback to STORAGE_METADATA_KEY FlatMetadata
-                    if let Some(storage_meta_json) = reader
-                        .metadata()
-                        .file_schema
-                        .metadata
-                        .get(STORAGE_METADATA_KEY)
-                    {
-                        let storage_metadata_vec: Vec<String> =
-                            serde_json::from_str(storage_meta_json).map_err(|e| {
-                                Error::index(format!("Failed to parse storage metadata: {}", e))
-                            })?;
-                        if let Some(first_meta) = storage_metadata_vec.first() {
-                            if let Ok(flat_meta) = serde_json::from_str::<FlatMetadata>(first_meta)
-                            {
-                                flat_meta.dim
-                            } else {
-                                return Err(Error::index(
-                                    "FLAT metadata missing in storage metadata".to_string(),
-                                ));
-                            }
-                        } else {
-                            return Err(Error::index(
-                                "FLAT metadata missing in storage metadata".to_string(),
-                            ));
-                        }
-                    } else {
+                else {
+                    return Err(Error::index(
+                        "FLAT column missing from IVF_HNSW_FLAT shard schema".to_string(),
+                    ));
+                };
+                let (d0, item_type) = match flat_field.data_type() {
+                    DataType::FixedSizeList(item, sz) => (*sz as usize, item.data_type().clone()),
+                    _ => {
                         return Err(Error::index(
-                            "FLAT column missing and no storage metadata".to_string(),
+                            "FLAT column is not a FixedSizeList in IVF_HNSW_FLAT shard schema"
+                                .to_string(),
                         ));
                     }
                 };
@@ -1197,7 +1178,8 @@ pub async fn merge_partial_vector_auxiliary_files(
                     return Err(Error::index("Dimension mismatch across shards".to_string()));
                 }
                 if v2w_opt.is_none() {
-                    let w = init_writer_for_flat(object_store, &aux_out, d0, dt, fv).await?;
+                    let w = init_writer_for_flat(object_store, &aux_out, d0, &item_type, dt, fv)
+                        .await?;
                     v2w_opt = Some(w);
                 }
             }
@@ -1523,7 +1505,9 @@ pub async fn merge_partial_vector_auxiliary_files(
 mod tests {
     use super::*;
 
-    use arrow_array::{FixedSizeListArray, Float32Array, RecordBatch, UInt8Array, UInt64Array};
+    use arrow_array::{
+        FixedSizeListArray, Float32Array, Float64Array, RecordBatch, UInt8Array, UInt64Array,
+    };
     use arrow_schema::Field;
     use bytes::Bytes;
     use futures::StreamExt;
@@ -1614,6 +1598,71 @@ mod tests {
         Ok(total_rows)
     }
 
+    async fn write_flat_partial_aux_f64(
+        store: &ObjectStore,
+        aux_path: &Path,
+        dim: i32,
+        lengths: &[u32],
+        base_row_id: u64,
+        distance_type: DistanceType,
+    ) -> Result<usize> {
+        let arrow_schema = ArrowSchema::new(vec![
+            (*ROW_ID_FIELD).clone(),
+            Field::new(
+                crate::vector::flat::storage::FLAT_COLUMN,
+                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float64, true)), dim),
+                true,
+            ),
+        ]);
+
+        let writer = store.create(aux_path).await?;
+        let mut v2w = V2Writer::try_new(
+            writer,
+            lance_core::datatypes::Schema::try_from(&arrow_schema)?,
+            V2WriterOptions::default(),
+        )?;
+        v2w.add_schema_metadata(DISTANCE_TYPE_KEY, distance_type.to_string());
+
+        let ivf_meta = pb::Ivf {
+            centroids: Vec::new(),
+            offsets: Vec::new(),
+            lengths: lengths.to_vec(),
+            centroids_tensor: None,
+            loss: None,
+        };
+        let buf = Bytes::from(ivf_meta.encode_to_vec());
+        let pos = v2w.add_global_buffer(buf).await?;
+        v2w.add_schema_metadata(IVF_METADATA_KEY, pos.to_string());
+
+        let total_rows: usize = lengths.iter().map(|v| *v as usize).sum();
+        let mut row_ids = Vec::with_capacity(total_rows);
+        let mut values = Vec::with_capacity(total_rows * dim as usize);
+
+        let mut current_row_id = base_row_id;
+        for (pid, len) in lengths.iter().enumerate() {
+            for _ in 0..*len {
+                row_ids.push(current_row_id);
+                current_row_id += 1;
+                for d in 0..dim {
+                    values.push(pid as f64 + d as f64 * 0.01);
+                }
+            }
+        }
+
+        let row_id_arr = UInt64Array::from(row_ids);
+        let value_arr = Float64Array::from(values);
+        let fsl = FixedSizeListArray::try_new_from_values(value_arr, dim).unwrap();
+        let batch = RecordBatch::try_new(
+            Arc::new(arrow_schema),
+            vec![Arc::new(row_id_arr), Arc::new(fsl)],
+        )
+        .unwrap();
+
+        v2w.write_batch(&batch).await?;
+        v2w.finish().await?;
+        Ok(total_rows)
+    }
+
     #[tokio::test]
     async fn test_merge_ivf_flat_success_basic() {
         let object_store = ObjectStore::memory();
@@ -1831,6 +1880,64 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn test_merge_ivf_flat_preserves_float64_schema() {
+        let object_store = ObjectStore::memory();
+        let index_dir = Path::from("index/float64_uuid");
+
+        let partial0 = index_dir.child("partial_0");
+        let partial1 = index_dir.child("partial_1");
+        let aux0 = partial0.child(INDEX_AUXILIARY_FILE_NAME);
+        let aux1 = partial1.child(INDEX_AUXILIARY_FILE_NAME);
+
+        let lengths = vec![2_u32, 2_u32];
+        let dim = 3_i32;
+
+        write_flat_partial_aux_f64(&object_store, &aux0, dim, &lengths, 0, DistanceType::L2)
+            .await
+            .unwrap();
+        write_flat_partial_aux_f64(&object_store, &aux1, dim, &lengths, 100, DistanceType::L2)
+            .await
+            .unwrap();
+
+        merge_partial_vector_auxiliary_files(
+            &object_store,
+            &[aux0.clone(), aux1.clone()],
+            &index_dir,
+            Arc::new(RecordingProgress::default()),
+        )
+        .await
+        .unwrap();
+
+        let aux_out = index_dir.child(INDEX_AUXILIARY_FILE_NAME);
+        let sched = ScanScheduler::new(
+            Arc::new(object_store.clone()),
+            SchedulerConfig::max_bandwidth(&object_store),
+        );
+        let fh = sched
+            .open_file(&aux_out, &CachedFileSize::unknown())
+            .await
+            .unwrap();
+        let reader = V2Reader::try_open(
+            fh,
+            None,
+            Arc::default(),
+            &lance_core::cache::LanceCache::no_cache(),
+            V2ReaderOptions::default(),
+        )
+        .await
+        .unwrap();
+
+        let flat_field = reader
+            .schema()
+            .field(crate::vector::flat::storage::FLAT_COLUMN)
+            .unwrap();
+        let DataType::FixedSizeList(item, _) = flat_field.data_type() else {
+            panic!("flat column should be a fixed size list");
+        };
+        assert_eq!(item.data_type(), &DataType::Float64);
+    }
+
     #[allow(clippy::too_many_arguments)]
     async fn write_pq_partial_aux(
         store: &ObjectStore,