From 26016cbad5ff41f5808cf4eb447c7914bf4d7409 Mon Sep 17 00:00:00 2001 From: Anoop Narang Date: Fri, 9 Oct 2026 13:25:54 +0530 Subject: [PATCH] fix(optimizer): keep scans that read virtual columns on ParquetSource The liquid read path has no notion of DataFusion's virtual columns. It carries the `TableSchema` across but never produces one, so a scan whose projection includes a virtual column reads back a batch that lacks it, and resolving the column fails with an Arrow schema error naming only the file's own fields. `SELECT file_row_index() FROM t` is enough to hit it: projection pushdown rewrites the function into a row-number virtual column on the scan. Decline the swap for such scans, leaving them on `ParquetSource`, which derives virtual columns from the parquet reader. The guard keys on the virtual columns the scan actually reads (projection unioned with the pushed-down filter), not on what the table declares, so a table that declares a row-number column keeps the cache for queries that never read it. No projection at all reads the whole table schema, virtual columns included. Row numbers are absolute positions in the file; a shifted one would silently attach a positional delete or lineage record to the wrong row, so declining is the sound answer until the liquid reader can generate them itself. --- src/datafusion-local/src/tests/mod.rs | 1 + .../src/tests/virtual_column.rs | 70 ++++++ src/datafusion/src/optimizers/mod.rs | 201 +++++++++++++++++- 3 files changed, 270 insertions(+), 2 deletions(-) create mode 100644 src/datafusion-local/src/tests/virtual_column.rs diff --git a/src/datafusion-local/src/tests/mod.rs b/src/datafusion-local/src/tests/mod.rs index 8bc6ebf1..a08b76c5 100644 --- a/src/datafusion-local/src/tests/mod.rs +++ b/src/datafusion-local/src/tests/mod.rs @@ -24,6 +24,7 @@ mod date_optimizer; mod filter_limit; mod nested_filter; mod variants; +mod virtual_column; const TEST_FILE: &str = "../../examples/nano_hits.parquet"; const OPENOBSERVE_FILE: &str = "../../dev/test_parquet/openobserve.parquet"; diff --git a/src/datafusion-local/src/tests/virtual_column.rs b/src/datafusion-local/src/tests/virtual_column.rs new file mode 100644 index 00000000..2f0d9231 --- /dev/null +++ b/src/datafusion-local/src/tests/virtual_column.rs @@ -0,0 +1,70 @@ +//! A scan that reads a virtual column. +//! +//! `file_row_index()` is answered by the parquet reader: projection pushdown +//! rewrites it into a row-number virtual column on the scan's `TableSchema`, +//! which `ParquetSource` derives from the parquet reader. The query must return +//! the same rows with and without the cache. + +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::util::pretty::pretty_format_batches; +use arrow_schema::{DataType, Field, Schema}; +use datafusion::prelude::{ParquetReadOptions, SessionConfig, SessionContext}; +use parquet::arrow::ArrowWriter; +use tempfile::TempDir; + +use crate::LiquidCacheLocalBuilder; + +/// 12000 rows, more than the 8192-row default batch, so row indexes span batches. +fn write_t1(path: &Path) { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let id: Int64Array = (0..12000i64).map(|i| i * 10).collect::>().into(); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(id)]).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(); +} + +async fn register(ctx: &SessionContext, parquet: &Path) { + ctx.register_parquet( + "t1", + parquet.to_str().unwrap(), + ParquetReadOptions::default(), + ) + .await + .unwrap(); +} + +async fn run(ctx: &SessionContext, sql: &str) -> String { + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + pretty_format_batches(&batches).unwrap().to_string() +} + +#[tokio::test] +async fn file_row_index_matches_datafusion() { + let dir = TempDir::new().unwrap(); + let parquet = dir.path().join("t1.parquet"); + write_t1(&parquet); + + let plain = SessionContext::new(); + register(&plain, &parquet).await; + + let cache_dir = dir.path().join("cache"); + std::fs::create_dir_all(&cache_dir).unwrap(); + let (liquid, _cache) = LiquidCacheLocalBuilder::new() + .with_cache_dir(cache_dir) + .build(SessionConfig::new()) + .await + .unwrap(); + register(&liquid, &parquet).await; + + let sql = "SELECT file_row_index() AS pos, id FROM t1 \ + WHERE id % 997 = 0 ORDER BY pos"; + let expected = run(&plain, sql).await; + // Twice: the second run is served from the cache when the scan is cached. + assert_eq!(run(&liquid, sql).await, expected); + assert_eq!(run(&liquid, sql).await, expected); +} diff --git a/src/datafusion/src/optimizers/mod.rs b/src/datafusion/src/optimizers/mod.rs index 7e30acde..d8950b11 100644 --- a/src/datafusion/src/optimizers/mod.rs +++ b/src/datafusion/src/optimizers/mod.rs @@ -2,13 +2,18 @@ mod lineage; -use std::sync::Arc; +use std::{collections::HashSet, sync::Arc}; use datafusion::{ catalog::memory::DataSourceExec, common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}, config::ConfigOptions, - datasource::{physical_plan::ParquetSource, source::DataSource}, + datasource::{ + physical_plan::{FileSource, ParquetSource}, + source::DataSource, + table_schema::TableSchema, + }, + physical_expr::{PhysicalExpr, projection::ProjectionExprs, utils::collect_columns}, physical_optimizer::PhysicalOptimizerRule, physical_plan::ExecutionPlan, }; @@ -107,6 +112,73 @@ pub fn rewrite_data_source_plan( rewrite_data_source_plan_with_hints(plan, cache, &ColumnLineages::default()) } +/// The virtual columns this scan reads that the liquid path cannot produce. +/// +/// The liquid read path has no notion of DataFusion's virtual columns. It carries +/// the [`TableSchema`] across faithfully but never produces one, so a scan that +/// reads a virtual column gets back a batch which simply lacks it, and a predicate +/// over one cannot be rewritten against the file schemas either. Such a scan has +/// to stay on `ParquetSource`, which derives virtual columns from the parquet +/// reader. +/// +/// Declining a scan costs it the cache, so this stays narrow: a virtual column the +/// scan actually reads, not the mere presence of one on the table. A provider that +/// declares a row-position column on every table keeps the cache for the queries +/// that never project one. No projection at all is the one broad case, and it is +/// not a guess — the scan then reads the whole table schema, virtual columns +/// included. +/// +/// Positional reads are what reaches here: applying positional deletes, and row +/// lineage, both project a reader-produced physical row position, which is an +/// absolute index into the file. A plausible but shifted position would associate +/// a delete with the wrong row, so declining is the sound answer while the liquid +/// reader cannot generate positions itself. +fn unproducible_virtual_columns( + table_schema: &TableSchema, + projection: Option<&ProjectionExprs>, + filter: Option<&Arc>, +) -> Option { + let virtual_columns = table_schema.virtual_columns(); + if virtual_columns.is_empty() { + return None; + } + + let mut needed: Vec<&str> = Vec::new(); + match projection { + None => needed.extend(virtual_columns.iter().map(|field| field.name().as_str())), + Some(projection) => { + let mut read: HashSet = projection + .expr_iter() + .flat_map(|expr| collect_columns(&expr)) + .map(|column| column.name().to_string()) + .collect(); + // The pushed-down predicate is rewritten against the file schemas in + // the reader too, so a conjunct over a virtual column fails exactly as + // a projection over one does. The row filter's own check cannot catch + // it: that resolves against the table schema, which does hold the + // virtual columns. + if let Some(filter) = filter { + read.extend( + collect_columns(filter) + .into_iter() + .map(|column| column.name().to_string()), + ); + } + needed.extend( + virtual_columns + .iter() + .map(|field| field.name().as_str()) + .filter(|name| read.contains(*name)), + ); + } + } + + if needed.is_empty() { + return None; + } + Some(needed.join("`, `")) +} + /// If `node` is a `DataSourceExec` over a `ParquetSource`, return an equivalent /// node backed by [`LiquidParquetSource`] carrying `hints`. fn convert_parquet_scan( @@ -119,6 +191,20 @@ fn convert_parquet_scan( let (file_scan_config, parquet_source) = data_source_exec.downcast_to_file_source::()?; + let pushed_filter = parquet_source.filter(); + if let Some(names) = unproducible_virtual_columns( + parquet_source.table_schema(), + parquet_source.projection(), + pushed_filter.as_ref(), + ) { + // At info: this silently turns the cache off for a scan, and the only + // symptom is that queries stop getting faster. + log::info!( + "liquid_cache scan BYPASS: the read path cannot produce virtual column(s) `{names}`" + ); + return None; + } + let new_source = LiquidParquetSource::from_parquet_source(parquet_source.clone(), cache.clone()) .with_lineages(Arc::new(hints)) @@ -341,4 +427,115 @@ mod tests { }; assert_eq!(pruning_metrics.pruned(), 1); } + + /// Declining a scan costs it the cache, so the guard must key on what the scan + /// reads, not on what the table declares. A row-position column present on the + /// table but absent from the projection keeps the cache; no projection at all + /// reads the whole table schema and does not. + #[test] + fn only_a_virtual_column_the_scan_reads_costs_the_cache() { + use arrow_schema::Fields; + use datafusion::physical_expr::expressions::{BinaryExpr, col, lit}; + use datafusion::physical_expr::projection::ProjectionExpr; + + let file_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("val", DataType::Int64, true), + ])); + let row_pos = Field::new("row_number", DataType::Int64, false) + .with_extension_type(parquet::arrow::RowNumber); + + // No virtual columns at all: nothing to refuse, whatever the projection. + let plain = TableSchema::builder(Arc::clone(&file_schema)).build(); + assert_eq!(unproducible_virtual_columns(&plain, None, None), None); + + let positional = TableSchema::builder(Arc::clone(&file_schema)) + .with_virtual_columns(Fields::from(vec![row_pos.clone()])) + .build(); + let full = positional.table_schema(); + + // Declared but not read: still cached. + let only_id = + ProjectionExprs::new(vec![ProjectionExpr::new(col("id", full).unwrap(), "id")]); + assert_eq!( + unproducible_virtual_columns(&positional, Some(&only_id), None), + None + ); + + // Read: refused, and named. + let reads_pos = ProjectionExprs::new(vec![ProjectionExpr::new( + col("row_number", full).unwrap(), + "row_number", + )]); + assert_eq!( + unproducible_virtual_columns(&positional, Some(&reads_pos), None).as_deref(), + Some("row_number") + ); + + // No projection reads the whole table schema, virtual columns included. + assert_eq!( + unproducible_virtual_columns(&positional, None, None).as_deref(), + Some("row_number") + ); + + // Read only by the pushed-down predicate: refused too. The reader rewrites + // the predicate against the file schemas as well, and the row filter's own + // check passes it because that resolves against the table schema. + let pos_predicate: Arc = Arc::new(BinaryExpr::new( + col("row_number", full).unwrap(), + Operator::Gt, + lit(0i64), + )); + assert_eq!( + unproducible_virtual_columns(&positional, Some(&only_id), Some(&pos_predicate)) + .as_deref(), + Some("row_number") + ); + } + + /// The guard at its call site, with an ordinary scan as a positive control, so + /// deleting it from `convert_parquet_scan` fails a test instead of silently + /// restoring a plan that cannot execute. + #[tokio::test] + async fn a_scan_reading_a_virtual_column_stays_on_parquet_source() { + use arrow_schema::Fields; + use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::physical_plan::FileScanConfigBuilder; + use datafusion::execution::object_store::ObjectStoreUrl; + + let file_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + + let scan = |table_schema: TableSchema| -> Arc { + let source = Arc::new(ParquetSource::new(table_schema)) as Arc; + let config = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(PartitionedFile::new("t.parquet", 16)) + .build(); + Arc::new(DataSourceExec::new(Arc::new(config))) + }; + + let tmp_dir = tempfile::tempdir().unwrap(); + let cache = make_cache(tmp_dir.path()).await; + + // Positive control: an ordinary scan is still handed to the cache. + let plain = scan(TableSchema::builder(Arc::clone(&file_schema)).build()); + assert!( + convert_parquet_scan(&plain, &cache, ColumnLineages::default(), true).is_some(), + "an ordinary scan must still convert to the liquid source" + ); + + // `ParquetSource::new` projects the whole table schema, so the positional + // scan reads the row-position column and cannot be served. + let positional = scan( + TableSchema::builder(file_schema) + .with_virtual_columns(Fields::from(vec![ + Field::new("row_number", DataType::Int64, false) + .with_extension_type(parquet::arrow::RowNumber), + ])) + .build(), + ); + assert!( + convert_parquet_scan(&positional, &cache, ColumnLineages::default(), true).is_none(), + "a scan reading a virtual column must stay on ParquetSource" + ); + } }