lance_file/versions/v1/
encoding.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Lance v1 column encoding grammar.
5//!
6//! These codecs own the v1 wire layout, including schema dictionary payloads.
7//! They are intentionally not exposed through the version-free I/O layer.
8
9use arrow_array::{
10    Array, ArrayRef,
11    types::{BinaryType, LargeBinaryType, LargeUtf8Type, Utf8Type},
12};
13use arrow_schema::DataType;
14use async_recursion::async_recursion;
15
16pub mod binary;
17pub mod dictionary;
18pub mod plain;
19
20use lance_arrow::DataTypeExt;
21use lance_core::{
22    Error, Result,
23    datatypes::{Field, Schema},
24};
25use lance_io::{
26    ReadBatchParams,
27    traits::{Reader, Writer},
28};
29
30use self::{
31    binary::{BinaryDecoder, BinaryEncoder},
32    plain::{PlainDecoder, PlainEncoder},
33};
34
35/// Decode a binary-like array from a v1 values region.
36pub async fn read_binary_array(
37    reader: &dyn Reader,
38    data_type: &DataType,
39    nullable: bool,
40    position: usize,
41    length: usize,
42    params: impl Into<ReadBatchParams>,
43) -> Result<ArrayRef> {
44    use arrow_schema::DataType::*;
45
46    let params = params.into();
47    match data_type {
48        Utf8 => {
49            BinaryDecoder::<Utf8Type>::new(reader, position, length, nullable)
50                .get(params)
51                .await
52        }
53        Binary => {
54            BinaryDecoder::<BinaryType>::new(reader, position, length, nullable)
55                .get(params)
56                .await
57        }
58        LargeUtf8 => {
59            BinaryDecoder::<LargeUtf8Type>::new(reader, position, length, nullable)
60                .get(params)
61                .await
62        }
63        LargeBinary => {
64            BinaryDecoder::<LargeBinaryType>::new(reader, position, length, nullable)
65                .get(params)
66                .await
67        }
68        _ => Err(lance_core::Error::invalid_input(format!(
69            "unsupported v1 binary data type: {data_type}"
70        ))),
71    }
72}
73
74/// Decode a fixed-stride array from a v1 values region.
75pub async fn read_fixed_stride_array(
76    reader: &dyn Reader,
77    data_type: &DataType,
78    position: usize,
79    length: usize,
80    params: impl Into<ReadBatchParams>,
81) -> Result<ArrayRef> {
82    if !lance_arrow::DataTypeExt::is_fixed_stride(data_type) {
83        return Err(lance_core::Error::schema(format!(
84            "{data_type} is not a fixed stride type"
85        )));
86    }
87    PlainDecoder::new(reader, data_type, position, length)?
88        .get(params.into())
89        .await
90}
91
92/// Persist every schema dictionary using the v1 value codecs.
93pub async fn write_schema_dictionaries(writer: &mut dyn Writer, schema: &mut Schema) -> Result<()> {
94    let max_field_id = schema.max_field_id().unwrap_or(-1);
95    for field_id in 0..=max_field_id {
96        let Some(field) = schema.mut_field_by_id(field_id) else {
97            continue;
98        };
99        if !field.data_type().is_dictionary() {
100            continue;
101        }
102
103        let dict_info = field.dictionary.as_mut().ok_or_else(|| {
104            Error::io(format!(
105                "v1 dictionary field '{}' is missing dictionary metadata",
106                field.name
107            ))
108        })?;
109        let values = dict_info.values.as_ref().ok_or_else(|| {
110            Error::invalid_input(format!(
111                "v1 dictionary field '{}' is missing dictionary values",
112                field.name
113            ))
114        })?;
115
116        let data_type = values.data_type();
117        let position = if data_type.is_numeric() {
118            PlainEncoder::new(writer, data_type)
119                .encode(&[values])
120                .await?
121        } else if data_type.is_binary_like() {
122            BinaryEncoder::new(writer).encode(&[values]).await?
123        } else {
124            return Err(Error::schema(format!(
125                "v1 dictionary values do not support data type {data_type}"
126            )));
127        };
128        dict_info.offset = position;
129        dict_info.length = values.len();
130    }
131    Ok(())
132}
133
134#[async_recursion]
135async fn populate_field_dictionary(field: &mut Field, reader: &dyn Reader) -> Result<()> {
136    if let DataType::Dictionary(_, value_type) = field.data_type() {
137        let dict_info = field.dictionary.as_mut().ok_or_else(|| {
138            Error::io(format!(
139                "v1 dictionary field '{}' is missing dictionary metadata",
140                field.name
141            ))
142        })?;
143        let values = if value_type.is_binary_like() {
144            read_binary_array(
145                reader,
146                value_type.as_ref(),
147                true,
148                dict_info.offset,
149                dict_info.length,
150                ..,
151            )
152            .await?
153        } else if matches!(
154            value_type.as_ref(),
155            DataType::Int8
156                | DataType::Int16
157                | DataType::Int32
158                | DataType::Int64
159                | DataType::UInt8
160                | DataType::UInt16
161                | DataType::UInt32
162                | DataType::UInt64
163        ) {
164            read_fixed_stride_array(
165                reader,
166                value_type.as_ref(),
167                dict_info.offset,
168                dict_info.length,
169                ..,
170            )
171            .await?
172        } else {
173            return Err(Error::schema(format!(
174                "v1 dictionary values do not support data type {value_type}"
175            )));
176        };
177        dict_info.values = Some(values);
178    } else {
179        for child in &mut field.children {
180            populate_field_dictionary(child, reader).await?;
181        }
182    }
183    Ok(())
184}
185
186/// Load every persisted v1 schema dictionary into its in-memory field.
187pub async fn populate_schema_dictionaries(schema: &mut Schema, reader: &dyn Reader) -> Result<()> {
188    for field in &mut schema.fields {
189        populate_field_dictionary(field, reader).await?;
190    }
191    Ok(())
192}
lance_file/versions/v1/encoding.rs

lance_file/versions/v1/
encoding.rs
