Add an example of embedding indexes inside a parquet file

[zhuqi-lucas]

Which issue does this PR close?

• Closes Add an example of embedding indexes *inside* a parquet file #16374

Rationale for this change

//! Example: embedding and using a custom “distinct values” index in Parquet files
//!
//! This example shows how to build and leverage a file‑level distinct‑values index
//! for pruning in DataFusion’s Parquet scans.
//!
//! Steps:
//! 1. Compute the distinct values for a target column and serialize them into bytes.
//! 2. Write each Parquet file with:
//!    - regular data pages for your column
//!    - the magic marker `IDX1` and a little‑endian length, to identify our custom index format
//!    - the serialized distinct‑values bytes
//!    - footer key/value metadata entries (`distinct_index_offset` and `distinct_index_length`)
//! 3. Read back each file’s footer metadata to locate and deserialize the index.
//! 4. Build a `DistinctIndexTable` (a custom `TableProvider`) that scans footers
//!    into a map of filename → `HashSet<String>` of distinct values.
//! 5. In `scan()`, prune out any Parquet files whose distinct set doesn’t match the
//!    `category = 'X'` filter, then only read data from the remaining files.
//!
//! This technique embeds a lightweight, application‑specific index directly in Parquet
//! metadata to achieve efficient file‑level pruning without modifying the Parquet format.
//!
//! And it's very efficient, since we don't add any additional info to the metadata, we write the custom index
//! after the data pages, and we only read it when needed.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

[alamb]

I think this writes the index into the footer itself (as an opaque string)

This has at least 2 downsides

1. The footer metadata will be much larger / longer to parse
2. A binary index must be converted to/from strings (as you are doing here with b64)

Is it possible to write the binary data directly into the parquet file?

Specifically, so then the metadata looks something like

// Find out where the current write position is
let offset_to_index_in_file = file.current_position()
file.write_all(distinct_index)?;

// now, finalize the file with the parquet metadata:
let props = WriterProperties::builder()
        .set_key_value_metadata(Some(vec![KeyValue::new(
            "distinct_index_data".into(),
            offset_to_index_in_file.to_string(),
        )]))
        .build();

I am not sure how easy this would be to do with the current API

[zhuqi-lucas]

Very good point @alamb! Thank you.

I will try to find a better solution, i agree the following downsides.

This has at least 2 downsides

1. The footer metadata will be much larger / longer to parse
2. A binary index must be converted to/from strings (as you are doing here with b64)

[zhuqi-lucas]

I tried today, but found it's hard for current API to support this, will try it again.

[zhuqi-lucas]

Try to using low level API, but it only works when we disable page index, if we setting page index, it will follow up the real row group data, and it conflicts with our embedding indexes.

[zhuqi-lucas]

The code is here:

// 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.

//! Example: embedding a "distinct values" index in a Parquet file's metadata
//!
//! 1. Read existing Parquet files
//! 2. Compute distinct values for a target column using DataFusion
//! 3. Serialize the distinct index to bytes and write to the new Parquet file
//!    with these encoded bytes appended as a custom metadata entry
//! 4. Read each new parquet file, extract and deserialize the index from footer
//! 5. Use the distinct index to prune files when querying

use arrow::array::{ArrayRef, StringArray};
use arrow::record_batch::RecordBatch;
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use async_trait::async_trait;
use datafusion::catalog::{Session, TableProvider};
use datafusion::common::{HashMap, HashSet, Result};
use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::memory::DataSourceExec;
use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource};
use datafusion::datasource::TableType;
use datafusion::execution::object_store::ObjectStoreUrl;
use datafusion::logical_expr::{Operator, TableProviderFilterPushDown};
use datafusion::parquet::arrow::ArrowSchemaConverter;
use datafusion::parquet::data_type::{ByteArray, ByteArrayType};
use datafusion::parquet::errors::ParquetError;
use datafusion::parquet::file::metadata::KeyValue;
use datafusion::parquet::file::properties::WriterProperties;
use datafusion::parquet::file::reader::{FileReader, SerializedFileReader};
use datafusion::parquet::file::writer::SerializedFileWriter;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::prelude::*;
use datafusion::scalar::ScalarValue;
use futures::AsyncWriteExt;
use std::fs::{create_dir_all, read_dir, File};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tempfile::TempDir;

