Skip to main content

lance_file/versions/v2_2/
mod.rs

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