lance_file/versions/v2_2/
mod.rs1use 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
93pub 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
105pub 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
114pub 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
124pub fn create_lazy_writer(
126 object_writer: Box<dyn ObjectWriter>,
127 options: FileWriterOptions,
128) -> Writer {
129 Writer::new_lazy(object_writer, options)
130}
131
132pub 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
141pub fn encode_self_described_batch(batch: &EncodedBatch) -> Result<Bytes> {
143 writer::concat_lance_footer(batch, true)
144}
145
146pub fn encode_mini_batch(batch: &EncodedBatch) -> Result<Bytes> {
148 writer::concat_lance_footer(batch, false)
149}