From 3d6b97612c9662639112c8cf726b298525fb1090 Mon Sep 17 00:00:00 2001 From: Andrew Lamb Date: Thu, 26 Sep 2024 10:40:47 -0400 Subject: [PATCH] Example of reading and writing parquet metadata outside the file --- parquet/Cargo.toml | 6 + parquet/examples/external_metadata.rs | 221 ++++++++++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 parquet/examples/external_metadata.rs diff --git a/parquet/Cargo.toml b/parquet/Cargo.toml index 491122298e08..c0ae009d0ed6 100644 --- a/parquet/Cargo.toml +++ b/parquet/Cargo.toml @@ -121,6 +121,12 @@ sysinfo = ["dep:sysinfo"] # Verify 32-bit CRC checksum when decoding parquet pages crc = ["dep:crc32fast"] + +[[example]] +name = "external_metadata" +required-features = ["arrow", "async"] +path = "./examples/external_metadata.rs" + [[example]] name = "read_parquet" required-features = ["arrow"] diff --git a/parquet/examples/external_metadata.rs b/parquet/examples/external_metadata.rs new file mode 100644 index 000000000000..6677cd363580 --- /dev/null +++ b/parquet/examples/external_metadata.rs @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_cast::pretty::pretty_format_batches; +use futures::TryStreamExt; +use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use parquet::arrow::{ArrowWriter, ParquetRecordBatchStreamBuilder}; +use parquet::file::metadata::{ParquetMetaData, ParquetMetaDataReader, ParquetMetaDataWriter}; +use parquet::file::properties::{EnabledStatistics, WriterProperties}; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use tempfile::TempDir; + +/// This example demonstrates advanced usage of the Parquet metadata APIs. +/// +/// For example, you can copy the metadata for parquet files stored on remote +/// object storage (e.g. S3) to a local file or an in-memory cache, use a query +/// engine like DataFusion to analyze the metadata to determine which file to +/// read, and then read any matching files with a single object store request. +/// +/// # Usecase +/// +/// 1. Read Parquet metadata from Parquet file stored on a remote source +/// +/// 2. Store that metadata somewhere "locally" that is faster to access than +/// the rest of a parquet data (e.g. a cache). +/// +/// 3. Use the metadata to determine of the file should be read, and if so, read +/// the remote Parquet file, without re-reading / reparsing the metadata footer. +/// +/// # Example Overview +/// 1. Reads the metadata of a Parquet file +/// +/// 2. Removes some column statistics from the metadata (to make them smaller) +/// +/// 3. Stores the metadata in a separate file +/// +/// 4. Reads the metadata from the separate file and uses that to read the +/// Parquet file +#[tokio::main(flavor = "current_thread")] +async fn main() -> parquet::errors::Result<()> { + let tempdir = TempDir::new().unwrap(); + println!("data in {tempdir:?}"); + let parquet_path = create_parquet_file(&tempdir); + let metadata_path = tempdir.path().join("thrift_metadata.dat"); + + // In this example, we use a tokio file to mimic an async remote data source + let mut remote_parquet_file = tokio::fs::File::open(&parquet_path).await?; + + let metadata = get_metadata_from_remote_parquet_file(&mut remote_parquet_file).await; + println!( + "Metadata from 'remote' Parquet file into memory: {} bytes", + metadata.memory_size() + ); + + // now slim down the metadata and write it to a "local" file + let metadata = prepare_metadata(metadata); + write_metadata_to_local_file(metadata, &metadata_path); + + // now read the metadata from the local file and use it to read the "remote" Parquet file + let metadata = read_metadata_from_local_file(&metadata_path); + println!("Read metadata from file: {metadata:#?}"); + + let batches = read_remote_parquet_file_with_metadata(remote_parquet_file, metadata).await; + + // display the results + let batches_string = pretty_format_batches(&batches).unwrap().to_string(); + let batches_lines: Vec<_> = batches_string.split('\n').collect(); + + assert_eq!( + batches_lines, + [ + "+-----+-------------+", + "| id | description |", + "+-----+-------------+", + "| 100 | oranges |", + "| 200 | apples |", + "| 201 | grapefruit |", + "| 300 | bannanas |", + "| 102 | grapes |", + "| 33 | pears |", + "+-----+-------------+", + ], + "actual output:\n\n{batches_lines:#?}" + ); + + Ok(()) +} + +/// Reads the metadata from a "remote" parquet file +/// +/// Note that this function models reading from a remote file source using a +/// tokio file. In a real application, you would implement [`MetadataFetch`] for +/// your own remote source. +/// +/// [`MetadataFetch`]: parquet::arrow::async_reader::MetadataFetch +async fn get_metadata_from_remote_parquet_file( + remote_file: &mut tokio::fs::File, +) -> ParquetMetaData { + // the remote source must know the total file size (e.g. from an object store LIST operation) + let file_size = remote_file.metadata().await.unwrap().len(); + + // tell the reader to read the page index + ParquetMetaDataReader::new() + .with_page_indexes(true) + .load_and_finish(remote_file, file_size as usize) + .await + .unwrap() +} + +/// modifies the metadata to reduce its size +fn prepare_metadata(metadata: ParquetMetaData) -> ParquetMetaData { + // TODO + metadata +} + +/// writes the metadata to a file +/// +/// The data is stored using the same thrift format as the Parquet file metadata +fn write_metadata_to_local_file(metadata: ParquetMetaData, file: impl AsRef) { + let file = File::create(file).unwrap(); + ParquetMetaDataWriter::new(file, &metadata) + .finish() + .unwrap() +} + +/// Reads the metadata from a file +/// +/// This function reads the format written by `write_metadata_to_file` +fn read_metadata_from_local_file(file: impl AsRef) -> ParquetMetaData { + let file = File::open(file).unwrap(); + ParquetMetaDataReader::new() + .with_column_indexes(true) + .with_offset_indexes(true) + .parse_and_finish(&file) + .unwrap() +} + +/// Reads the "remote" Parquet file using the metadata +/// +/// This shows how to read the Parquet file using previously read metadata +/// instead of the metadata in the Parquet file itself. This avoids an IO / +/// having to fetch and decode the metadata from the Parquet file before +/// beginning to read it. +/// +/// Note that this function models reading from a remote file source using a +/// tokio file. In a real application, you would implement [`AsyncFileReader`] +/// for your own remote source. +/// +/// In this example, we simply read the results as Arrow record batches +/// +/// [`AsyncFileReader`]: parquet::arrow::async_reader::AsyncFileReader +async fn read_remote_parquet_file_with_metadata( + remote_file: tokio::fs::File, + metadata: ParquetMetaData, +) -> Vec { + let options = ArrowReaderOptions::new() + // tell the reader to read the page index + .with_page_index(true); + // create a reader with pre-existing metadata + let arrow_reader_metadata = ArrowReaderMetadata::try_new(metadata.into(), options).unwrap(); + let reader = + ParquetRecordBatchStreamBuilder::new_with_metadata(remote_file, arrow_reader_metadata) + .build() + .unwrap(); + + reader.try_collect::>().await.unwrap() +} + +/// Make a new parquet file in the temporary directory, and returns the path +fn create_parquet_file(tmpdir: &TempDir) -> PathBuf { + let path = tmpdir.path().join("example.parquet"); + let new_file = File::create(&path).unwrap(); + + let batch = RecordBatch::try_from_iter(vec![ + ( + "id", + Arc::new(Int32Array::from(vec![100, 200, 201, 300, 102, 33])) as ArrayRef, + ), + ( + "description", + Arc::new(StringArray::from(vec![ + "oranges", + "apples", + "grapefruit", + "bannanas", + "grapes", + "pears", + ])), + ), + ]) + .unwrap(); + + let props = WriterProperties::builder() + // ensure we write the page index level statistics + .set_statistics_enabled(EnabledStatistics::Page) + .build(); + + let mut writer = ArrowWriter::try_new(new_file, batch.schema(), Some(props)).unwrap(); + + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + + path +}