/// We should disable page index support in the Parquet reader
/// when we ennable this feature, since we are using a custom index.
///
/// Example creating parquet file that
/// contains specialized indexes that
/// are ignored by other readers
///
/// ```text
///         ┌──────────────────────┐
///         │┌───────────────────┐ │
///         ││     DataPage      │ │      Standard Parquet
///         │└───────────────────┘ │      Data / pages
///         │┌───────────────────┐ │
///         ││     DataPage      │ │
///         │└───────────────────┘ │
///         │        ...           │
///         │                      │
///         │┌───────────────────┐ │
///         ││     DataPage      │ │
///         │└───────────────────┘ │
///         │┏━━━━━━━━━━━━━━━━━━━┓ │
///         │┃                   ┃ │        key/value metadata
///         │┃   Special Index   ┃◀┼────    that points at the
///         │┃                   ┃ │     │  special index
///         │┗━━━━━━━━━━━━━━━━━━━┛ │
///         │╔═══════════════════╗ │     │
///         │║                   ║ │
///         │║  Parquet Footer   ║ │     │  Footer includes
///         │║                   ║ ┼──────  thrift-encoded
///         │║                   ║ │        ParquetMetadata
///         │╚═══════════════════╝ │
///         └──────────────────────┘
///
///               Parquet File
/// ```
/// DistinctIndexTable is a custom TableProvider that reads Parquet files
#[derive(Debug)]
struct DistinctIndexTable {
    schema: SchemaRef,
    index: HashMap<String, HashSet<String>>,
    dir: PathBuf,
}

impl DistinctIndexTable {
    /// Scan a directory, read each file's footer metadata into a map
    fn try_new(dir: impl Into<PathBuf>, schema: SchemaRef) -> Result<Self> {
        let dir = dir.into();
        let mut index = HashMap::new();

        for entry in read_dir(&dir)? {
            let path = entry?.path();
            if path.extension().and_then(|s| s.to_str()) != Some("parquet") {
                continue;
            }
            let file_name = path.file_name().unwrap().to_string_lossy().to_string();

            let distinct_set = read_distinct_index(&path)?;

            println!("Read distinct index for {}: {:?}", file_name, distinct_set);
            index.insert(file_name, distinct_set);
        }

        Ok(Self { schema, index, dir })
    }
}

pub struct IndexedParquetWriter<W: Write + Seek> {
    writer: SerializedFileWriter<W>,
}

impl<W: Write + Seek + Send> IndexedParquetWriter<W> {
    pub fn try_new(
        sink: W,
        schema: Arc<Schema>,
        props: WriterProperties,
    ) -> Result<Self> {
        let schema_desc = ArrowSchemaConverter::new().convert(schema.as_ref())?;
        let props_ptr = Arc::new(props);
        let writer =
            SerializedFileWriter::new(sink, schema_desc.root_schema_ptr(), props_ptr)?;
        Ok(Self { writer })
    }
}

const INDEX_MAGIC: &[u8] = b"IDX1";

fn write_file_with_index(path: &Path, values: &[&str]) -> Result<()> {
    let field = Field::new("category", DataType::Utf8, false);
    let schema = Arc::new(Schema::new(vec![field.clone()]));
    let arr: ArrayRef = Arc::new(StringArray::from(values.to_vec()));
    let batch = RecordBatch::try_new(schema.clone(), vec![arr])?;

    let distinct: HashSet<_> = values.iter().copied().collect();
    let serialized = distinct.into_iter().collect::<Vec<_>>().join("\n");
    let index_bytes = serialized.into_bytes();

    let props = WriterProperties::builder().build();
    let file = File::create(path)?;

    let mut writer = IndexedParquetWriter::try_new(file, schema.clone(), props)?;

    {
        let mut rg_writer = writer.writer.next_row_group()?;
        let mut ser_col_writer = rg_writer
            .next_column()?
            .ok_or_else(|| ParquetError::General("No column writer".into()))?;

        let col_writer = ser_col_writer.typed::<ByteArrayType>();
        let values_bytes: Vec<ByteArray> = batch
            .column(0)
            .as_any()
            .downcast_ref::<StringArray>()
            .unwrap()
            .iter()
            .map(|opt| ByteArray::from(opt.unwrap()))
            .collect();

        println!("Writing values: {:?}", values_bytes);
        col_writer.write_batch(&values_bytes, None, None)?;
        ser_col_writer.close()?;
        rg_writer.close()?;
    }

    let offset = writer.writer.inner().seek(SeekFrom::Current(0))?;

    let index_len = index_bytes.len() as u64;
    writer.writer.inner().write_all(b"IDX1")?;
    writer.writer.inner().write_all(&index_len.to_le_bytes())?;

    writer.writer.inner().write_all(&index_bytes)?;

    writer.writer.append_key_value_metadata(KeyValue::new(
        "distinct_index_offset".to_string(),
        offset.to_string(),
    ));
    writer.writer.append_key_value_metadata(KeyValue::new(
        "distinct_index_length".to_string(),
        index_bytes.len().to_string(),
    ));

    writer.writer.close()?;

    println!("Finished writing file to {}", path.display());
    Ok(())
}

