Skip to main content

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 tracing::instrument;
9
10use crate::vector::transform::Transformer;
11
12use super::storage::FLAT_COLUMN;
13
14#[derive(Debug)]
15pub struct FlatTransformer {
16    input_column: String,
17}
18
19impl FlatTransformer {
20    pub fn new(input_column: impl AsRef<str>) -> Self {
21        Self {
22            input_column: input_column.as_ref().to_owned(),
23        }
24    }
25}
26
27impl Transformer for FlatTransformer {
28    #[instrument(name = "FlatTransformer::transform", level = "debug", skip_all)]
29    fn transform(&self, batch: &RecordBatch) -> crate::Result<RecordBatch> {
30        let input_arr = batch
31            .column_by_name(&self.input_column)
32            .ok_or(Error::index(format!(
33                "FlatTransform: column {} not found in batch",
34                self.input_column
35            )))?;
36        let field = Field::new(
37            FLAT_COLUMN,
38            input_arr.data_type().clone(),
39            input_arr.is_nullable(),
40        );
41        // rename the column to FLAT_COLUMN
42        let batch = batch
43            .drop_column(&self.input_column)?
44            .try_with_column(field, input_arr.clone())?;
45        Ok(batch)
46    }
47}