fix

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

rust/lance-index/src/vector/ivf.rs

@@ -25,7 +25,7 @@ use super::quantizer::Quantization;
 use super::residual::ResidualTransform;
 use super::sq::transform::SQTransformer;
 use super::sq::ScalarQuantizer;
-use super::transform::KeepFiniteVectors;
+use super::transform::{DropColumn, KeepFiniteVectors};
 use super::{quantizer::Quantizer, residual::compute_residual};
 use super::{PART_ID_COLUMN, PQ_CODE_COLUMN, SQ_CODE_COLUMN};
 
@@ -200,6 +200,7 @@ impl IvfTransformer {
             vector_column,
             PQ_CODE_COLUMN,
         )));
+
         Self::new(centroids, distance_type, transforms)
     }
 

rust/lance-index/src/vector/v3/shuffler.rs

@@ -117,6 +117,7 @@ impl Shuffler for IvfShuffler {
                         .column_by_name(PART_ID_COLUMN)
                         .expect("Partition ID column not found")
                         .as_primitive();
+                    let batch = batch.drop_column(PART_ID_COLUMN)?;
 
                     let mut partition_buffers =
                         (0..num_partitions).map(|_| Vec::new()).collect::<Vec<_>>();

rust/lance/src/index/vector/builder.rs

@@ -500,6 +500,18 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
                 if num_rows == 0 {
                     return Ok((0, 0));
                 }
+                for batch in batches.iter() {
+                    if batch.schema() != batches[0].schema() {
+                        return Err(Error::invalid_input(
+                            format!(
+                                "batch schema mismatch: expected {:?}, got {:?}",
+                                batches[0].schema(),
+                                batch.schema()
+                            ),
+                            location!(),
+                        ));
+                    }
+                }
                 let batch = arrow::compute::concat_batches(&batches[0].schema(), batches.iter())?;
 
                 Self::build_partition(
@@ -602,7 +614,7 @@ impl<S: IvfSubIndex + 'static, Q: Quantization + 'static> IvfIndexBuilder<S, Q>
             //         Ok(batch.try_with_column(PART_ID_FIELD.clone(), Arc::new(part_ids))?)
             //     })
             //     .collect::<Result<Vec<_>>>()?;
-            batches.extend(part_batches);
+            batches.extend(part_batches.map(|b| b.with_metadata(HashMap::new()).unwrap()));
         }
 
         if reader.partition_size(part_id)? > 0 {