Skip to main content

lance_file/versions/v2_3/
mod.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Lance v2.3 encoding composition.
5
6use std::sync::{
7    Arc,
8    atomic::{AtomicBool, Ordering},
9};
10
11use bytes::Bytes;
12use lance_core::{
13    Error, Result,
14    datatypes::{Field, Schema},
15};
16use lance_encoding::{
17    compression_config::CompressionParams,
18    encoder::{
19        ColumnIndexSequence, EncodedBatch, FieldEncoder, FieldEncodingContext,
20        FieldEncodingStrategy,
21        structural::{
22            PrimitiveFieldEncoding, PrimitivePageEncoding, try_create_binary_blob, try_create_list,
23            try_create_map, try_create_struct, try_create_structural_blob,
24            try_create_structural_fixed_size_list,
25        },
26    },
27};
28use lance_io::traits::Writer as ObjectWriter;
29
30use crate::writer::FileWriterOptions;
31
32mod compression;
33mod writer;
34
35pub use writer::Writer;
36
37static WARNED_ON_UNSTABLE_FORMAT: AtomicBool = AtomicBool::new(false);
38
39#[derive(Debug)]
40struct FieldStrategy {
41    primitive: PrimitiveFieldEncoding,
42}
43
44impl FieldEncodingStrategy for FieldStrategy {
45    fn create_field_encoder(
46        &self,
47        field: &Field,
48        column_index: &mut ColumnIndexSequence,
49        context: &FieldEncodingContext<'_>,
50    ) -> Result<Box<dyn FieldEncoder>> {
51        if let Some(encoder) =
52            try_create_binary_blob(&self.primitive, field, column_index, context)?
53        {
54            return Ok(encoder);
55        }
56        if let Some(encoder) =
57            try_create_structural_blob(&self.primitive, field, column_index, context)?
58        {
59            return Ok(encoder);
60        }
61        if field.is_blob() {
62            return Err(Error::invalid_input_source(
63                format!(
64                    "Blob encoding is not available for field '{}' with data type {}",
65                    field.name,
66                    field.data_type()
67                )
68                .into(),
69            ));
70        }
71        if let Some(encoder) = try_create_map(field, column_index, context)? {
72            return Ok(encoder);
73        }
74        if let Some(encoder) = try_create_structural_fixed_size_list(field, column_index, context)?
75        {
76            return Ok(encoder);
77        }
78        if let Some(encoder) = self.primitive.try_create(field, column_index, context)? {
79            return Ok(encoder);
80        }
81        if let Some(encoder) = try_create_list(field, column_index, context)? {
82            return Ok(encoder);
83        }
84        if let Some(encoder) = try_create_struct(field, column_index, context)? {
85            return Ok(encoder);
86        }
87        Err(Error::not_supported_source(
88            format!(
89                "Lance v2.3 has no field encoding for '{}' with data type {}",
90                field.name,
91                field.data_type()
92            )
93            .into(),
94        ))
95    }
96}
97
98/// Compose the v2.3 field encoding mechanisms.
99pub fn encoding_strategy(params: CompressionParams) -> Arc<dyn FieldEncodingStrategy> {
100    let compression = Arc::new(compression::Strategy::new(params));
101    Arc::new(FieldStrategy {
102        primitive: PrimitiveFieldEncoding::new([
103            PrimitivePageEncoding::sparse(compression.clone()),
104            PrimitivePageEncoding::constant(),
105            PrimitivePageEncoding::dense_u32(compression),
106        ]),
107    })
108}
109
110fn warn_unstable_format() {
111    if WARNED_ON_UNSTABLE_FORMAT
112        .compare_exchange(false, true, Ordering::Relaxed, Ordering::Relaxed)
113        .is_ok()
114    {
115        log::warn!(
116            "You have requested an unstable format version.  Files written with this format version may not be readable in the future!  This is a development feature and should only be used for experimentation and never for production data."
117        );
118    }
119}
120
121/// Create a v2.3 writer with an explicit schema.
122pub fn create_writer(
123    object_writer: Box<dyn ObjectWriter>,
124    schema: Schema,
125    options: FileWriterOptions,
126) -> Result<Writer> {
127    warn_unstable_format();
128    Writer::try_new(object_writer, schema, options)
129}
130
131/// Create a v2.3 writer with explicit compression tuning.
132pub fn create_writer_with_compression(
133    object_writer: Box<dyn ObjectWriter>,
134    schema: Schema,
135    options: FileWriterOptions,
136    compression: CompressionParams,
137) -> Result<Writer> {
138    warn_unstable_format();
139    Writer::try_new_with_compression(object_writer, schema, options, compression)
140}
141
142/// Create a v2.3 writer whose schema is inferred from the first batch.
143pub fn create_lazy_writer(
144    object_writer: Box<dyn ObjectWriter>,
145    options: FileWriterOptions,
146) -> Writer {
147    warn_unstable_format();
148    Writer::new_lazy(object_writer, options)
149}
150
151/// Create a lazy v2.3 writer with explicit compression tuning.
152pub fn create_lazy_writer_with_compression(
153    object_writer: Box<dyn ObjectWriter>,
154    options: FileWriterOptions,
155    compression: CompressionParams,
156) -> Writer {
157    warn_unstable_format();
158    Writer::new_lazy_with_compression(object_writer, options, compression)
159}
160
161/// Encode a self-described v2.3 batch.
162pub fn encode_self_described_batch(batch: &EncodedBatch) -> Result<Bytes> {
163    writer::concat_lance_footer(batch, true)
164}
165
166/// Encode a mini-lance v2.3 batch.
167pub fn encode_mini_batch(batch: &EncodedBatch) -> Result<Bytes> {
168    writer::concat_lance_footer(batch, false)
169}