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
79/// Serialize an execution description using its format's wire grammar.
80pub(crate) fn encode_page_encoding(encoding: &PageEncoding) -> Result<Vec<u8>> {
81    let envelope = match encoding {
82        PageEncoding::Legacy(layout) => Any::from_msg(layout)?,
83        PageEncoding::Structural(layout) => {
84            let layout = lance_encoding::format::pb21::PageLayout::try_from(layout.clone())?;
85            Any::from_msg(&layout)?
86        }
87        PageEncoding::Planar(layout) => Any::from_msg(layout)?,
88    };
89    Ok(envelope.encode_to_vec())
90}
91
92fn column_info_to_metadata(column: &ColumnInfo) -> Result<pbfile::ColumnMetadata> {
93    let pages = column
94        .page_infos
95        .iter()
96        .map(|page| {
97            let encoding = encode_page_encoding(&page.encoding)?;
98            let (buffer_offsets, buffer_sizes) =
99                page.buffer_offsets_and_sizes.iter().copied().unzip();
100            Ok(pbfile::column_metadata::Page {
101                buffer_offsets,
102                buffer_sizes,
103                encoding: Some(pbfile::Encoding {
104                    location: Some(pbfile::encoding::Location::Direct(pbfile::DirectEncoding {
105                        encoding,
106                    })),
107                }),
108                length: page.num_rows,
109                priority: page.priority,
110            })
111        })
112        .collect::<Result<Vec<_>>>()?;
113    let (buffer_offsets, buffer_sizes) = column.buffer_offsets_and_sizes.iter().copied().unzip();
114    let encoding = Any::from_msg(&column.encoding)?.encode_to_vec();
115    Ok(pbfile::ColumnMetadata {
116        pages,
117        buffer_offsets,
118        buffer_sizes,
119        encoding: Some(pbfile::Encoding {
120            location: Some(pbfile::encoding::Location::Direct(pbfile::DirectEncoding {
121                encoding,
122            })),
123        }),
124    })
125}
126
127impl From<versions::v2_0::Writer> for FileWriter {
128    fn from(writer: versions::v2_0::Writer) -> Self {
129        Self::V2_0(Box::new(writer))
130    }
131}
132
133impl From<versions::v2_1::Writer> for FileWriter {
134    fn from(writer: versions::v2_1::Writer) -> Self {
135        Self::V2_1(Box::new(writer))
136    }
137}
138
139impl From<versions::v2_2::Writer> for FileWriter {
140    fn from(writer: versions::v2_2::Writer) -> Self {
141        Self::V2_2(Box::new(writer))
142    }
143}
144
145impl From<versions::v2_3::Writer> for FileWriter {
146    fn from(writer: versions::v2_3::Writer) -> Self {
147        Self::V2_3(Box::new(writer))
148    }
149}
150
151impl FileWriter {
152    /// Spill page metadata to a sidecar file.
153    pub fn with_page_metadata_spill(self, object_store: Arc<ObjectStore>, path: Path) -> Self {
154        match self {
155            Self::V2_0(writer) => Self::V2_0(Box::new(
156                (*writer).with_page_metadata_spill(object_store, path),
157            )),
158            Self::V2_1(writer) => Self::V2_1(Box::new(
159                (*writer).with_page_metadata_spill(object_store, path),
160            )),
161            Self::V2_2(writer) => Self::V2_2(Box::new(
162                (*writer).with_page_metadata_spill(object_store, path),
163            )),
164            Self::V2_3(writer) => Self::V2_3(Box::new(
165                (*writer).with_page_metadata_spill(object_store, path),
166            )),
167        }
168    }
169
170    /// Schedule batches of data to be written to the file.
171    pub async fn write_batches(
172        &mut self,
173        batches: impl Iterator<Item = &RecordBatch>,
174    ) -> Result<()> {
175        for batch in batches {
176            self.write_batch(batch).await?;
177        }
178        Ok(())
179    }
180
181    /// Schedule a batch of data to be written to the file.
182    pub async fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
183        match self {
184            Self::V2_0(writer) => writer.write_batch(batch).await,
185            Self::V2_1(writer) => writer.write_batch(batch).await,
186            Self::V2_2(writer) => writer.write_batch(batch).await,
187            Self::V2_3(writer) => writer.write_batch(batch).await,
188        }
189    }
190
191    /// Write one top-level column.
192    pub async fn write_column(&mut self, column_index: usize, array: ArrayRef) -> Result<()> {
193        match self {
194            Self::V2_0(writer) => writer.write_column(column_index, array).await,
195            Self::V2_1(writer) => writer.write_column(column_index, array).await,
196            Self::V2_2(writer) => writer.write_column(column_index, array).await,
197            Self::V2_3(writer) => writer.write_column(column_index, array).await,
198        }
199    }
200
201    /// Append a buffer whose page or column metadata is supplied externally.
202    pub async fn write_external_buffer(&mut self, bytes: &[u8]) -> Result<(u64, u64)> {
203        match self {
204            Self::V2_0(writer) => writer.write_external_buffer(bytes).await,
205            Self::V2_1(writer) => writer.write_external_buffer(bytes).await,
206            Self::V2_2(writer) => writer.write_external_buffer(bytes).await,
207            Self::V2_3(writer) => writer.write_external_buffer(bytes).await,
208        }
209    }
210
211    /// Add a metadata entry to the schema.
212    pub fn add_schema_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
213        let key = key.into();
214        let value = value.into();
215        match self {
216            Self::V2_0(writer) => writer.add_schema_metadata(key, value),
217            Self::V2_1(writer) => writer.add_schema_metadata(key, value),
218            Self::V2_2(writer) => writer.add_schema_metadata(key, value),
219            Self::V2_3(writer) => writer.add_schema_metadata(key, value),
220        }
221    }
222
223    /// Prepare a writer from encoded columns whose buffers were produced externally.
224    pub fn initialize_with_external_columns(
225        &mut self,
226        schema: Schema,
227        columns: &[Arc<ColumnInfo>],
228        rows_written: u64,
229    ) -> Result<()> {
230        let column_metadata = columns
231            .iter()
232            .map(|column| column_info_to_metadata(column))
233            .collect::<Result<Vec<_>>>()?;
234        match self {
235            Self::V2_0(writer) => {
236                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
237            }
238            Self::V2_1(writer) => {
239                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
240            }
241            Self::V2_2(writer) => {
242                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
243            }
244            Self::V2_3(writer) => {
245                writer.initialize_with_external_metadata(schema, column_metadata, rows_written)
246            }
247        }
248        Ok(())
249    }
250
251    /// Add an arbitrary global buffer and return its one-based index.
252    pub async fn add_global_buffer(&mut self, buffer: Bytes) -> Result<u32> {
253        match self {
254            Self::V2_0(writer) => writer.add_global_buffer(buffer).await,
255            Self::V2_1(writer) => writer.add_global_buffer(buffer).await,
256            Self::V2_2(writer) => writer.add_global_buffer(buffer).await,
257            Self::V2_3(writer) => writer.add_global_buffer(buffer).await,
258        }
259    }
260
261    /// Finish the file and close its object writer.
262    pub async fn finish(&mut self) -> Result<FileWriteSummary> {
263        match self {
264            Self::V2_0(writer) => writer.finish().await,
265            Self::V2_1(writer) => writer.finish().await,
266            Self::V2_2(writer) => writer.finish().await,
267            Self::V2_3(writer) => writer.finish().await,
268        }
269    }
270
271    /// Abandon the file write.
272    pub async fn abort(&mut self) {
273        match self {
274            Self::V2_0(writer) => writer.abort().await,
275            Self::V2_1(writer) => writer.abort().await,
276            Self::V2_2(writer) => writer.abort().await,
277            Self::V2_3(writer) => writer.abort().await,
278        }
279    }
280
281    /// Return the current object-writer position.
282    pub async fn tell(&mut self) -> Result<u64> {
283        match self {
284            Self::V2_0(writer) => writer.tell().await,
285            Self::V2_1(writer) => writer.tell().await,
286            Self::V2_2(writer) => writer.tell().await,
287            Self::V2_3(writer) => writer.tell().await,
288        }
289    }
290
291    /// Return the field-id to physical-column mapping.
292    pub fn field_id_to_column_indices(&self) -> &[(u32, u32)] {
293        match self {
294            Self::V2_0(writer) => writer.field_id_to_column_indices(),
295            Self::V2_1(writer) => writer.field_id_to_column_indices(),
296            Self::V2_2(writer) => writer.field_id_to_column_indices(),
297            Self::V2_3(writer) => writer.field_id_to_column_indices(),
298        }
299    }
300}
301
302#[cfg(test)]
303#[path = "writer_tests.rs"]
304mod writer_tests;
lance_file/writer.rs

lance_file/
writer.rs
