From 8225b6883e9f56979be3d24e9ff920e5a1380d1c Mon Sep 17 00:00:00 2001 From: Anoop Date: Fri, 9 Oct 2026 13:58:38 +0530 Subject: [PATCH] fix(optimizer): keep scans that read virtual columns on ParquetSource (#523) Queries that read a DataFusion virtual column fail once the liquid cache is on. ```sql SELECT file_row_index() AS pos, id FROM t ``` Plain DataFusion returns each row's position in its file. Through `LiquidCacheLocalBuilder` it errors: ``` SchemaError("Unable to get field named \"__datafusion_file_row_index\". Valid fields: [\"id\"]") ``` ## Why `file_row_index()` is answered by the parquet reader: projection pushdown rewrites it into a row-number virtual column (arrow's `RowNumber` extension type) on the scan's `TableSchema`, and `ParquetSource` asks the parquet reader to produce it. Table providers can declare such columns directly too, for example to apply positional deletes. `LocalModeOptimizer` swaps every parquet scan to `LiquidParquetSource`. The liquid read path carries the `TableSchema` across but never produces a virtual column, so the batch comes back without it and resolving the column fails. ## The fix `convert_parquet_scan` declines the swap when the scan reads a virtual column, so the scan stays on `ParquetSource` and logs a `BYPASS` line at info. The check covers what the scan actually reads (the projection, unioned with the pushed-down filter), not what the table declares, so a table carrying a row-number column keeps the cache for queries that never read it. A scan with no projection reads the whole table schema and is declined. Declining rather than producing positions in the liquid reader is deliberate for now: a row number is an absolute position in the file, and a plausible-but-shifted one would silently attach a delete or lineage record to the wrong row. Generating them in the liquid reader would be the follow-up that lets these scans use the cache. The filter half of the check is defensive. DataFusion 55's `ParquetSource::try_pushdown_filters` already refuses predicates that reference a virtual column, so today it should not trigger; it covers a source that does push one down, since the liquid reader would fail rewriting it against the file schemas. ## Tests - `datafusion-local/src/tests/virtual_column.rs`: runs the `file_row_index()` query above (12000 rows, spanning batches) through plain DataFusion and through the liquid cache, twice, and compares results. - `optimizers` unit tests: the guard keys on what is read, not what is declared (projection, filter, and no-projection cases), and `convert_parquet_scan` keeps a virtual-column scan on `ParquetSource` with an ordinary scan as a positive control. With only the new test applied to `main`: ``` $ cargo test -p liquid-cache-datafusion-local virtual_column test tests::virtual_column::file_row_index_matches_datafusion ... FAILED called `Result::unwrap()` on an `Err` value: ArrowError(SchemaError("Unable to get field named \"__datafusion_file_row_index\". Valid fields: [\"id\"]"), Some("")) ``` With the fix, `liquid-cache-datafusion` (62) and `liquid-cache-datafusion-local` (32) pass, and clippy is clean with `-D warnings`. --- 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" + ); + } }