lance_nd/
options.rs

1//! Lance encoding options for the columns of a table.
2//!
3//! Lance reads the encoding of a column from Arrow field metadata. The
4//! structural encoding is on the top-level field. Compression, byte stream
5//! split and run-length encoding are on each leaf field. Lance stores this
6//! metadata in the table schema, so it applies to all later appends.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use arrow::datatypes::{DataType, Field, Schema};
12use lance::Dataset;
13
14use crate::error::{Error, Result};
15use crate::format::check_format;
16
17// The keys of `lance_encoding::constants`.
18const COMPRESSION_KEY: &str = "lance-encoding:compression";
19const COMPRESSION_LEVEL_KEY: &str = "lance-encoding:compression-level";
20const BYTE_STREAM_SPLIT_KEY: &str = "lance-encoding:bss";
21const RLE_THRESHOLD_KEY: &str = "lance-encoding:rle-threshold";
22const STRUCTURAL_ENCODING_KEY: &str = "lance-encoding:structural-encoding";
23const KEYS: [&str; 5] = [
24    COMPRESSION_KEY,
25    COMPRESSION_LEVEL_KEY,
26    BYTE_STREAM_SPLIT_KEY,
27    RLE_THRESHOLD_KEY,
28    STRUCTURAL_ENCODING_KEY,
29];
30
31/// The prefix of Lance encoding keys in field metadata.
32pub(crate) const LANCE_ENCODING_PREFIX: &str = "lance-encoding:";
33
34/// The prefix of the table schema metadata keys that hold the column defaults.
35pub(crate) const COLUMN_DEFAULTS_PREFIX: &str = "lance_nd.column_defaults.";
36
37/// The general compression of the values of a column.
38#[derive(Debug, Clone, Copy, PartialEq, Eq)]
39pub enum Compression {
40    /// No general compression.
41    None,
42    /// LZ4 compression.
43    Lz4,
44    /// Zstandard compression. Lance selects the level when `level` is `None`.
45    Zstd {
46        /// The Zstandard compression level.
47        level: Option<i32>,
48    },
49}
50
51/// The page layout of a column.
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub enum StructuralEncoding {
54    /// Small blocks of values. Lance prefers it for narrow values.
55    MiniBlock,
56    /// Values zipped with their structure. Lance prefers it for wide values.
57    FullZip,
58}
59
60/// Byte stream split of floating point values before compression.
61#[derive(Debug, Clone, Copy, PartialEq, Eq)]
62pub enum ByteStreamSplit {
63    /// Never split.
64    Off,
65    /// Always split floating point values.
66    On,
67    /// Lance decides from the data.
68    Auto,
69}
70
71/// Lance encoding options of one column. `None` leaves the choice to Lance.
72#[derive(Debug, Clone, Default, PartialEq)]
73pub struct ColumnOptions {
74    /// The general compression of the values.
75    pub compression: Option<Compression>,
76    /// Byte stream split of floating point values.
77    pub byte_stream_split: Option<ByteStreamSplit>,
78    /// Lance uses run-length encoding when `runs < values * rle_threshold`.
79    pub rle_threshold: Option<f64>,
80    /// The page layout of the column.
81    pub structural_encoding: Option<StructuralEncoding>,
82}
83
84impl ColumnOptions {
85    /// Set the general compression.
86    pub fn with_compression(mut self, compression: Compression) -> Self {
87        self.compression = Some(compression);
88        self
89    }
90
91    /// Set the byte stream split mode.
92    pub fn with_byte_stream_split(mut self, mode: ByteStreamSplit) -> Self {
93        self.byte_stream_split = Some(mode);
94        self
95    }
96
97    /// Set the run-length encoding threshold, from 0.0 to 1.0.
98    pub fn with_rle_threshold(mut self, threshold: f64) -> Self {
99        self.rle_threshold = Some(threshold);
100        self
101    }
102
103    /// Set the structural encoding.
104    pub fn with_structural_encoding(mut self, encoding: StructuralEncoding) -> Self {
105        self.structural_encoding = Some(encoding);
106        self
107    }
108
109    /// The options of `self`, with the options of `defaults` where `self` has none.
110    pub(crate) fn or(&self, defaults: &ColumnOptions) -> ColumnOptions {
111        ColumnOptions {
112            compression: self.compression.or(defaults.compression),
113            byte_stream_split: self.byte_stream_split.or(defaults.byte_stream_split),
114            rle_threshold: self.rle_threshold.or(defaults.rle_threshold),
115            structural_encoding: self.structural_encoding.or(defaults.structural_encoding),
116        }
117    }
118
119    /// The Lance metadata for the leaf fields of a column.
120    pub(crate) fn leaf_metadata(&self) -> HashMap<String, String> {
121        let mut metadata = HashMap::new();
122        match self.compression {
123            Some(Compression::None) => insert(&mut metadata, COMPRESSION_KEY, "none"),
124            Some(Compression::Lz4) => insert(&mut metadata, COMPRESSION_KEY, "lz4"),
125            Some(Compression::Zstd { level }) => {
126                insert(&mut metadata, COMPRESSION_KEY, "zstd");
127                if let Some(level) = level {
128                    insert(&mut metadata, COMPRESSION_LEVEL_KEY, &level.to_string());
129                }
130            }
131            None => {}
132        }
133        let mode = self.byte_stream_split.map(|mode| match mode {
134            ByteStreamSplit::Off => "off",
135            ByteStreamSplit::On => "on",
136            ByteStreamSplit::Auto => "auto",
137        });
138        if let Some(mode) = mode {
139            insert(&mut metadata, BYTE_STREAM_SPLIT_KEY, mode);
140        }
141        if let Some(threshold) = self.rle_threshold {
142            insert(&mut metadata, RLE_THRESHOLD_KEY, &threshold.to_string());
143        }
144        metadata
145    }
146
147    /// The Lance metadata for the top-level field of a column.
148    pub(crate) fn root_metadata(&self) -> HashMap<String, String> {
149        let mut metadata = HashMap::new();
150        let encoding = self.structural_encoding.map(|encoding| match encoding {
151            StructuralEncoding::MiniBlock => "miniblock",
152            StructuralEncoding::FullZip => "fullzip",
153        });
154        if let Some(encoding) = encoding {
155            insert(&mut metadata, STRUCTURAL_ENCODING_KEY, encoding);
156        }
157        metadata
158    }
159
160    /// The options as table schema metadata, see [`COLUMN_DEFAULTS_PREFIX`].
161    pub(crate) fn to_table_metadata(&self) -> HashMap<String, String> {
162        self.leaf_metadata()
163            .into_iter()
164            .chain(self.root_metadata())
165            .map(|(key, value)| (format!("{COLUMN_DEFAULTS_PREFIX}{key}"), value))
166            .collect()
167    }
168
169    /// Read the options from table schema metadata. Other keys are ignored.
170    pub(crate) fn from_table_metadata(metadata: &HashMap<String, String>) -> Result<Self> {
171        let get = |key: &str| metadata.get(&format!("{COLUMN_DEFAULTS_PREFIX}{key}"));
172        let invalid = |key: &str, value: &str| {
173            Error::InvalidTable(format!(
174                "column default {key} has the invalid value '{value}'"
175            ))
176        };
177        let level = get(COMPRESSION_LEVEL_KEY)
178            .map(|v| v.parse().map_err(|_| invalid(COMPRESSION_LEVEL_KEY, v)))
179            .transpose()?;
180        let compression = get(COMPRESSION_KEY)
181            .map(|v| match v.as_str() {
182                "none" => Ok(Compression::None),
183                "lz4" => Ok(Compression::Lz4),
184                "zstd" => Ok(Compression::Zstd { level }),
185                _ => Err(invalid(COMPRESSION_KEY, v)),
186            })
187            .transpose()?;
188        let byte_stream_split = get(BYTE_STREAM_SPLIT_KEY)
189            .map(|v| match v.as_str() {
190                "off" => Ok(ByteStreamSplit::Off),
191                "on" => Ok(ByteStreamSplit::On),
192                "auto" => Ok(ByteStreamSplit::Auto),
193                _ => Err(invalid(BYTE_STREAM_SPLIT_KEY, v)),
194            })
195            .transpose()?;
196        let rle_threshold = get(RLE_THRESHOLD_KEY)
197            .map(|v| v.parse().map_err(|_| invalid(RLE_THRESHOLD_KEY, v)))
198            .transpose()?;
199        let structural_encoding = get(STRUCTURAL_ENCODING_KEY)
200            .map(|v| match v.as_str() {
201                "miniblock" => Ok(StructuralEncoding::MiniBlock),
202                "fullzip" => Ok(StructuralEncoding::FullZip),
203                _ => Err(invalid(STRUCTURAL_ENCODING_KEY, v)),
204            })
205            .transpose()?;
206        Ok(Self {
207            compression,
208            byte_stream_split,
209            rle_threshold,
210            structural_encoding,
211        })
212    }
213}
214
215fn insert(metadata: &mut HashMap<String, String>, key: &str, value: &str) {
216    metadata.insert(key.to_string(), value.to_string());
217}
218
219/// The column defaults of the table. New columns get them.
220pub fn column_defaults(dataset: &Dataset) -> Result<ColumnOptions> {
221    check_format(dataset)?;
222    ColumnOptions::from_table_metadata(&dataset.schema().metadata)
223}
224
225/// Replace the column defaults of the table. Existing columns keep their encoding.
226#[tracing::instrument(skip_all, fields(uri = %dataset.uri()), err)]
227pub async fn set_column_defaults(dataset: &mut Dataset, defaults: &ColumnOptions) -> Result<()> {
228    check_format(dataset)?;
229    let mut values = defaults.to_table_metadata();
230    let entries: Vec<(String, Option<String>)> = KEYS
231        .iter()
232        .map(|key| {
233            let key = format!("{COLUMN_DEFAULTS_PREFIX}{key}");
234            let value = values.remove(&key);
235            (key, value)
236        })
237        .collect();
238    dataset.update_schema_metadata(entries).await?;
239    Ok(())
240}
241
242/// `schema` with the encoding metadata of each column. The options of a
243/// column are its entry in `columns`, then `defaults`.
244pub(crate) fn with_column_options(
245    schema: &Schema,
246    defaults: &ColumnOptions,
247    columns: &HashMap<String, ColumnOptions>,
248) -> Schema {
249    let fields: Vec<Field> = schema
250        .fields()
251        .iter()
252        .map(|field| {
253            let options = match columns.get(field.name()) {
254                Some(options) => options.or(defaults),
255                None => defaults.clone(),
256            };
257            let leaf = options.leaf_metadata();
258            let mut field = with_leaf_metadata(field, &leaf);
259            let mut metadata = field.metadata().clone();
260            metadata.extend(options.root_metadata());
261            field.set_metadata(metadata);
262            field
263        })
264        .collect();
265    Schema::new_with_metadata(fields, schema.metadata().clone())
266}
267
268/// `field` with `leaf` added to the metadata of each leaf field.
269fn with_leaf_metadata(field: &Field, leaf: &HashMap<String, String>) -> Field {
270    let data_type = match field.data_type() {
271        DataType::Struct(children) => DataType::Struct(
272            children
273                .iter()
274                .map(|child| with_leaf_metadata(child, leaf))
275                .collect(),
276        ),
277        DataType::List(item) => DataType::List(Arc::new(with_leaf_metadata(item, leaf))),
278        _ => {
279            let mut metadata = field.metadata().clone();
280            metadata.extend(leaf.clone());
281            return field.clone().with_metadata(metadata);
282        }
283    };
284    field.clone().with_data_type(data_type)
285}
286
287/// `field` without the Lance encoding keys, also in nested fields.
288pub(crate) fn without_encoding_metadata(field: &Field) -> Field {
289    let data_type = match field.data_type() {
290        DataType::Struct(children) => DataType::Struct(
291            children
292                .iter()
293                .map(|child| without_encoding_metadata(child))
294                .collect(),
295        ),
296        DataType::List(item) => DataType::List(Arc::new(without_encoding_metadata(item))),
297        data_type => data_type.clone(),
298    };
299    let metadata = field
300        .metadata()
301        .iter()
302        .filter(|(key, _)| !key.starts_with(LANCE_ENCODING_PREFIX))
303        .map(|(key, value)| (key.clone(), value.clone()))
304        .collect();
305    field
306        .clone()
307        .with_data_type(data_type)
308        .with_metadata(metadata)
309}
lance_nd/options.rs

lance_nd/
options.rs
