lance_index/vector/flat/
transform.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use arrow_array::RecordBatch;
5use arrow_schema::Field;
6use lance_arrow::RecordBatchExt;
7use lance_core::Error;
8use snafu::location;
9use tracing::instrument;
10
11use crate::vector::transform::Transformer;
12
13use super::storage::FLAT_COLUMN;
14
15#[derive(Debug)]
16pub struct FlatTransformer {
17    input_column: String,
18}
19
20impl FlatTransformer {
21    pub fn new(input_column: impl AsRef<str>) -> Self {
22        Self {
23            input_column: input_column.as_ref().to_owned(),
24        }
25    }
26}
27
28impl Transformer for FlatTransformer {
29    #[instrument(name = "FlatTransformer::transform", level = "debug", skip_all)]
30    fn transform(&self, batch: &RecordBatch) -> crate::Result<RecordBatch> {
31        let input_arr = batch
32            .column_by_name(&self.input_column)
33            .ok_or(Error::Index {
34                message: format!(
35                    "FlatTransform: column {} not found in batch",
36                    self.input_column
37                ),
38                location: location!(),
39            })?;
40        let field = Field::new(
41            FLAT_COLUMN,
42            input_arr.data_type().clone(),
43            input_arr.is_nullable(),
44        );
45        // rename the column to FLAT_COLUMN
46        let batch = batch
47            .drop_column(&self.input_column)?
48            .try_with_column(field, input_arr.clone())?;
49        Ok(batch)
50    }
51}
lance_index/vector/flat/transform.rs

lance_index/vector/flat/
transform.rs
