lance_file/
writer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use arrow_array::{ArrayRef, RecordBatch};
7use bytes::Bytes;
8use lance_core::{Error, Result, datatypes::Schema};
9use lance_encoding::decoder::{ColumnInfo, PageEncoding};
10use lance_io::object_store::ObjectStore;
11use object_store::path::Path;
12use prost::Message;
13use prost_types::Any;
14
15use crate::{format::pbfile, versions};
16
17mod field_types;
18pub(crate) mod structural;
19
20pub(crate) use field_types::ExpectedTypes;
21pub use field_types::FieldTypeMismatch;
22
23/// Page buffers in current Lance files are aligned to 64 bytes.
24pub(crate) const PAGE_BUFFER_ALIGNMENT: usize = 64;
25pub(crate) const ENV_LANCE_FILE_WRITER_MAX_PAGE_BYTES: &str = "LANCE_FILE_WRITER_MAX_PAGE_BYTES";
26
27/// Summary of a completed Lance file write.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct FileWriteSummary {
30    /// The number of rows written to the file.
31    pub num_rows: u64,
32    /// The final size of the file in bytes.
33    pub size_bytes: u64,
34}
35
36/// Runtime options shared by all current-format writers.
37///
38/// These options control buffering and execution only. Select the file grammar
39/// by constructing a writer under [`crate::versions`].
40#[derive(Debug, Clone, Default)]
41pub struct FileWriterOptions {
42    /// How many bytes to use for buffering column data.
43    ///
44    /// The budget is divided evenly across columns. The default is 8 MiB per
45    /// column.
46    pub data_cache_bytes: Option<u64>,
47    /// A best-effort maximum encoded page size.
48    pub max_page_bytes: Option<u64>,
49    /// Keep input arrays instead of copying buffered slices.
50    ///
51    /// Do not enable this for arrays arriving through the Arrow C data
52    /// interface because a small child array can keep an entire batch alive.
53    pub keep_original_array: Option<bool>,
54}
55
56impl FileWriterOptions {
57    pub(crate) fn validate(&self) -> Result<()> {
58        if self.max_page_bytes == Some(0) {
59            return Err(Error::invalid_input(
60                "max_page_bytes must be greater than 0, got 0",
61            ));
62        }
63        Ok(())
64    }
65}
66
67/// A type-erased current-format file writer.
68///
69/// This enum exists for callers that select a concrete file version at
70/// runtime. Each variant owns the complete implementation for exactly one file
71/// grammar; this type only forwards operations without adding format policy.
72pub enum FileWriter {
73    V2_0(Box<versions::v2_0::Writer>),
74    V2_1(Box<versions::v2_1::Writer>),
75    V2_2(Box<versions::v2_2::Writer>),
76    V2_3(Box<versions::v2_3::Writer>),
77}
78
79fn column_info_to_metadata(column: &ColumnInfo) -> Result<pbfile::ColumnMetadata> {
80    let pages = column
81        .page_infos
82        .iter()
83        .map(|page| {
84            let encoding = match &page.encoding {
85                PageEncoding::Legacy(encoding) => Any::from_msg(encoding)?.encode_to_vec(),
86                PageEncoding::Structural(encoding) => Any::from_msg(encoding)?.encode_to_vec(),
87            };
88            let (buffer_offsets, buffer_sizes) =
89                page.buffer_offsets_and_sizes.iter().copied().unzip();
90            Ok(pbfile::column_metadata::Page {
91                buffer_offsets,
92                buffer_sizes,
93                encoding: Some(pbfile::Encoding {
94                    location: Some(pbfile::encoding::Location::Direct(pbfile::DirectEncoding {
95                        encoding,
96                    })),
97                }),
98                length: page.num_rows,
99                priority: page.priority,
100            })
101        })
102        .collect::<Result<Vec<_>>>()?;
103    let (buffer_offsets, buffer_sizes) = column.buffer_offsets_and_sizes.iter().copied().unzip();
104    let encoding = Any::from_msg(&column.encoding)?.encode_to_vec();
105    Ok(pbfile::ColumnMetadata {
106        pages,
107        buffer_offsets,
108        buffer_sizes,
109        encoding: Some(pbfile::Encoding {
110            location: Some(pbfile::encoding::Location::Direct(pbfile::DirectEncoding {
111                encoding,
112            })),
113        }),
114    })
115}
116
117impl From<versions::v2_0::Writer> for FileWriter {
118    fn from(writer: versions::v2_0::Writer) -> Self {
119        Self::V2_0(Box::new(writer))
120    }
121}
122
123impl From<versions::v2_1::Writer> for FileWriter {
124    fn from(writer: versions::v2_1::Writer) -> Self {
125        Self::V2_1(Box::new(writer))
126    }
127}
128
129impl From<versions::v2_2::Writer> for FileWriter {
130    fn from(writer: versions::v2_2::Writer) -> Self {
131        Self::V2_2(Box::new(writer))
132    }
133}
134
135impl From<versions::v2_3::Writer> for FileWriter {
136    fn from(writer: versions::v2_3::Writer) -> Self {
137        Self::V2_3(Box::new(writer))
138    }
139}
140
141impl FileWriter {
142    /// Spill page metadata to a sidecar file.
143    pub fn with_page_metadata_spill(self, object_store: Arc<ObjectStore>, path: Path) -> Self {
144        match self {
145            Self::V2_0(writer) => Self::V2_0(Box::new(
146                (*writer).with_page_metadata_spill(object_store, path),
147            )),
148            Self::V2_1(writer) => Self::V2_1(Box::new(
149                (*writer).with_page_metadata_spill(object_store, path),
150            )),
151            Self::V2_2(writer) => Self::V2_2(Box::new(
152                (*writer).with_page_metadata_spill(object_store, path),
153            )),
154            Self::V2_3(writer) => Self::V2_3(Box::new(
155                (*writer).with_page_metadata_spill(object_store, path),
156            )),
157        }
158    }
159
160    /// Schedule batches of data to be written to the file.
161    pub async fn write_batches(
162        &mut self,
163        batches: impl Iterator<Item = &RecordBatch>,
164    ) -> Result<()> {
165        for batch in batches {
166            self.write_batch(batch).await?;
167        }
168        Ok(())
169    }
170
171    /// Schedule a batch of data to be written to the file.
172    pub async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
173        match self {
174            Self::V2_0(writer) => writer.write_batch(batch).await,
175            Self::V2_1(writer) => writer.write_batch(batch).await,
176            Self::V2_2(writer) => writer.write_batch(batch).await,
177            Self::V2_3(writer) => writer.write_batch(batch).await,
178        }
179    }
180
181    /// Write one top-level column.
182    pub async fn write_column(&mut self, column_index: usize, array: ArrayRef) -> Result<()> {
183        match self {
184            Self::V2_0(writer) => writer.write_column(column_index, array).await,
185            Self::V2_1(writer) => writer.write_column(column_index, array).await,
186            Self::V2_2(writer) => writer.write_column(column_index, array).await,
187            Self::V2_3(writer) => writer.write_column(column_index, array).await,
188        }
189    }
190
191    /// Append a buffer whose page or column metadata is supplied externally.
192    pub async fn write_external_buffer(&mut self, bytes: &[u8]) -> Result<(u64, u64)> {
193        match self {
194            Self::V2_0(writer) => writer.write_external_buffer(bytes).await,
195            Self::V2_1(writer) => writer.write_external_buffer(bytes).await,
196            Self::V2_2(writer) => writer.write_external_buffer(bytes).await,
197            Self::V2_3(writer) => writer.write_external_buffer(bytes).await,
198        }
199    }
200
201    /// Add a metadata entry to the schema.
202    pub fn add_schema_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
203        let key = key.into();
204        let value = value.into();
205        match self {
206            Self::V2_0(writer) => writer.add_schema_metadata(key, value),
207            Self::V2_1(writer) => writer.add_schema_metadata(key, value),
208            Self::V2_2(writer) => writer.add_schema_metadata(key, value),
209            Self::V2_3(writer) => writer.add_schema_metadata(key, value),
210        }
211    }
212
213    /// Prepare a writer from encoded columns whose buffers were produced externally.
214    pub fn initialize_with_external_columns(
215        &mut self,
216        schema: Schema,
217        columns: &[Arc<ColumnInfo>],
218        rows_written: u64,
219    ) -> Result<()> {
220        let column_metadata = columns
221            .iter()
222            .map(|column| column_info_to_metadata(column))
223            .collect::<Result<Vec<_>>>()?;
224        match self {
225            Self::V2_0(writer) => {
226                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
227            }
228            Self::V2_1(writer) => {
229                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
230            }
231            Self::V2_2(writer) => {
232                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
233            }
234            Self::V2_3(writer) => {
235                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
236            }
237        }
238        Ok(())
239    }
240
241    /// Add an arbitrary global buffer and return its one-based index.
242    pub async fn add_global_buffer(&mut self, buffer: Bytes) -> Result<u32> {
243        match self {
244            Self::V2_0(writer) => writer.add_global_buffer(buffer).await,
245            Self::V2_1(writer) => writer.add_global_buffer(buffer).await,
246            Self::V2_2(writer) => writer.add_global_buffer(buffer).await,
247            Self::V2_3(writer) => writer.add_global_buffer(buffer).await,
248        }
249    }
250
251    /// Finish the file and close its object writer.
252    pub async fn finish(&mut self) -> Result<FileWriteSummary> {
253        match self {
254            Self::V2_0(writer) => writer.finish().await,
255            Self::V2_1(writer) => writer.finish().await,
256            Self::V2_2(writer) => writer.finish().await,
257            Self::V2_3(writer) => writer.finish().await,
258        }
259    }
260
261    /// Abandon the file write.
262    pub async fn abort(&mut self) {
263        match self {
264            Self::V2_0(writer) => writer.abort().await,
265            Self::V2_1(writer) => writer.abort().await,
266            Self::V2_2(writer) => writer.abort().await,
267            Self::V2_3(writer) => writer.abort().await,
268        }
269    }
270
271    /// Return the current object-writer position.
272    pub async fn tell(&mut self) -> Result<u64> {
273        match self {
274            Self::V2_0(writer) => writer.tell().await,
275            Self::V2_1(writer) => writer.tell().await,
276            Self::V2_2(writer) => writer.tell().await,
277            Self::V2_3(writer) => writer.tell().await,
278        }
279    }
280
281    /// Return the field-id to physical-column mapping.
282    pub fn field_id_to_column_indices(&self) -> &[(u32, u32)] {
283        match self {
284            Self::V2_0(writer) => writer.field_id_to_column_indices(),
285            Self::V2_1(writer) => writer.field_id_to_column_indices(),
286            Self::V2_2(writer) => writer.field_id_to_column_indices(),
287            Self::V2_3(writer) => writer.field_id_to_column_indices(),
288        }
289    }
290}
291
292#[cfg(test)]
293#[path = "writer_tests.rs"]
294mod writer_tests;
lance_file/writer.rs

lance_file/
writer.rs