fn read_distinct_index(path: &Path) -> Result<HashSet<String>, ParquetError> {
    let mut file = File::open(path)?;

    let file_size = file.metadata()?.len();
    println!(
        "Reading index from {} (size: {})",
        path.display(),
        file_size
    );

    let reader = SerializedFileReader::new(file.try_clone()?)?;
    let meta = reader.metadata().file_metadata();

    let offset = meta
        .key_value_metadata()
        .and_then(|kvs| kvs.iter().find(|kv| kv.key == "distinct_index_offset"))
        .and_then(|kv| kv.value.as_ref())
        .ok_or_else(|| ParquetError::General("Missing index offset".into()))?
        .parse::<u64>()
        .map_err(|e| ParquetError::General(e.to_string()))?;

    let length = meta
        .key_value_metadata()
        .and_then(|kvs| kvs.iter().find(|kv| kv.key == "distinct_index_length"))
        .and_then(|kv| kv.value.as_ref())
        .ok_or_else(|| ParquetError::General("Missing index length".into()))?
        .parse::<usize>()
        .map_err(|e| ParquetError::General(e.to_string()))?;

    println!("Reading index at offset: {}, length: {}", offset, length);

    file.seek(SeekFrom::Start(offset))?;

    let mut magic_buf = [0u8; 4];
    file.read_exact(&mut magic_buf)?;
    if &magic_buf != INDEX_MAGIC {
        return Err(ParquetError::General("Invalid index magic".into()));
    }

    let mut len_buf = [0u8; 8];
    file.read_exact(&mut len_buf)?;
    let stored_len = u64::from_le_bytes(len_buf) as usize;

    if stored_len != length {
        return Err(ParquetError::General("Index length mismatch".into()));
    }

    let mut index_buf = vec![0u8; length];
    file.read_exact(&mut index_buf)?;

    let s =
        String::from_utf8(index_buf).map_err(|e| ParquetError::General(e.to_string()))?;

    Ok(s.lines().map(|s| s.to_string()).collect())
}

/// Implement TableProvider for DistinctIndexTable, using the distinct index to prune files
#[async_trait]
impl TableProvider for DistinctIndexTable {
    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
    fn table_type(&self) -> TableType {
        TableType::Base
    }

    /// Prune files before reading: only keep files whose distinct set contains the filter value
    async fn scan(
        &self,
        _ctx: &dyn Session,
        _proj: Option<&Vec<usize>>,
        filters: &[Expr],
        _limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        // Look for a single `category = 'X'` filter
        let mut target: Option<String> = None;

        if filters.len() == 1 {
            if let Expr::BinaryExpr(expr) = &filters[0] {
                if expr.op == Operator::Eq {
                    if let (
                        Expr::Column(c),
                        Expr::Literal(ScalarValue::Utf8(Some(v)), _),
                    ) = (&*expr.left, &*expr.right)
                    {
                        if c.name == "category" {
                            println!("Filtering for category: {v}");
                            target = Some(v.clone());
                        }
                    }
                }
            }
        }
        // Determine which files to scan
        let keep: Vec<String> = self
            .index
            .iter()
            .filter(|(_f, set)| target.as_ref().is_none_or(|v| set.contains(v)))
            .map(|(f, _)| f.clone())
            .collect();

        println!("Pruned files: {:?}", keep.clone());

        // Build ParquetSource for kept files
        let url = ObjectStoreUrl::parse("file://")?;
        let source = Arc::new(ParquetSource::default());
        let mut builder = FileScanConfigBuilder::new(url, self.schema.clone(), source);
        for file in keep {
            let path = self.dir.join(&file);
            let len = std::fs::metadata(&path)?.len();
            builder = builder.with_file(PartitionedFile::new(
                path.to_str().unwrap().to_string(),
                len,
            ));
        }
        Ok(DataSourceExec::from_data_source(builder.build()))
    }

