lance_index/vector/
flat.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Flat Vector Index.
5//!
6
7use std::sync::Arc;
8
9use arrow::{array::AsArray, buffer::NullBuffer};
10use arrow_array::{Array, ArrayRef, Float32Array, RecordBatch, make_array};
11use arrow_schema::{DataType, Field as ArrowField};
12use lance_arrow::*;
13use lance_core::{Error, ROW_ID, Result};
14use lance_linalg::distance::{DistanceType, multivec_distance};
15use tracing::instrument;
16
17use super::DIST_COL;
18
19pub mod index;
20pub mod storage;
21pub mod transform;
22
23fn distance_field() -> ArrowField {
24    ArrowField::new(DIST_COL, DataType::Float32, true)
25}
26
27/// Get a column from a RecordBatch, supporting nested field paths.
28///
29/// This function handles:
30/// - Simple column names: "column"
31/// - Nested paths: "parent.child" or "parent.child.grandchild"
32/// - Backtick-escaped field names: "parent.`field.with.dots`"
33fn get_column_from_batch(batch: &RecordBatch, column: &str) -> Result<ArrayRef> {
34    // Try to get the column directly first (fast path for simple columns)
35    if let Some(col) = batch.column_by_name(column) {
36        return Ok(col.clone());
37    }
38
39    // Parse the field path using Lance's field path parsing logic
40    // This properly handles backtick-escaped field names
41    let parts = lance_core::datatypes::parse_field_path(column)
42        .map_err(|e| Error::schema(format!("Failed to parse field path '{}': {}", column, e)))?;
43
44    if parts.is_empty() {
45        return Err(Error::schema(format!(
46            "Invalid empty field path: {}",
47            column
48        )));
49    }
50
51    // Get the root column
52    let mut current_array: ArrayRef = batch
53        .column_by_name(&parts[0])
54        .ok_or_else(|| {
55            Error::schema(format!(
56                "Column '{}' does not exist in batch (looking for root field '{}')",
57                column, parts[0]
58            ))
59        })?
60        .clone();
61
62    // Navigate through nested struct fields
63    for part in &parts[1..] {
64        let struct_array = current_array
65            .as_any()
66            .downcast_ref::<arrow_array::StructArray>()
67            .ok_or_else(|| {
68                Error::schema(format!(
69                    "Cannot access nested field '{}' in column '{}': parent is not a struct",
70                    part, column
71                ))
72            })?;
73
74        current_array = struct_array
75            .column_by_name(part)
76            .ok_or_else(|| {
77                Error::schema(format!(
78                    "Nested field '{}' does not exist in column '{}'",
79                    part, column
80                ))
81            })?
82            .clone();
83    }
84
85    Ok(current_array)
86}
87
88#[instrument(level = "debug", skip_all)]
89pub async fn compute_distance(
90    key: ArrayRef,
91    dt: DistanceType,
92    column: &str,
93    mut batch: RecordBatch,
94) -> Result<RecordBatch> {
95    if batch.column_by_name(DIST_COL).is_some() {
96        // Ignore the distance calculated from inner vector index.
97        batch = batch.drop_column(DIST_COL)?;
98    }
99
100    let vectors = get_column_from_batch(&batch, column)?;
101
102    let validity_buffer = if let Some(rowids) = batch.column_by_name(ROW_ID) {
103        NullBuffer::union(rowids.nulls(), vectors.nulls())
104    } else {
105        vectors.nulls().cloned()
106    };
107
108    tokio::task::spawn_blocking(move || {
109        // A selection vector may have been applied to _rowid column, so we need to
110        // push that onto vectors if possible.
111
112        let vectors = vectors
113            .into_data()
114            .into_builder()
115            .null_bit_buffer(validity_buffer.map(|b| b.buffer().clone()))
116            .build()
117            .map(make_array)?;
118        let distances = match vectors.data_type() {
119            DataType::FixedSizeList(_, _) => {
120                let vectors = vectors.as_fixed_size_list();
121                dt.arrow_batch_func()(key.as_ref(), vectors)? as ArrayRef
122            }
123            DataType::List(_) => {
124                let vectors = vectors.as_list();
125                let dists = multivec_distance(key.as_ref(), vectors, dt)?;
126                Arc::new(Float32Array::from(dists))
127            }
128            _ => {
129                unreachable!()
130            }
131        };
132
133        batch
134            .try_with_column(distance_field(), distances)
135            .map_err(|e| Error::execution(format!("Failed to adding distance column: {}", e)))
136    })
137    .await?
138}
lance_index/vector/flat.rs

lance_index/vector/
flat.rs
