mz_persist/indexed/columnar/
parquet.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Apache Parquet encodings and utils for persist data
11
12use std::io::Write;
13
14use arrow::record_batch::{RecordBatch, RecordBatchReader};
15use differential_dataflow::trace::Description;
16use mz_ore::bytes::SegmentedBytes;
17use mz_ore::cast::CastFrom;
18use mz_persist_types::Codec64;
19use mz_persist_types::parquet::EncodingConfig;
20use parquet::arrow::ArrowWriter;
21use parquet::arrow::arrow_reader::{
22    ArrowReaderMetadata, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder,
23};
24use parquet::basic::Encoding;
25use parquet::file::metadata::{KeyValue, ParquetMetaData};
26use parquet::file::properties::{EnabledStatistics, WriterProperties, WriterVersion};
27use timely::PartialOrder;
28use timely::progress::{Antichain, Timestamp};
29
30use crate::error::Error;
31use crate::generated::persist::ProtoBatchFormat;
32use crate::generated::persist::proto_batch_part_inline::FormatMetadata as ProtoFormatMetadata;
33use crate::indexed::columnar::arrow::{decode_arrow_batch, encode_arrow_batch};
34use crate::indexed::encoding::{
35    BlobTraceBatchPart, BlobTraceUpdates, decode_trace_inline_meta, encode_trace_inline_meta,
36    validate_trace_updates,
37};
38use crate::metrics::{ColumnarMetrics, ParquetColumnMetrics};
39
40const INLINE_METADATA_KEY: &str = "MZ:inline";
41
42/// Encodes a [`BlobTraceBatchPart`] into the Parquet format.
43pub fn encode_trace_parquet<W: Write + Send, T: Timestamp + Codec64>(
44    w: &mut W,
45    batch: &BlobTraceBatchPart<T>,
46    metrics: &ColumnarMetrics,
47    cfg: &EncodingConfig,
48) -> Result<(), Error> {
49    // Better to error now than write out an invalid batch.
50    batch.validate()?;
51
52    let inline_meta = encode_trace_inline_meta(batch);
53    encode_parquet_kvtd(w, inline_meta, &batch.updates, metrics, cfg)
54}
55
56/// Decodes a BlobTraceBatchPart from the Parquet format.
57pub fn decode_trace_parquet<T: Timestamp + Codec64>(
58    buf: SegmentedBytes,
59    metrics: &ColumnarMetrics,
60) -> Result<BlobTraceBatchPart<T>, Error> {
61    let mut reader = BlobTraceBatchPartReader::new(buf, usize::MAX)?;
62    let updates = match reader.next_updates(metrics) {
63        Some(updates) => updates?,
64        None => {
65            let empty = RecordBatch::new_empty(reader.reader.schema());
66            reader.decode_batch(empty, metrics)?
67        }
68    };
69    // We enforce an invariant that we have a single RowGroup.
70    if reader.next_updates(metrics).is_some() {
71        return Err(Error::String("found more than one RowGroup".to_string()));
72    }
73    Ok(BlobTraceBatchPart {
74        desc: reader.desc,
75        index: reader.index,
76        updates,
77    })
78}
79
80/// Incremental decoder over a parquet-encoded [`BlobTraceBatchPart`].
81///
82/// Decodes at most `batch_rows` rows per call to [`Self::next_updates`], so
83/// only the encoded bytes and one decoded batch are resident at a time. Each
84/// batch is validated like [`BlobTraceBatchPart::validate`] validates a whole
85/// part.
86///
87/// Successive batches cover the part's rows in their encoded order, without
88/// gaps or overlap. A caller that consolidates adjacent updates can therefore
89/// carry its pending update across a batch boundary and merge the same
90/// updates as it would on the whole part.
91#[derive(Debug)]
92pub struct BlobTraceBatchPartReader<T> {
93    desc: Description<T>,
94    index: u64,
95    num_rows: usize,
96    /// Drop the `k_s` and `v_s` columns of the deprecated
97    /// `StructuredMigration(1)` format.
98    project_v1: bool,
99    reader: ParquetRecordBatchReader,
100}
101
102impl<T: Timestamp + Codec64> BlobTraceBatchPartReader<T> {
103    /// Opens `buf` for decoding in batches of at most `batch_rows` rows.
104    ///
105    /// A `batch_rows` of 0 is treated as 1.
106    pub fn new(buf: SegmentedBytes, batch_rows: usize) -> Result<Self, Error> {
107        let metadata = ArrowReaderMetadata::load(&buf, Default::default())?;
108        let inline = metadata
109            .metadata()
110            .file_metadata()
111            .key_value_metadata()
112            .and_then(|x| x.iter().find(|x| x.key == INLINE_METADATA_KEY));
113        let (format, inline) = decode_trace_inline_meta(inline.and_then(|x| x.value.as_ref()))?;
114        let project_v1 = match format {
115            ProtoBatchFormat::Unknown => return Err("unknown format".into()),
116            ProtoBatchFormat::ArrowKvtd => {
117                return Err("ArrowKVTD format not supported in parquet".into());
118            }
119            ProtoBatchFormat::ParquetKvtd => false,
120            // Even though `format_metadata` is optional, we expect it when
121            // our format is ParquetStructured.
122            ProtoBatchFormat::ParquetStructured => match inline.format_metadata {
123                None => return Err("missing field 'format_metadata'".into()),
124                Some(ProtoFormatMetadata::StructuredMigration(v @ 1..=3)) => v == 1,
125                unknown => Err(format!("unkown ProtoFormatMetadata, {unknown:?}"))?,
126            },
127        };
128        let desc = inline.desc.map_or_else(
129            || {
130                Description::new(
131                    Antichain::from_elem(T::minimum()),
132                    Antichain::from_elem(T::minimum()),
133                    Antichain::from_elem(T::minimum()),
134                )
135            },
136            |x| x.into(),
137        );
138
139        // Checked here as well as per batch, since a part without rows yields
140        // no batches.
141        if PartialOrder::less_equal(desc.upper(), desc.lower()) {
142            return Err(format!("invalid desc: {:?}", desc).into());
143        }
144
145        let builder = ParquetRecordBatchReaderBuilder::new_with_metadata(buf, metadata);
146        let row_groups = builder.metadata().row_groups();
147        if row_groups.len() > 1 {
148            return Err(Error::String("found more than 1 RowGroup".to_string()));
149        }
150        let num_rows = match row_groups.first() {
151            Some(row_group) => usize::try_from(row_group.num_rows())
152                .map_err(|_| Error::String("found negative rows".to_string()))?,
153            None => 0,
154        };
155        let reader = builder
156            .with_batch_size(batch_rows.min(num_rows).max(1))
157            .build()?;
158        Ok(BlobTraceBatchPartReader {
159            desc,
160            index: inline.index,
161            num_rows,
162            project_v1,
163            reader,
164        })
165    }
166
167    /// The part's inline description.
168    pub fn desc(&self) -> &Description<T> {
169        &self.desc
170    }
171
172    /// The total number of rows in the part, decoded or not.
173    pub fn num_rows(&self) -> usize {
174        self.num_rows
175    }
176
177    /// Decodes the next batch of rows, or returns `None` once all rows have
178    /// been returned.
179    pub fn next_updates(
180        &mut self,
181        metrics: &ColumnarMetrics,
182    ) -> Option<Result<BlobTraceUpdates, Error>> {
183        let batch = match self.reader.next()? {
184            Ok(batch) => batch,
185            Err(e) => return Some(Err(Error::String(e.to_string()))),
186        };
187        Some(self.decode_batch(batch, metrics))
188    }
189
190    fn decode_batch(
191        &self,
192        mut batch: RecordBatch,
193        metrics: &ColumnarMetrics,
194    ) -> Result<BlobTraceUpdates, Error> {
195        // Version 1 is a deprecated format so we just ignored the k_s and v_s columns.
196        if self.project_v1 && batch.num_columns() > 4 {
197            batch = batch.project(&[0, 1, 2, 3])?;
198        }
199        let updates = decode_arrow_batch(&batch, metrics).map_err(|e| e.to_string())?;
200        validate_trace_updates(&self.desc, &updates)?;
201        Ok(updates)
202    }
203}
204
205/// Encodes [`BlobTraceUpdates`] to Parquet using the [`parquet`] crate.
206pub fn encode_parquet_kvtd<W: Write + Send>(
207    w: &mut W,
208    inline_base64: String,
209    updates: &BlobTraceUpdates,
210    metrics: &ColumnarMetrics,
211    cfg: &EncodingConfig,
212) -> Result<(), Error> {
213    let metadata = KeyValue::new(INLINE_METADATA_KEY.to_string(), inline_base64);
214
215    // Note: most of these settings are the defaults from `arrow2` which we
216    // previously used and maintain until we tune with benchmarking.
217    let properties = WriterProperties::builder()
218        .set_dictionary_enabled(cfg.use_dictionary)
219        .set_encoding(Encoding::PLAIN)
220        .set_statistics_enabled(EnabledStatistics::None)
221        .set_compression(cfg.compression.into())
222        .set_writer_version(WriterVersion::PARQUET_2_0)
223        .set_data_page_size_limit(1024 * 1024)
224        .set_max_row_group_row_count(None)
225        .set_key_value_metadata(Some(vec![metadata]))
226        .build();
227
228    let batch = encode_arrow_batch(updates);
229    let format = match updates {
230        BlobTraceUpdates::Row(_) => "k,v,t,d",
231        BlobTraceUpdates::Both(_, _) => "k,v,t,d,k_s,v_s",
232        BlobTraceUpdates::Structured { .. } => "t,d,k_s,v_s",
233    };
234
235    let mut writer = ArrowWriter::try_new(w, batch.schema(), Some(properties))?;
236    writer.write(&batch)?;
237
238    writer.flush()?;
239    let bytes_written = writer.bytes_written();
240    let file_metadata = writer.close()?;
241
242    report_parquet_metrics(metrics, &file_metadata, bytes_written, format);
243
244    Ok(())
245}
246
247/// Best effort reporting of metrics from the resulting [`parquet::format::FileMetaData`] returned
248/// from the [`ArrowWriter`].
249fn report_parquet_metrics(
250    metrics: &ColumnarMetrics,
251    metadata: &ParquetMetaData,
252    bytes_written: usize,
253    format: &'static str,
254) {
255    metrics
256        .parquet()
257        .num_row_groups
258        .with_label_values(&[format])
259        .inc_by(u64::cast_from(metadata.row_groups().len()));
260    metrics
261        .parquet()
262        .encoded_size
263        .with_label_values(&[format])
264        .inc_by(u64::cast_from(bytes_written));
265
266    let report_column_size = |col_name: &str, metrics: &ParquetColumnMetrics| {
267        let (uncomp, comp) = metadata
268            .row_groups()
269            .iter()
270            .map(|row_group| row_group.columns().iter())
271            .flatten()
272            .filter(|m| m.column_path().parts().first().map(|s| s.as_str()) == Some(col_name))
273            .map(|m| (m.uncompressed_size(), m.compressed_size()))
274            .fold((0, 0), |(tot_u, tot_c), (u, c)| (tot_u + u, tot_c + c));
275
276        let uncomp = uncomp.try_into().unwrap_or(0u64);
277        let comp = comp.try_into().unwrap_or(0u64);
278
279        metrics.report_sizes(uncomp, comp);
280    };
281
282    report_column_size("k", &metrics.parquet().k_metrics);
283    report_column_size("v", &metrics.parquet().v_metrics);
284    report_column_size("t", &metrics.parquet().t_metrics);
285    report_column_size("d", &metrics.parquet().d_metrics);
286    report_column_size("k_s", &metrics.parquet().k_s_metrics);
287    report_column_size("v_s", &metrics.parquet().v_s_metrics);
288}
mz_persist/indexed/columnar/parquet.rs

mz_persist/indexed/columnar/
parquet.rs
