chore: sync upstream (virtual-column guard, datafusion-contrib/liquid-cache#523)

src/datafusion-local/src/tests/mod.rs

@@ -26,6 +26,7 @@ mod filter_limit;
 mod nested_filter;
 mod unevaluable_conjunct;
 mod variants;
+mod virtual_column;
 
 const TEST_FILE: &str = "../../examples/nano_hits.parquet";
 const OPENOBSERVE_FILE: &str = "../../dev/test_parquet/openobserve.parquet";

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::<Vec<_>>().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);
+}

src/datafusion/src/optimizers/mod.rs

@@ -2,8 +2,7 @@
 
 mod lineage;
 
-use std::collections::HashSet;
-use std::sync::Arc;
+use std::{collections::HashSet, sync::Arc};
 
 use datafusion::{
     arrow::datatypes::SchemaRef,
@@ -622,9 +621,8 @@ fn convert_parquet_scan(
         parquet_source.projection(),
         pushed_filter.as_ref(),
     ) {
-        // At info, like the admission gate's BYPASS line: this silently turns the
-        // cache off for a scan, and the only symptom is that queries stop getting
-        // faster.
+        // 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}`"
         );
@@ -731,116 +729,6 @@ mod tests {
             .unwrap()
     }
 
-    /// 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("__ducklake_row_pos", DataType::Int64, true);
-
-        // 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("__ducklake_row_pos", full).unwrap(),
-            "__ducklake_row_pos",
-        )]);
-        assert_eq!(
-            unproducible_virtual_columns(&positional, Some(&reads_pos), None).as_deref(),
-            Some("__ducklake_row_pos")
-        );
-
-        // No projection reads the whole table schema, virtual columns included.
-        assert_eq!(
-            unproducible_virtual_columns(&positional, None, None).as_deref(),
-            Some("__ducklake_row_pos")
-        );
-
-        // 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<dyn PhysicalExpr> = Arc::new(BinaryExpr::new(
-            col("__ducklake_row_pos", full).unwrap(),
-            Operator::Gt,
-            lit(0i64),
-        ));
-        assert_eq!(
-            unproducible_virtual_columns(&positional, Some(&only_id), Some(&pos_predicate))
-                .as_deref(),
-            Some("__ducklake_row_pos")
-        );
-    }
-
-    /// 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::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<dyn ExecutionPlan> {
-            let source = Arc::new(ParquetSource::new(table_schema)) as Arc<dyn FileSource>;
-            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(
-                    "__ducklake_row_pos",
-                    DataType::Int64,
-                    true,
-                )]))
-                .build(),
-        );
-        assert!(
-            convert_parquet_scan(&positional, &cache, ColumnLineages::default(), true).is_none(),
-            "a scan reading a virtual column must stay on ParquetSource"
-        );
-    }
-
     /// The admission gate bypasses a scan whose estimated footprint exceeds the
     /// budget (large expansion here forces that), and caches one that fits.
     #[tokio::test]
@@ -1160,6 +1048,117 @@ 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<dyn PhysicalExpr> = 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<dyn ExecutionPlan> {
+            let source = Arc::new(ParquetSource::new(table_schema)) as Arc<dyn FileSource>;
+            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"
+        );
+    }
 }
 
 /// Pure unit tests for the footprint byte-math (no cache / no t4 mount, so they