    fn supports_filters_pushdown(
        &self,
        fs: &[&Expr],
    ) -> Result<Vec<TableProviderFilterPushDown>> {
        // Mark as inexact since pruning is file‑granular
        Ok(vec![TableProviderFilterPushDown::Inexact; fs.len()])
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    // 1. Create temp dir and write 3 Parquet files with different category sets
    let tmp = TempDir::new()?;
    let dir = tmp.path();
    create_dir_all(dir)?;
    write_file_with_index(&dir.join("a.parquet"), &["foo", "bar", "foo"])?;
    write_file_with_index(&dir.join("b.parquet"), &["baz", "qux"])?;
    write_file_with_index(&dir.join("c.parquet"), &["foo", "quux", "quux"])?;

    // 2. Register our custom TableProvider
    let field = Field::new("category", DataType::Utf8, false);
    let schema_ref = Arc::new(Schema::new(vec![field]));
    let provider = Arc::new(DistinctIndexTable::try_new(dir, schema_ref.clone())?);
    let ctx = SessionContext::new();

    ctx.register_table("t", provider)?;

    // 3. Run a query: only files containing 'foo' get scanned
    let df = ctx.sql("SELECT * FROM t").await?;
    df.show().await?;

    // 3. Run a query: only files containing 'foo' get scanned
    let df = ctx.sql("SELECT * FROM t WHERE category = 'foo'").await?;
    df.show().await?;

    Ok(())
}

[alamb]

This looks super cool @zhuqi-lucas

Try to using low level API, but it only works when we disable page index, if we setting page index, it will follow up the real row group data, and it conflicts with our embedding indexes.

I don't fully understand this concern -- I would probably have to play around with it some more

Are you willing to update this PR with this new example? I have some ideas on the various APIs we could use (like we could potentially encapsulate the index writing some more)

We could also then file a ticket upstream i arrow-rs with a description of what wasn't working with page indexes

[zhuqi-lucas]

Thank you @alamb , updated the code without page index using low level API, i will continue debugging the case that our self defined index with page index.

[alamb]

Thank you @zhuqi-lucas -- this is really neat. I left some thoughts -- let me know what you think

[zhuqi-lucas]

Thank you @alamb Currently, i disable the page index for reading, so this example will succeed, but if we enable page index, it will fail due to:

1. We are writing the self defined index just after the data.
2. But it seems, the page index offset info will write to the same place.
3. I can't find a solution until now, need some help.

Thanks!

[zhuqi-lucas]

Thank you @alamb,
I am excited to update today that i resolve the page index conflicts by adding new API in arrow-rs which can write bytes to the buf, and it can make the buf-wirtten metrics consistent, and the buf-wirtten will be used by page index also, so it's safe now, and i enable the page index now for the example, the testing result is good!

I am currently using this arrow-rs branch before the code merge:
apache/arrow-rs#7714

The example print logs, it's good, thanks!

Writing values: [ByteArray { data: "foo" }, ByteArray { data: "bar" }, ByteArray { data: "foo" }]
Writing custom index at offset: 68, length: 7
Finished writing file to /var/folders/q7/zjtv8rvx2hz0_t_rjjq8p9k00000gp/T/.tmp9zCIJt/a.parquet
Writing values: [ByteArray { data: "baz" }, ByteArray { data: "qux" }]
Writing custom index at offset: 68, length: 7
Finished writing file to /var/folders/q7/zjtv8rvx2hz0_t_rjjq8p9k00000gp/T/.tmp9zCIJt/b.parquet
Writing values: [ByteArray { data: "foo" }, ByteArray { data: "quux" }, ByteArray { data: "quux" }]
Writing custom index at offset: 70, length: 8
Finished writing file to /var/folders/q7/zjtv8rvx2hz0_t_rjjq8p9k00000gp/T/.tmp9zCIJt/c.parquet
Reading index from /var/folders/q7/zjtv8rvx2hz0_t_rjjq8p9k00000gp/T/.tmp9zCIJt/a.parquet (size: 363)
Reading index at offset: 68, length: 7
Read distinct index for a.parquet: "a.parquet"
Reading index from /var/folders/q7/zjtv8rvx2hz0_t_rjjq8p9k00000gp/T/.tmp9zCIJt/b.parquet (size: 363)
Reading index at offset: 68, length: 7
Read distinct index for b.parquet: "b.parquet"
Reading index from /var/folders/q7/zjtv8rvx2hz0_t_rjjq8p9k00000gp/T/.tmp9zCIJt/c.parquet (size: 368)
Reading index at offset: 70, length: 8
Read distinct index for c.parquet: "c.parquet"
Filtering for category: foo
Pruned files: ["c.parquet", "a.parquet"]
+----------+
| category |
+----------+
| foo      |
| foo      |
| foo      |
+----------+

[alamb]

The example print logs, it's good, thanks!

this is so cool!

[alamb]

Thank you @zhuqi-lucas -- this is (really) cool. It is definitely blog post worthy (we have too many cool things that are blog worthy recently - and not enough time to write the blogs!)

Anyhow I left some other suggestions and will prioritiize getting this PR in upstream

• Support write to buffer api for SerializedFileWriter arrow-rs#7714

[alamb]

This is amazing -- thank you

[alamb]

A small software engineering suggestion would be to encapsulate the custom distinct index into a struct to make it perhaps clearer what was specific to that index and what was required parquet plumbing

Something like

struct DistinctIndex {
  inner: HashSet<String>,
}

impl DistinctIndex {
  // serialize the distinct index to a writer
  fn serialize<W: Write>(&self) -> Result<()> { ... }

