From f6ea8431b39681880c33b86bcd62a1072c753916 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Thu, 30 Jul 2026 14:44:17 -0400 Subject: [PATCH] Parquet arrow schema --- crates/datafusion-app/src/local.rs | 4 + .../datafusion-functions-parquet/Cargo.toml | 2 +- .../src/arrow_schema.rs | 280 ++++++++++++++++++ .../datafusion-functions-parquet/src/lib.rs | 2 + docs/features.md | 6 + 5 files changed, 293 insertions(+), 1 deletion(-) create mode 100644 crates/datafusion-functions-parquet/src/arrow_schema.rs diff --git a/crates/datafusion-app/src/local.rs b/crates/datafusion-app/src/local.rs index 35a89e2..3619029 100644 --- a/crates/datafusion-app/src/local.rs +++ b/crates/datafusion-app/src/local.rs @@ -123,6 +123,10 @@ impl ExecutionContext { "parquet_kv_metadata", Arc::new(datafusion_functions_parquet::ParquetKvMetadataFunc {}), ); + session_ctx.register_udtf( + "parquet_arrow_schema", + Arc::new(datafusion_functions_parquet::ParquetArrowSchemaFunc {}), + ); session_ctx.register_udtf( "parquet_bloom_filter", Arc::new(datafusion_functions_parquet::ParquetBloomFilterFunc {}), diff --git a/crates/datafusion-functions-parquet/Cargo.toml b/crates/datafusion-functions-parquet/Cargo.toml index 6acf047..b9b4474 100644 --- a/crates/datafusion-functions-parquet/Cargo.toml +++ b/crates/datafusion-functions-parquet/Cargo.toml @@ -13,11 +13,11 @@ version = "0.1.0" [dependencies] arrow = { version = "58" } async-trait = "0.1.41" +base64 = "0.22" datafusion = { version = "54" } parquet = { default-features = false, version = "58.3" } [dev-dependencies] -base64 = "0.22" parquet = { features = ["arrow"], version = "58" } tempfile = "3" tokio = { features = ["macros", "rt"], version = "1" } diff --git a/crates/datafusion-functions-parquet/src/arrow_schema.rs b/crates/datafusion-functions-parquet/src/arrow_schema.rs new file mode 100644 index 0000000..fd3dabc --- /dev/null +++ b/crates/datafusion-functions-parquet/src/arrow_schema.rs @@ -0,0 +1,280 @@ +// 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::{BooleanArray, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use async_trait::async_trait; +use base64::prelude::{Engine, BASE64_STANDARD}; +use datafusion::catalog::Session; +use datafusion::catalog::TableFunctionImpl; +use datafusion::common::{plan_err, Column}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::TableProvider; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::Expr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::scalar::ScalarValue; +use parquet::file::reader::FileReader; +use parquet::file::serialized_reader::SerializedFileReader; +use std::fs::File; +use std::sync::Arc; + +// Same value as `parquet::arrow::ARROW_SCHEMA_META_KEY`, which is unavailable +// here because this crate builds `parquet` without its `arrow` feature. +const ARROW_SCHEMA_META_KEY: &str = "ARROW:schema"; + +#[derive(Debug)] +struct ParquetArrowSchemaTable { + schema: SchemaRef, + batch: RecordBatch, +} + +#[async_trait] +impl TableProvider for ParquetArrowSchemaTable { + fn schema(&self) -> arrow::datatypes::SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> datafusion::logical_expr::TableType { + datafusion::logical_expr::TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + TableProvider::schema(self), + projection.cloned(), + )?) + } +} + +/// `parquet_arrow_schema` table-valued function +/// +/// Decodes the Arrow schema stored in a Parquet file's `ARROW:schema` +/// key-value metadata entry (a base64-encoded Arrow IPC schema message) and +/// returns one row per field. +/// +/// Example: +/// ```sql +/// SELECT * FROM parquet_arrow_schema('file.parquet'); +/// ``` +#[derive(Debug)] +pub struct ParquetArrowSchemaFunc {} + +impl TableFunctionImpl for ParquetArrowSchemaFunc { + fn call(&self, exprs: &[Expr]) -> Result> { + let filename = match exprs.first() { + Some(Expr::Literal(ScalarValue::Utf8(Some(s)), _)) => s, + Some(Expr::Column(Column { name, .. })) => name, + _ => { + return plan_err!("parquet_arrow_schema requires a string argument as its input"); + } + }; + + let file = File::open(filename.clone())?; + let reader = SerializedFileReader::new(file)?; + let metadata = reader.metadata(); + + let encoded = metadata + .file_metadata() + .key_value_metadata() + .and_then(|kvs| { + kvs.iter() + .find(|kv| kv.key == ARROW_SCHEMA_META_KEY) + .and_then(|kv| kv.value.clone()) + }); + let Some(encoded) = encoded else { + return plan_err!("no {ARROW_SCHEMA_META_KEY} key-value metadata found in {filename}"); + }; + + let decoded = BASE64_STANDARD.decode(encoded).map_err(|e| { + DataFusionError::Execution(format!( + "failed to base64 decode {ARROW_SCHEMA_META_KEY} value in {filename}: {e}" + )) + })?; + let arrow_schema = arrow::ipc::convert::try_schema_from_ipc_buffer(&decoded)?; + + let schema = Arc::new(Schema::new(vec![ + Field::new("filename", DataType::Utf8, true), + Field::new("field_name", DataType::Utf8, true), + Field::new("data_type", DataType::Utf8, true), + Field::new("nullable", DataType::Boolean, true), + Field::new("metadata", DataType::Utf8, true), + ])); + + let mut filename_arr: Vec> = vec![]; + let mut field_name_arr: Vec> = vec![]; + let mut data_type_arr: Vec> = vec![]; + let mut nullable_arr: Vec> = vec![]; + let mut metadata_arr: Vec> = vec![]; + + for field in arrow_schema.fields() { + filename_arr.push(Some(filename.clone())); + field_name_arr.push(Some(field.name().clone())); + data_type_arr.push(Some(field.data_type().to_string())); + nullable_arr.push(Some(field.is_nullable())); + metadata_arr.push(if field.metadata().is_empty() { + None + } else { + Some(format!("{:?}", field.metadata())) + }); + } + + let rb = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(filename_arr)), + Arc::new(StringArray::from(field_name_arr)), + Arc::new(StringArray::from(data_type_arr)), + Arc::new(BooleanArray::from(nullable_arr)), + Arc::new(StringArray::from(metadata_arr)), + ], + )?; + + Ok(Arc::new(ParquetArrowSchemaTable { schema, batch: rb })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Int32Array, TimestampNanosecondArray}; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use arrow::record_batch::RecordBatch; + use datafusion::prelude::SessionContext; + use parquet::arrow::arrow_writer::ArrowWriterOptions; + use parquet::arrow::ArrowWriter; + use std::path::Path; + + fn make_ctx() -> SessionContext { + let ctx = SessionContext::new(); + ctx.register_udtf("parquet_arrow_schema", Arc::new(ParquetArrowSchemaFunc {})); + ctx + } + + /// Write a parquet file at `path` with a nullable int column and a + /// timezone-aware timestamp column (the latter is only fully described by + /// the embedded Arrow schema, not the Parquet schema). + fn write_parquet(path: &Path) { + let tz = "America/New_York"; + let schema = Arc::new(Schema::new(vec![ + Field::new("v", DataType::Int32, true), + Field::new( + "ts", + DataType::Timestamp(TimeUnit::Nanosecond, Some(tz.into())), + false, + ), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from(vec![1i32])), + Arc::new(TimestampNanosecondArray::from(vec![1_000_000_000i64]).with_timezone(tz)), + ], + ) + .unwrap(); + let file = std::fs::File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + + fn str_val(batch: &RecordBatch, row: usize, col: &str) -> Option { + let array = batch + .column(batch.schema().index_of(col).unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + array.is_valid(row).then(|| array.value(row).to_string()) + } + + /// One row per field with name, data type (including timezone), and nullability. + #[tokio::test] + async fn test_fields_returned() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t.parquet"); + write_parquet(&path); + + let ctx = make_ctx(); + let sql = format!("SELECT * FROM parquet_arrow_schema('{}')", path.display()); + let result = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let batch = arrow::compute::concat_batches(&result[0].schema(), &result).unwrap(); + assert_eq!(batch.num_rows(), 2); + + assert_eq!(str_val(&batch, 0, "field_name").as_deref(), Some("v")); + assert_eq!(str_val(&batch, 0, "data_type").as_deref(), Some("Int32")); + assert_eq!(str_val(&batch, 1, "field_name").as_deref(), Some("ts")); + let ts_type = str_val(&batch, 1, "data_type").unwrap(); + assert!( + ts_type.contains("Timestamp") && ts_type.contains("America/New_York"), + "unexpected data_type: {ts_type}" + ); + + let nullable = batch + .column(batch.schema().index_of("nullable").unwrap()) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(nullable.value(0)); + assert!(!nullable.value(1)); + + assert_eq!( + str_val(&batch, 0, "filename").as_deref(), + Some(path.display().to_string().as_str()) + ); + } + + /// A file written without the embedded Arrow schema produces an error. + #[tokio::test] + async fn test_missing_arrow_schema() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t.parquet"); + + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32, false)])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1i32]))]) + .unwrap(); + let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = ArrowWriter::try_new_with_options(file, schema, options).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let ctx = make_ctx(); + let sql = format!("SELECT * FROM parquet_arrow_schema('{}')", path.display()); + let err = ctx.sql(&sql).await.unwrap_err().to_string(); + assert!( + err.contains("no ARROW:schema key-value metadata found"), + "{err}" + ); + } + + /// Missing argument produces a plan error. + #[tokio::test] + async fn test_missing_argument() { + let ctx = make_ctx(); + let err = ctx.sql("SELECT * FROM parquet_arrow_schema()").await; + assert!(err.is_err()); + } +} diff --git a/crates/datafusion-functions-parquet/src/lib.rs b/crates/datafusion-functions-parquet/src/lib.rs index d8de84a..b1a6259 100644 --- a/crates/datafusion-functions-parquet/src/lib.rs +++ b/crates/datafusion-functions-parquet/src/lib.rs @@ -37,9 +37,11 @@ use parquet::file::statistics::Statistics; use std::fs::File; use std::sync::Arc; +mod arrow_schema; mod bloom_filter; mod kv_metadata; mod page_index; +pub use arrow_schema::ParquetArrowSchemaFunc; pub use bloom_filter::{ParquetBloomFilterCheckFunc, ParquetBloomFilterFunc}; pub use kv_metadata::ParquetKvMetadataFunc; pub use page_index::ParquetPageIndexFunc; diff --git a/docs/features.md b/docs/features.md index 080b838..fd35a9e 100644 --- a/docs/features.md +++ b/docs/features.md @@ -18,6 +18,12 @@ Or for reading the key-value metadata from a file's footer: SELECT * FROM parquet_kv_metadata('my_parquet_file.parquet') ``` +The Arrow schema embedded in the `ARROW:schema` key-value metadata entry (written by Arrow-based writers) can be decoded and printed with one row per field: + +```sql +SELECT * FROM parquet_arrow_schema('my_parquet_file.parquet') +``` + Bloom filter details (presence, location, and size per column chunk) can be inspected with: ```sql