lance_file/versions/v2_3/
mod.rs1use 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
98pub 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
121pub 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
131pub 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
142pub 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
151pub 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
161pub fn encode_self_described_batch(batch: &EncodedBatch) -> Result<Bytes> {
163 writer::concat_lance_footer(batch, true)
164}
165
166pub fn encode_mini_batch(batch: &EncodedBatch) -> Result<Bytes> {
168 writer::concat_lance_footer(batch, false)
169}