  // create a new distinct index from the specified bytes
  fn new_from_bytes(serialized: &[u8]) -> Result<Self> {.. }
}

[zhuqi-lucas]

Thank you @alamb! This is really a better way, perfect!

[alamb]

I wonder if the example would be simpler if it used ArrowWriter directly? It is not super clear to me why this example needs to use the lower level APIs directly

Or if it needs to use the lower level APIs it probably would be good to explain why you can't use the normal Arrow writer here directly

[zhuqi-lucas]

Thank you @alamb , this is a very good point, i was using low level because we don't have consistent buf API, but now we have, i will try to use ArrowWriter in follow-up try!

[alamb]

that is very cool

[alamb]

You can also potentially use PruningPredicate::literal_guarantee to do this analysis rather than repeating it here

However, doing this walk explicitly in the example might also be a good idea to show how it could be done generall

[alamb]

We shouldn't merge this PR until arrow-rs is released but otherwise it is good

[zhuqi-lucas]

Thank you @zhuqi-lucas -- this is (really) cool. It is definitely blog post worthy (we have too many cool things that are blog worthy recently - and not enough time to write the blogs!)

Anyhow I left some other suggestions and will prioritiize getting this PR in upstream

• Support write to buffer api for SerializedFileWriter arrow-rs#7714

Thank you @alamb , i can try to write a blog about this!

[2010YOUY01]

wow this is so cool!

I have a question (and I think it's worth adding to the comment for people like me that's not familiar with parquet internals):
How does it ensure that this extra index can be safely ignored by other readers? If another parquet reader implementation decides to do a sequential whole file scan, will it read into the extra custom data?

[zhuqi-lucas]

wow this is so cool!

I have a question (and I think it's worth adding to the comment for people like me that's not familiar with parquet internals): How does it ensure that this extra index can be safely ignored by other readers? If another parquet reader implementation decides to do a sequential whole file scan, will it read into the extra custom data?

Thank you for the review and great question, @2010YOUY01!

Short answer:
Because we append our custom index before the Parquet footer and never modify the existing metadata schema, Parquet readers will still:

1. Seek to the end of file and read the last 8 bytes, which consist of:
   • A 4‑byte little‑endian footer length
   • The magic marker PAR1
2. Jump back by that length to parse the Thrift‑encoded footer (and its key‑value list).

Any bytes you append ahead of the footer (i.e. after the data pages but before writing footer and magic) are simply skipped over by steps (1)&(2), because readers never scan from the file start—they always locate the footer via the trailer magic and length.

Why key/value metadata is safe:

• We only add two new keys (distinct_index_offset and distinct_index_length) into the existing footer metadata map.

• All standard readers will see unknown keys and either ignore them or surface them as “extra metadata,” but they will not attempt to deserialize our custom binary blob.

• On our side, we:

  1. Read the Parquet footer as usual.
  2. Extract our two key/value entries for offset+length.
  3. seek(offset) + read_exact(length) to load the custom index and deserialize it.

Because every compliant Parquet reader must interpret the PAR1 magic and footer length, none of them will ever “spill over” into our custom index or treat it as data pages.

I’ll add these details into the code comments. We’re also planning a blog post on Parquet indexing internals suggested by @alamb , thanks!

[alamb]

FYI @XiangpengHao and @@JigaoLuo -- here is another example of the somewhat crazy things you can do with parquet

[alamb]

How does it ensure that this extra index can be safely ignored by other readers? If another parquet reader implementation decides to do a sequential whole file scan, will it read into the extra custom data?

I agree with what @zhuqi-lucas says too

The way I think about this is that the parquet file's footer contains pointers (offsets) to the actual data in the file. There is no requirement that the footer points to all bytes within the file

There are other interesting things that can be done with this setup too (for example, concatenating parquet files together without having to re-encode the data (you can just copy the bytes around and rewrite the footer)

[zhuqi-lucas]

This is amazing @alamb ! Thanks!

There are other interesting things that can be done with this setup too (for example, concatenating parquet files together without having to re-encode the data (you can just copy the bytes around and rewrite the footer)