diff --git a/crates/core/src/table/mod.rs b/crates/core/src/table/mod.rs index 7bb54d2c..f4c3d9c3 100644 --- a/crates/core/src/table/mod.rs +++ b/crates/core/src/table/mod.rs @@ -795,8 +795,8 @@ impl Table { /// Snapshot streams batches as they are read from each file slice. Incremental /// streaming is not yet supported and returns an `Unsupported` error. /// - /// For MOR tables with log files, streaming falls back to a collect-and-merge that - /// yields the merged result as a single batch. + /// For MOR file slices with log files, streaming falls back to a collect-and-merge + /// that yields that file slice's merged result as a single batch. /// /// # Example /// ```ignore diff --git a/crates/core/tests/table_read_tests.rs b/crates/core/tests/table_read_tests.rs index bbfa76b3..0eede116 100644 --- a/crates/core/tests/table_read_tests.rs +++ b/crates/core/tests/table_read_tests.rs @@ -22,11 +22,11 @@ //! organized by table version (v6, v8+) and query type. use arrow::compute::concat_batches; -use arrow_array::RecordBatch; +use arrow_array::{Int64Array, RecordBatch, StringArray}; use futures::StreamExt; use futures::stream::BoxStream; use hudi_core::config::read::HudiReadConfig; -use hudi_core::error::Result; +use hudi_core::error::{CoreError, Result}; use hudi_core::table::{QueryType, ReadOptions, Table}; use hudi_test::{QuickstartTripsTable, SampleTable}; @@ -40,6 +40,80 @@ async fn collect_stream_batches( Ok(batches) } +fn sample_rows(records: &[RecordBatch]) -> Result> { + if records.is_empty() { + return Ok(Vec::new()); + } + let schema = records[0].schema(); + let batch = concat_batches(&schema, records)?; + Ok(SampleTable::sample_data_order_by_id(&batch) + .into_iter() + .map(|(id, name, is_active)| (id, name.to_string(), is_active)) + .collect()) +} + +fn txn_rows(records: &[RecordBatch]) -> Vec<(String, String, i64)> { + let mut rows = Vec::new(); + for record_batch in records { + let txn_ids = record_batch + .column_by_name("txn_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let txn_types = record_batch + .column_by_name("txn_type") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let txn_timestamps = record_batch + .column_by_name("txn_ts") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + for row_idx in 0..record_batch.num_rows() { + rows.push(( + txn_ids.value(row_idx).to_string(), + txn_types.value(row_idx).to_string(), + txn_timestamps.value(row_idx), + )); + } + } + rows.sort_unstable(); + rows +} + +fn assert_batch_size_respected( + batches: &[RecordBatch], + batch_size: usize, + expected_rows: usize, + context: &str, +) { + assert!( + !batches.is_empty(), + "{context}: expected at least one batch" + ); + assert!( + batches.iter().all(|batch| batch.num_rows() <= batch_size), + "{context}: each batch must have at most {batch_size} rows, got {:?}", + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>() + ); + if expected_rows > batch_size { + assert!( + batches.len() > 1, + "{context}: batch_size={batch_size} should split {expected_rows} rows" + ); + } + let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(total_rows, expected_rows, "{context}: total rows"); +} + /// Test helper module for v6 tables (pre-1.0 spec) mod v6_tables { use super::*; @@ -195,6 +269,30 @@ mod v6_tables { Ok(()) } + #[tokio::test] + async fn test_partition_filter_reduces_file_slices() -> Result<()> { + let base_url = SampleTable::V6SimplekeygenNonhivestyle.url_to_cow(); + let hudi_table = Table::new(base_url.path()).await?; + + let all_slices = hudi_table.get_file_slices(&ReadOptions::new()).await?; + let filtered_options = ReadOptions::new().with_filters([("byteField", "=", "10")])?; + let filtered_slices = hudi_table.get_file_slices(&filtered_options).await?; + + assert!( + filtered_slices.len() < all_slices.len(), + "partition filter must reduce file slices: {} filtered vs {} unfiltered", + filtered_slices.len(), + all_slices.len() + ); + + let records = hudi_table.read(&filtered_options).await?; + let schema = records[0].schema(); + let records = concat_batches(&schema, &records)?; + let sample_data = SampleTable::sample_data_order_by_id(&records); + assert_eq!(sample_data, vec![(1, "Alice", false), (3, "Carol", true)]); + Ok(()) + } + #[tokio::test] async fn test_simple_keygen_hivestyle_no_metafields() -> Result<()> { for base_url in SampleTable::V6SimplekeygenHivestyleNoMetafields.urls() { @@ -770,21 +868,15 @@ mod v8_tables { let base_url = SampleTable::V8Nonpartitioned.url_to_cow(); let hudi_table = Table::new(base_url.path()).await?; - // Request small batch size let options = ReadOptions::new().with_batch_size(1)?; let mut stream = hudi_table.read_stream(&options).await?; - // Collect all batches from stream let mut batches = Vec::new(); while let Some(result) = stream.next().await { batches.push(result?); } - // With batch_size=1 and 4 rows, we expect multiple batches, but the - // exact number depends on both the batch_size setting and the Parquet - // file's internal row group structure. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, 4, "Total rows should match expected count"); + assert_batch_size_respected(&batches, 1, 4, "V8 snapshot stream batch_size=1"); Ok(()) } @@ -834,6 +926,7 @@ mod v8_tables { .create_file_group_reader_with_options(None, std::iter::empty::<(&str, &str)>())?; let options = ReadOptions::new(); let file_slice = &file_slices[0]; + let expected = fg_reader.read_file_slice(file_slice, &options).await?; let mut stream = fg_reader .read_file_slice_stream(file_slice, &options) .await?; @@ -845,10 +938,11 @@ mod v8_tables { } assert!(!batches.is_empty(), "Should produce at least one batch"); - - // Verify we got records - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!(total_rows > 0, "Should read at least one row"); + assert_eq!( + sample_rows(&[expected])?, + sample_rows(&batches)?, + "file-slice streaming should match eager file-slice read" + ); Ok(()) } @@ -873,8 +967,7 @@ mod v8_tables { batches.push(result?); } - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, 4, "Should read all 4 rows"); + assert_batch_size_respected(&batches, 1, 4, "V8 file-slice stream batch_size=1"); Ok(()) } @@ -886,6 +979,7 @@ mod v8_tables { let hudi_table = Table::new(base_url.path()).await?; let options = ReadOptions::new(); + let file_slices = hudi_table.get_file_slices(&options).await?; let mut stream = hudi_table.read_stream(&options).await?; let mut batches = Vec::new(); @@ -894,7 +988,11 @@ mod v8_tables { } assert!(!batches.is_empty(), "Should produce batches from MOR table"); - + assert_eq!( + batches.len(), + file_slices.len(), + "MOR snapshot stream should emit one merged batch per file slice" + ); // Verify total row count - should have 6 rows (8 inserts - 2 deletes) let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total_rows, 6, "Should have 6 rows (8 inserts - 2 deleted)"); @@ -922,41 +1020,6 @@ mod v8_tables { /// Test helper module for v9 tables (1.1 spec) mod v9_tables { use super::*; - use arrow_array::{Int64Array, StringArray}; - - fn txn_rows(records: &[RecordBatch]) -> Vec<(String, String, i64)> { - let mut rows = Vec::new(); - for record_batch in records { - let txn_ids = record_batch - .column_by_name("txn_id") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let txn_types = record_batch - .column_by_name("txn_type") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - let txn_timestamps = record_batch - .column_by_name("txn_ts") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - - for i in 0..record_batch.num_rows() { - rows.push(( - txn_ids.value(i).to_string(), - txn_types.value(i).to_string(), - txn_timestamps.value(i), - )); - } - } - rows.sort_unstable(); - rows - } fn read_optimized_options() -> ReadOptions { ReadOptions::new().with_hudi_option(HudiReadConfig::UseReadOptimizedMode.as_ref(), "true") @@ -1440,9 +1503,8 @@ mod v9_tables { let options = ReadOptions::new().with_batch_size(1)?; let stream = hudi_table.read_stream(&options).await?; let batches = collect_stream_batches(stream).await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, 3, "Total rows should match expected count"); + assert_batch_size_respected(&batches, 1, 3, "V9 snapshot stream batch_size=1"); Ok(()) } @@ -1498,13 +1560,17 @@ mod v9_tables { .create_file_group_reader_with_options(None, std::iter::empty::<(&str, &str)>())?; let options = ReadOptions::new(); let file_slice = &file_slices[0]; + let expected = fg_reader.read_file_slice(file_slice, &options).await?; let stream = fg_reader .read_file_slice_stream(file_slice, &options) .await?; let batches = collect_stream_batches(stream).await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!(total_rows > 0, "Should read at least one row"); + assert_eq!( + txn_rows(&[expected]), + txn_rows(&batches), + "file-slice streaming should match eager file-slice read" + ); Ok(()) } @@ -1518,14 +1584,21 @@ mod v9_tables { let fg_reader = hudi_table .create_file_group_reader_with_options(None, std::iter::empty::<(&str, &str)>())?; + let expected = fg_reader + .read_file_slice(file_slice, &ReadOptions::new()) + .await?; let options = ReadOptions::new().with_batch_size(1)?; let stream = fg_reader .read_file_slice_stream(file_slice, &options) .await?; let batches = collect_stream_batches(stream).await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!(total_rows > 0, "Should read at least one row"); + assert_batch_size_respected( + &batches, + 1, + expected.num_rows(), + "V9 file-slice stream batch_size=1", + ); Ok(()) } @@ -1631,16 +1704,11 @@ mod streaming_queries { let base_url = SampleTable::V6Nonpartitioned.url_to_cow(); let hudi_table = Table::new(base_url.path()).await?; - // Request small batch size let options = ReadOptions::new().with_batch_size(1)?; let stream = hudi_table.read_stream(&options).await?; let batches = collect_stream_batches(stream).await?; - // With batch_size=1 and 4 rows, we expect multiple batches, but the - // exact number depends on both the batch_size setting and the Parquet - // file's internal row group structure. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, 4, "Total rows should match expected count"); + assert_batch_size_respected(&batches, 1, 4, "V6 snapshot stream batch_size=1"); Ok(()) } @@ -1685,16 +1753,18 @@ mod streaming_queries { .create_file_group_reader_with_options(None, std::iter::empty::<(&str, &str)>())?; let options = ReadOptions::new(); let file_slice = &file_slices[0]; + let expected = fg_reader.read_file_slice(file_slice, &options).await?; let stream = fg_reader .read_file_slice_stream(file_slice, &options) .await?; let batches = collect_stream_batches(stream).await?; assert!(!batches.is_empty(), "Should produce at least one batch"); - - // Verify we got records - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert!(total_rows > 0, "Should read at least one row"); + assert_eq!( + sample_rows(&[expected])?, + sample_rows(&batches)?, + "file-slice streaming should match eager file-slice read" + ); Ok(()) } @@ -1708,15 +1778,13 @@ mod streaming_queries { let fg_reader = hudi_table .create_file_group_reader_with_options(None, std::iter::empty::<(&str, &str)>())?; - // Test with small batch size let options = ReadOptions::new().with_batch_size(1)?; let stream = fg_reader .read_file_slice_stream(file_slice, &options) .await?; let batches = collect_stream_batches(stream).await?; - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - assert_eq!(total_rows, 4, "Should read all 4 rows"); + assert_batch_size_respected(&batches, 1, 4, "V6 file-slice stream batch_size=1"); Ok(()) } @@ -1727,17 +1795,42 @@ mod streaming_queries { let hudi_table = Table::new(base_url.path()).await?; let options = ReadOptions::new(); + let file_slices = hudi_table.get_file_slices(&options).await?; let stream = hudi_table.read_stream(&options).await?; let batches = collect_stream_batches(stream).await?; assert!(!batches.is_empty(), "Should produce batches from MOR table"); - + assert_eq!( + batches.len(), + file_slices.len(), + "MOR snapshot stream should emit one merged batch per file slice" + ); // Verify total row count let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total_rows, 8, "Should have 8 rows (8 inserts)"); Ok(()) } + #[tokio::test] + async fn test_read_stream_incremental_returns_unsupported() -> Result<()> { + let base_url = SampleTable::V6Nonpartitioned.url_to_cow(); + let hudi_table = Table::new(base_url.path()).await?; + + let err = match hudi_table + .read_stream(&ReadOptions::new().with_query_type(QueryType::Incremental)) + .await + { + Ok(_) => panic!("incremental streaming should be unsupported"), + Err(err) => err, + }; + + assert!( + matches!(err, CoreError::Unsupported(_)), + "expected Unsupported for incremental streaming, got: {err}" + ); + Ok(()) + } + #[tokio::test] async fn test_read_snapshot_stream_with_projection() -> Result<()> { let base_url = SampleTable::V6Nonpartitioned.url_to_cow(); @@ -1876,9 +1969,9 @@ mod streaming_queries { #[tokio::test] async fn test_read_snapshot_stream_with_as_of_timestamp() -> Result<()> { - // Cross-validate streaming time-travel against the eager API: with the - // same `as_of_timestamp`, both must return the same row count. (Before - // the fix, streaming silently used the latest commit and could diverge.) + // Cross-validate streaming time-travel against the eager API. The + // fixture intentionally has fewer rows at the first commit than at the + // latest commit, so a stream that silently reads latest would diverge. let base_url = SampleTable::V6Nonpartitioned.url_to_cow(); let hudi_table = Table::new(base_url.path()).await?; let first_commit = hudi_table @@ -1891,12 +1984,21 @@ mod streaming_queries { let options = ReadOptions::new().with_as_of_timestamp(&first_commit); let eager = hudi_table.read(&options).await?; - let eager_rows: usize = eager.iter().map(|b| b.num_rows()).sum(); + let eager_rows = sample_rows(&eager)?; + + let latest = hudi_table.read(&ReadOptions::new()).await?; + let latest_rows = sample_rows(&latest)?; let stream = hudi_table.read_stream(&options).await?; let stream_batches = collect_stream_batches(stream).await?; - let stream_rows: usize = stream_batches.iter().map(|b| b.num_rows()).sum(); + let stream_rows = sample_rows(&stream_batches)?; + assert_eq!(eager_rows.len(), 3, "first commit should have 3 rows"); + assert_ne!( + eager_rows.len(), + latest_rows.len(), + "fixture must differ between first commit and latest snapshot" + ); assert_eq!(eager_rows, stream_rows); Ok(()) } @@ -2175,11 +2277,10 @@ mod manual_reader_matches_table_read { let expected = hudi_table.read(&options).await?; let actual = read_via_manual_reader(&hudi_table, &options).await?; - let expected_rows: usize = expected.iter().map(|b| b.num_rows()).sum(); - let actual_rows: usize = actual.iter().map(|b| b.num_rows()).sum(); assert_eq!( - expected_rows, actual_rows, - "incremental (explicit range up to first commit): row counts must match" + txn_rows(&expected), + txn_rows(&actual), + "incremental (explicit range up to first commit): rows must match" ); let options = ReadOptions::new() @@ -2190,11 +2291,10 @@ mod manual_reader_matches_table_read { let expected = hudi_table.read(&options).await?; let actual = read_via_manual_reader(&hudi_table, &options).await?; - let expected_rows: usize = expected.iter().map(|b| b.num_rows()).sum(); - let actual_rows: usize = actual.iter().map(|b| b.num_rows()).sum(); assert_eq!( - expected_rows, actual_rows, - "incremental (first..second commit): row counts must match" + txn_rows(&expected), + txn_rows(&actual), + "incremental (first..second commit): rows must match" ); Ok(()) @@ -2218,11 +2318,10 @@ mod manual_reader_matches_table_read { let expected = hudi_table.read(&options).await?; let actual = read_via_manual_reader(&hudi_table, &options).await?; - let expected_rows: usize = expected.iter().map(|b| b.num_rows()).sum(); - let actual_rows: usize = actual.iter().map(|b| b.num_rows()).sum(); assert_eq!( - expected_rows, actual_rows, - "incremental (start only, end defaults to latest): row counts must match" + txn_rows(&expected), + txn_rows(&actual), + "incremental (start only, end defaults to latest): rows must match" ); Ok(()) @@ -2563,10 +2662,9 @@ mod lance_tables { #[tokio::test] async fn test_v9_lance_nonpartitioned_cow_read_uses_extension_fallback_without_format_config() -> Result<()> { - // Safe because each `path_to_cow()` call extracts the zip into its own - // tempdir (see `hudi_test::extract_test_table`); the in-place edit does - // not leak into other tests. - let table_path = SampleTable::V9LanceNonpartitioned.path_to_cow(); + // Safe because this test asks for a fresh extraction before mutating + // hoodie.properties in place. + let table_path = SampleTable::V9LanceNonpartitioned.path_to_cow_fresh(); let props_path = std::path::Path::new(&table_path).join(".hoodie/hoodie.properties"); let props = std::fs::read_to_string(&props_path).unwrap(); let props_without_format = props @@ -2640,6 +2738,32 @@ mod mdt_enabled_tables { Ok(()) } + #[tokio::test] + async fn test_mdt_read_matches_nometa_read() -> Result<()> { + let mdt_url = SampleTable::V9TxnsNonpartMeta.url_to_mor_avro(); + let mdt_table = Table::new(mdt_url.path()).await?; + assert!( + mdt_table.is_metadata_table_enabled(), + "metadata-enabled fixture should use MDT listing" + ); + + let nometa_url = SampleTable::V9TxnsNonpartNometa.url_to_mor_avro(); + let nometa_table = Table::new(nometa_url.path()).await?; + assert!( + !nometa_table.is_metadata_table_enabled(), + "nometa fixture should use storage listing" + ); + + let mdt_rows = txn_rows(&mdt_table.read(&ReadOptions::new()).await?); + let nometa_rows = txn_rows(&nometa_table.read(&ReadOptions::new()).await?); + + assert_eq!( + mdt_rows, nometa_rows, + "MDT-backed read should return the same rows as storage listing" + ); + Ok(()) + } + /// Test MDT partition key normalization for non-partitioned tables. /// The metadata table stores "." as partition key, but external API should see "". /// For non-partitioned tables, we use a fast path that directly fetches "." without diff --git a/crates/datafusion/Cargo.toml b/crates/datafusion/Cargo.toml index ed99be4f..6e6bb333 100644 --- a/crates/datafusion/Cargo.toml +++ b/crates/datafusion/Cargo.toml @@ -50,6 +50,7 @@ datafusion-physical-expr = { workspace = true } # runtime / async async-trait = { workspace = true } +futures = { workspace = true } tokio = { workspace = true } url = { workspace = true } @@ -58,3 +59,4 @@ log = { workspace = true } [dev-dependencies] hudi-test = { path = "../test", features = ["datafusion"] } +parquet = { workspace = true } diff --git a/crates/datafusion/src/hudi_exec.rs b/crates/datafusion/src/hudi_exec.rs new file mode 100644 index 00000000..ee57d9df --- /dev/null +++ b/crates/datafusion/src/hudi_exec.rs @@ -0,0 +1,612 @@ +/* + * 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. + */ +//! Custom DataFusion execution plan for reading Hudi tables through +//! [`FileGroupReader`], supporting all base file formats and MOR log merging. + +use std::any::Any; +use std::fmt; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + SendableRecordBatchStream, +}; +use datafusion_common::DataFusionError::Execution; +use datafusion_common::stats::Precision; +use datafusion_common::{ColumnStatistics, DataFusionError, Result, Statistics}; +use futures::stream::{self, BoxStream, TryStreamExt}; +use futures::{Stream, StreamExt}; + +use crate::{external_error, inexact_usize_from_u64}; +use hudi_core::file_group::file_slice::FileSlice; +use hudi_core::file_group::reader::FileGroupReader; +use hudi_core::table::ReadOptions; + +/// DataFusion execution plan that reads Hudi file slices through +/// [`FileGroupReader`]. +/// +/// Used for non-Parquet base file formats (Lance) and MOR snapshot +/// queries where base + log file merging is required. Parquet-only +/// COW and MOR read-optimized queries continue to use DataFusion's +/// native `ParquetSource` path for row-group/page-level pruning. +#[derive(Debug)] +pub struct HudiScanExec { + file_slice_partitions: Vec>>, + file_group_reader: Arc, + read_options: ReadOptions, + input_partitions: usize, + file_slice_read_concurrency: usize, + projected_schema: SchemaRef, + projection: Option>, + limit: Option, + properties: PlanProperties, + metrics: ExecutionPlanMetricsSet, +} + +impl HudiScanExec { + #[allow(clippy::too_many_arguments)] + pub fn new( + file_slice_partitions: Vec>, + file_group_reader: Arc, + read_options: ReadOptions, + input_partitions: usize, + file_slice_read_concurrency: usize, + schema: SchemaRef, + projection: Option>, + limit: Option, + ) -> Self { + let projected_schema = if let Some(ref proj) = projection { + let fields: Vec<_> = proj.iter().map(|&i| schema.field(i).clone()).collect(); + Arc::new(arrow_schema::Schema::new(fields)) + } else { + schema.clone() + }; + + // Empty input is normalized to one empty partition so that + // `Partitioning::UnknownPartitioning(0)` doesn't propagate downstream + // (DataFusion's planners typically expect at least one partition; some + // operators panic on zero). `execute(0)` then returns an empty stream + // via the `if file_slices.is_empty()` short-circuit below. + let partitions: Vec>> = if file_slice_partitions.is_empty() { + vec![Arc::new(vec![])] + } else { + file_slice_partitions.into_iter().map(Arc::new).collect() + }; + let n_partitions = partitions.len(); + + let properties = PlanProperties::new( + datafusion::physical_expr::EquivalenceProperties::new(projected_schema.clone()), + Partitioning::UnknownPartitioning(n_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + ); + + Self { + file_slice_partitions: partitions, + file_group_reader, + read_options, + input_partitions, + file_slice_read_concurrency: file_slice_read_concurrency.max(1), + projected_schema, + projection, + limit, + properties, + metrics: ExecutionPlanMetricsSet::new(), + } + } + + #[cfg(test)] + pub(crate) fn read_options(&self) -> &ReadOptions { + &self.read_options + } +} + +struct LimitBatchStream { + inner: BoxStream<'static, Result>, + remaining: usize, +} + +impl LimitBatchStream { + fn new(stream: S, limit: usize) -> Self + where + S: Stream> + Send + 'static, + { + Self { + inner: stream.boxed(), + remaining: limit, + } + } +} + +impl Stream for LimitBatchStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.remaining == 0 { + return Poll::Ready(None); + } + + match this.inner.as_mut().poll_next(cx) { + Poll::Ready(Some(Ok(batch))) => { + let row_count = batch.num_rows(); + if row_count > this.remaining { + let limited = batch.slice(0, this.remaining); + this.remaining = 0; + Poll::Ready(Some(Ok(limited))) + } else { + this.remaining -= row_count; + Poll::Ready(Some(Ok(batch))) + } + } + other => other, + } + } +} + +struct BaselineMetricStream { + inner: BoxStream<'static, Result>, + baseline_metrics: BaselineMetrics, +} + +impl BaselineMetricStream { + fn new(stream: S, baseline_metrics: BaselineMetrics) -> Self + where + S: Stream> + Send + 'static, + { + Self { + inner: stream.boxed(), + baseline_metrics, + } + } +} + +impl Stream for BaselineMetricStream { + type Item = Result; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + let _timer = this.baseline_metrics.elapsed_compute().timer(); + let poll = this.inner.as_mut().poll_next(cx); + this.baseline_metrics.record_poll(poll) + } +} + +impl DisplayAs for HudiScanExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let total_slices: usize = self.file_slice_partitions.iter().map(|p| p.len()).sum(); + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "HudiScanExec: input_partitions={}, partitions={}, file_slices={}, file_slice_read_concurrency={}, projection={:?}, limit={:?}", + self.input_partitions, + self.file_slice_partitions.len(), + total_slices, + self.file_slice_read_concurrency, + self.projection, + self.limit, + ) + } + _ => { + write!(f, "HudiScanExec") + } + } + } +} + +impl ExecutionPlan for HudiScanExec { + fn name(&self) -> &str { + "HudiScanExec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.is_empty() { + Ok(self) + } else { + Err(Execution( + "HudiScanExec is a leaf node and does not accept children".to_string(), + )) + } + } + + fn metrics(&self) -> Option { + Some(self.metrics.clone_inner()) + } + + // Default `cardinality_effect()` (`Unknown`) is intentional. The + // `CardinalityEffect::Equal` variant describes pass-through nodes that + // preserve their input cardinality; for a leaf scan whose row count + // comes from `partition_statistics()`, `Unknown` is the conservative + // and correct value. + + // We deliberately do NOT override `repartitioned`. Partition count is + // fixed at `scan()` time from `state.target_partitions` (or the + // `hoodie.read.input.partitions` config) and we re-chunk file slices + // accordingly. Re-chunking after the fact would require redistributing + // already-bound `FileSlice` lists, which is not currently worth the + // complexity. Default `Ok(None)` declines re-partitioning. + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + let file_slices = self + .file_slice_partitions + .get(partition) + .ok_or_else(|| { + Execution(format!( + "HudiScanExec partition {partition} out of range (have {})", + self.file_slice_partitions.len() + )) + })? + .clone(); + + let projected_schema = self.projected_schema.clone(); + let baseline_metrics = BaselineMetrics::new(&self.metrics, partition); + + if file_slices.is_empty() || self.limit == Some(0) { + let stream = BaselineMetricStream::new(futures::stream::empty(), baseline_metrics); + return Ok(Box::pin(RecordBatchStreamAdapter::new( + projected_schema, + stream, + ))); + } + + let reader = self.file_group_reader.clone(); + let options = self.read_options.clone(); + let concurrency = self + .file_slice_read_concurrency + .min(file_slices.len()) + .max(1); + + let stream = stream::iter(0..file_slices.len()) + .map(move |idx| { + let file_slice = file_slices[idx].clone(); + let reader = reader.clone(); + let options = options.clone(); + async move { + let inner_stream = + reader + .read_file_slice_stream(&file_slice, &options) + .await + .map_err(|e| external_error("Failed to read file slice", e))?; + Ok::<_, DataFusionError>( + inner_stream.map_err(|e| external_error("Failed to read batch", e)), + ) + } + }) + // Scan output is unordered; SQL ordering is provided by explicit + // SortExec nodes. Keep both stages on the same small knob: one cap + // for async stream construction and one cap for active slice + // streams. Raising this can multiply memory pressure across input + // partitions on wide MOR scans. + .buffer_unordered(concurrency) + .try_flatten_unordered(concurrency) + .boxed(); + let stream = if let Some(limit) = self.limit { + LimitBatchStream::new(stream, limit).boxed() + } else { + stream + }; + let stream = BaselineMetricStream::new(stream, baseline_metrics); + + Ok(Box::pin(RecordBatchStreamAdapter::new( + projected_schema, + stream, + ))) + } + + fn statistics(&self) -> Result { + Ok(self.aggregate_file_slice_statistics()) + } + + fn partition_statistics(&self, partition: Option) -> Result { + let column_statistics = + vec![ColumnStatistics::new_unknown(); self.projected_schema.fields().len()]; + + let partitions: Box + '_> = match partition { + None => Box::new( + self.file_slice_partitions + .iter() + .map(|slices| slices.as_slice()), + ), + Some(idx) => match self.file_slice_partitions.get(idx) { + Some(slices) => Box::new(std::iter::once(slices.as_slice())), + None => return Ok(Statistics::new_unknown(&self.projected_schema)), + }, + }; + + Ok(Self::aggregate_partitions(partitions, column_statistics)) + } + + fn with_fetch(&self, limit: Option) -> Option> { + Some(Arc::new(Self { + file_slice_partitions: self.file_slice_partitions.clone(), + file_group_reader: self.file_group_reader.clone(), + read_options: self.read_options.clone(), + input_partitions: self.input_partitions, + file_slice_read_concurrency: self.file_slice_read_concurrency, + projected_schema: self.projected_schema.clone(), + projection: self.projection.clone(), + limit, + properties: self.properties.clone(), + // `with_fetch` is a planner-time clone used before execution. Metrics + // belong to the cloned plan instance and must start empty. + metrics: ExecutionPlanMetricsSet::new(), + })) + } + + fn fetch(&self) -> Option { + self.limit + } +} + +impl HudiScanExec { + fn aggregate_file_slice_statistics(&self) -> Statistics { + let column_statistics = + vec![ColumnStatistics::new_unknown(); self.projected_schema.fields().len()]; + Self::aggregate_partitions( + self.file_slice_partitions + .iter() + .map(|slices| slices.as_slice()), + column_statistics, + ) + } + + fn aggregate_partitions<'a, I>( + partitions: I, + column_statistics: Vec, + ) -> Statistics + where + I: IntoIterator, + { + let mut total_rows: u64 = 0; + let mut total_byte_size: u64 = 0; + let mut have_row_estimate = false; + let mut have_byte_estimate = false; + + for slices in partitions { + for file_slice in slices { + if let Some(meta) = &file_slice.base_file.file_metadata { + if meta.num_records > 0 { + total_rows = total_rows.saturating_add(meta.num_records as u64); + have_row_estimate = true; + } + if meta.size > 0 { + total_byte_size = total_byte_size.saturating_add(meta.size); + have_byte_estimate = true; + } + } + for log_file in &file_slice.log_files { + if let Some(meta) = &log_file.file_metadata + && meta.size > 0 + { + total_byte_size = total_byte_size.saturating_add(meta.size); + have_byte_estimate = true; + } + } + } + } + + let num_rows = if have_row_estimate { + inexact_usize_from_u64(total_rows) + } else { + Precision::Absent + }; + let total_byte_size = if have_byte_estimate { + inexact_usize_from_u64(total_byte_size) + } else { + Precision::Absent + }; + + Statistics { + num_rows, + total_byte_size, + column_statistics, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_array::Int32Array; + use arrow_schema::Schema; + use arrow_schema::{DataType, Field}; + use hudi_core::config::util::empty_options; + use hudi_core::file_group::base_file::BaseFile; + use hudi_core::storage::file_metadata::FileMetadata; + use std::collections::BTreeSet; + use std::fs::canonicalize; + use std::path::Path; + use std::str::FromStr; + use url::Url; + + fn int_batch(values: Vec) -> RecordBatch { + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])), + vec![Arc::new(Int32Array::from(values))], + ) + .unwrap() + } + + fn int_values(batch: &RecordBatch) -> Vec { + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + (0..values.len()).map(|idx| values.value(idx)).collect() + } + + fn file_slice_with_meta( + file_name: &str, + size: u64, + num_records: i64, + byte_size: i64, + ) -> FileSlice { + let mut bf = BaseFile::from_str(file_name).unwrap(); + bf.file_metadata = Some(FileMetadata { + name: file_name.to_string(), + size, + byte_size, + num_records, + }); + FileSlice { + base_file: bf, + log_files: BTreeSet::new(), + partition_path: String::new(), + base_file_column_stats: None, + } + } + + #[test] + fn test_aggregate_partitions_sums_rows_and_bytes() { + let partitions = [ + vec![ + file_slice_with_meta("fileA-0_0-1-1_20250101000000000.parquet", 100, 10, 200), + file_slice_with_meta("fileB-0_0-1-1_20250101000000000.parquet", 300, 20, 600), + ], + vec![file_slice_with_meta( + "fileC-0_0-1-1_20250101000000000.parquet", + 500, + 30, + 1000, + )], + ]; + + let stats = HudiScanExec::aggregate_partitions( + partitions.iter().map(Vec::as_slice), + vec![ColumnStatistics::new_unknown(); 2], + ); + + assert_eq!(stats.num_rows, Precision::Inexact(60)); + assert_eq!(stats.total_byte_size, Precision::Inexact(900)); + assert_eq!(stats.column_statistics.len(), 2); + } + + #[test] + fn test_aggregate_partitions_returns_absent_when_metadata_missing() { + let mut bf = BaseFile::from_str("fileA-0_0-1-1_20250101000000000.parquet").unwrap(); + bf.file_metadata = None; + let slices = [vec![FileSlice { + base_file: bf, + log_files: BTreeSet::new(), + partition_path: String::new(), + base_file_column_stats: None, + }]]; + + let stats = HudiScanExec::aggregate_partitions( + slices.iter().map(Vec::as_slice), + vec![ColumnStatistics::new_unknown()], + ); + + assert!(matches!(stats.num_rows, Precision::Absent)); + assert!(matches!(stats.total_byte_size, Precision::Absent)); + } + + #[test] + fn test_aggregate_partitions_byte_size_only_when_records_unknown() { + // Mirrors the Lance case: file listing reports `size` but + // FileStatsEstimator has not populated `num_records`. + let slices = [vec![file_slice_with_meta( + "fileA-0_0-1-1_20250101000000000.lance", + 500, + 0, + 0, + )]]; + + let stats = HudiScanExec::aggregate_partitions( + slices.iter().map(Vec::as_slice), + vec![ColumnStatistics::new_unknown()], + ); + + assert!(matches!(stats.num_rows, Precision::Absent)); + assert_eq!(stats.total_byte_size, Precision::Inexact(500)); + } + + #[tokio::test] + async fn test_limit_batch_stream_truncates_and_stops() { + let batches = vec![ + Ok(int_batch(vec![1, 2, 3])), + Ok(int_batch(vec![4, 5, 6])), + Ok(int_batch(vec![7, 8, 9])), + ]; + + let limited = LimitBatchStream::new(stream::iter(batches), 4) + .try_collect::>() + .await + .unwrap(); + + assert_eq!(limited.len(), 2); + assert_eq!(int_values(&limited[0]), vec![1, 2, 3]); + assert_eq!(int_values(&limited[1]), vec![4]); + } + + #[tokio::test] + async fn test_metrics_override_is_wired() { + let base_url = + Url::from_file_path(canonicalize(Path::new("tests/data/table_props_valid")).unwrap()) + .unwrap(); + let reader = Arc::new( + FileGroupReader::new_with_options(base_url.as_str(), empty_options()) + .await + .unwrap(), + ); + let exec = HudiScanExec::new( + vec![], + reader, + ReadOptions::new(), + 1, + 1, + Arc::new(Schema::empty()), + None, + None, + ); + + assert!(exec.metrics().is_some()); + } +} diff --git a/crates/datafusion/src/lib.rs b/crates/datafusion/src/lib.rs index 600980e1..7f0744a7 100644 --- a/crates/datafusion/src/lib.rs +++ b/crates/datafusion/src/lib.rs @@ -17,10 +17,12 @@ * under the License. */ +pub(crate) mod hudi_exec; pub(crate) mod util; use std::any::Any; use std::collections::HashMap; +use std::error::Error; use std::fmt::Debug; use std::sync::Arc; @@ -39,20 +41,54 @@ use datafusion::logical_expr::Operator; use datafusion::physical_plan::ExecutionPlan; use datafusion_common::DFSchema; use datafusion_common::DataFusionError::Execution; -use datafusion_common::Statistics; use datafusion_common::config::TableParquetOptions; use datafusion_common::stats::Precision; +use datafusion_common::{DataFusionError, Statistics}; +use datafusion_expr::utils::split_conjunction; use datafusion_expr::{CreateExternalTable, Expr, TableProviderFilterPushDown, TableType}; use datafusion_physical_expr::create_physical_expr; use log::warn; +use crate::hudi_exec::HudiScanExec; use crate::util::expr::exprs_to_filters; -use hudi_core::config::read::HudiReadConfig::{InputPartitions, UseReadOptimizedMode}; +use hudi_core::config::read::HudiReadConfig::{ + FileSliceReadConcurrency, InputPartitions, UseReadOptimizedMode, +}; use hudi_core::config::table::{BaseFileFormatValue, HudiTableConfig}; use hudi_core::config::util::empty_options; +use hudi_core::config::{ConfigParser, HudiConfigs}; +use hudi_core::file_group::file_slice::FileSlice; use hudi_core::storage::util::{get_scheme_authority, join_url_segments}; use hudi_core::table::{ReadOptions, Table as HudiTable}; +fn default_file_slice_read_concurrency() -> usize { + match FileSliceReadConcurrency.default_value() { + Some(value) => value.into(), + None => unreachable!("FileSliceReadConcurrency has a default value defined in hudi-core"), + } +} + +pub(crate) fn inexact_usize_from_u64(value: u64) -> Precision { + match usize::try_from(value) { + Ok(value) => Precision::Inexact(value), + Err(_) => Precision::Absent, + } +} + +pub(crate) fn external_error(context: impl Into, error: E) -> DataFusionError +where + E: Error + Send + Sync + 'static, +{ + DataFusionError::External(Box::new(error)).context(context) +} + +fn filter_field_matches_partition_column(filter_field: &str, partition_column: &str) -> bool { + filter_field == partition_column + || filter_field + .rsplit_once('.') + .is_some_and(|(_, name)| name == partition_column) +} + /// Create a `HudiDataSource`. /// Used for Datafusion to query Hudi tables /// @@ -87,15 +123,36 @@ use hudi_core::table::{ReadOptions, Table as HudiTable}; pub struct HudiDataSource { table: Arc, /// Cached table schema (with meta fields) for synchronous access in `TableProvider::schema()`. + /// This provider is a construction-time metadata snapshot; create a new + /// provider to observe schema/stat changes from later commits. schema: SchemaRef, /// Cached partition schema for determining partition columns. /// This is cached at construction since partition schema rarely changes /// and is needed synchronously in `supports_filters_pushdown`. partition_schema: Schema, /// Cached table-level statistics for join ordering and broadcast decisions. + /// + /// Consumed in two places: + /// 1. `TableProvider::statistics()` - returned to the optimizer (note: + /// DataFusion's main planner does not currently consult this method; + /// see the trait docs). + /// 2. `scan_parquet` - passed to `FileScanConfigBuilder::with_statistics(...)`, + /// which IS consumed by the optimizer for join planning on the + /// Parquet/COW fast path. + /// + /// The `HudiScanExec` path (Lance and MOR snapshot) derives statistics + /// per-execution from `FileSlice` metadata via `aggregate_partitions` + /// rather than reusing this table-level cache, since per-partition stats + /// are more useful for that path. cached_stats: Option, /// Number of input partitions for scan planning, extracted from read options. input_partitions: usize, + /// Read-optimized mode requested when constructing the provider. + read_optimized_mode: bool, + /// Maximum number of file-slice streams polled concurrently within a scan partition. + file_slice_read_concurrency: usize, + /// Explicit base file format from table config, if present. + base_file_format: Option, } impl std::fmt::Debug for HudiDataSource { @@ -142,18 +199,58 @@ impl HudiDataSource { })?, None => 0, }; + let read_optimized_mode: bool = match all_options + .iter() + .find(|(k, _)| k == UseReadOptimizedMode.as_ref()) + { + Some((_, v)) => v.parse().map_err(|e| { + Execution(format!( + "Invalid value '{v}' for {}: {e}", + UseReadOptimizedMode.as_ref() + )) + })?, + None => false, + }; + let file_slice_read_concurrency: usize = match all_options + .iter() + .find(|(k, _)| k == FileSliceReadConcurrency.as_ref()) + { + Some((_, v)) => { + let parsed = v.parse().map_err(|_| { + Execution(format!( + "Invalid value '{v}' for {}: expected a positive integer", + FileSliceReadConcurrency.as_ref() + )) + })?; + if parsed == 0 { + return Err(Execution(format!( + "Invalid value '0' for {}: expected a positive integer", + FileSliceReadConcurrency.as_ref() + ))); + } + parsed + } + None => default_file_slice_read_concurrency(), + }; let table = HudiTable::new_with_options(base_uri, all_options) .await - .map_err(|e| Execution(format!("Failed to create Hudi table: {e}")))?; - - let base_file_format: String = table - .hudi_configs - .get_or_default(HudiTableConfig::BaseFileFormat) - .into(); - if !base_file_format.eq_ignore_ascii_case(BaseFileFormatValue::Parquet.as_ref()) { - return Err(Execution(format!( - "Unsupported base file format '{base_file_format}' for HudiDataSource; only parquet is supported" - ))); + .map_err(|e| external_error("Failed to create Hudi table", e))?; + + let base_file_format = + BaseFileFormatValue::from_configs(&table.hudi_configs).map_err(|e| { + external_error( + format!( + "Invalid {} config", + HudiTableConfig::BaseFileFormat.as_ref() + ), + e, + ) + })?; + if matches!(base_file_format, Some(BaseFileFormatValue::HFile)) { + return Err(Execution( + "HFile is only supported for Hudi metadata tables, not regular DataFusion scans" + .to_string(), + )); } // Cache schema with meta fields at construction for synchronous access @@ -182,13 +279,9 @@ impl HudiDataSource { let cached_stats = match table.compute_table_stats(None).await { Some((num_rows, total_byte_size)) => { let num_fields = schema.fields().len(); - // Saturate on 32-bit targets where `u64` cannot fit into `usize`; - // on 64-bit this is a lossless conversion. - let num_rows = usize::try_from(num_rows).unwrap_or(usize::MAX); - let total_byte_size = usize::try_from(total_byte_size).unwrap_or(usize::MAX); Some(Statistics { - num_rows: Precision::Inexact(num_rows), - total_byte_size: Precision::Inexact(total_byte_size), + num_rows: inexact_usize_from_u64(num_rows), + total_byte_size: inexact_usize_from_u64(total_byte_size), column_statistics: vec![ datafusion_common::ColumnStatistics::new_unknown(); num_fields @@ -204,6 +297,9 @@ impl HudiDataSource { partition_schema, cached_stats, input_partitions, + read_optimized_mode, + file_slice_read_concurrency, + base_file_format, }) } @@ -211,6 +307,115 @@ impl HudiDataSource { self.input_partitions } + #[cfg(test)] + fn get_file_slice_read_concurrency(&self) -> usize { + self.file_slice_read_concurrency + } + + /// Returns the effective `UseReadOptimizedMode` value cached from the + /// table options at construction time. Used by [`Self::use_parquet_source`] + /// to decide whether MOR snapshot semantics are required. + fn effective_read_optimized(&self) -> bool { + self.read_optimized_mode + } + + fn scan_read_options( + &self, + pushdown_filters: Vec<(String, String, String)>, + read_optimized: bool, + ) -> Result { + let mut read_options = ReadOptions::new() + .with_filters(pushdown_filters) + .map_err(|e| external_error("Invalid pushdown filter", e))?; + if read_optimized { + read_options = read_options.with_hudi_option(UseReadOptimizedMode.as_ref(), "true"); + } + Ok(read_options) + } + + /// Build the [`ReadOptions`] passed to `HudiScanExec` for per-slice reads. + /// + /// Hudi table-level planning has already applied partition filters. When + /// partition fields are dropped from data files, passing those filters to + /// `FileGroupReader` would fail its strict batch-schema validation. + fn read_options_for_hudi_exec( + hudi_configs: &HudiConfigs, + options: &ReadOptions, + ) -> ReadOptions { + let drops_partition_columns: bool = hudi_configs + .get_or_default(HudiTableConfig::DropsPartitionFields) + .into(); + if !drops_partition_columns || options.filters.is_empty() { + return options.clone(); + } + + let partition_columns: Vec = hudi_configs + .get_or_default(HudiTableConfig::PartitionFields) + .into(); + let mut applicable = options.clone(); + applicable.filters = options + .filters + .iter() + .filter(|filter| { + !partition_columns + .iter() + .any(|p| filter_field_matches_partition_column(&filter.field, p)) + }) + .cloned() + .collect(); + applicable + } + + /// Returns true iff every file slice has a `.parquet` base file. + /// Empty input returns `false` so the scan is routed to `HudiScanExec`, + /// which handles the empty case via `RecordBatchStreamAdapter::new(.., empty())`. + fn file_slices_are_parquet(file_slices: &[FileSlice]) -> Result { + if file_slices.is_empty() { + return Ok(false); + } + for file_slice in file_slices { + let relative_path = file_slice.base_file_relative_path().map_err(|e| { + external_error( + format!("Failed to get base file relative path for {file_slice:?}"), + e, + ) + })?; + if !BaseFileFormatValue::Parquet.matches_extension(&relative_path) { + return Ok(false); + } + } + Ok(true) + } + + /// Decides whether a scan can be served by DataFusion's native + /// `ParquetSource` (cheap row-group / page pruning) or must go through + /// [`HudiScanExec`]. + /// + /// Routing matrix: + /// - Parquet COW → `ParquetSource` + /// - Parquet MOR read-optimized → `ParquetSource` + /// - Parquet MOR snapshot → `HudiScanExec` (base + log merging) + /// - Lance and any other base format → `HudiScanExec` + fn use_parquet_source( + &self, + read_options: &ReadOptions, + file_slices: &[FileSlice], + ) -> Result { + let parquet_base_files = match &self.base_file_format { + Some(format) => matches!(format, BaseFileFormatValue::Parquet), + None => Self::file_slices_are_parquet(file_slices)?, + }; + if !parquet_base_files { + return Ok(false); + } + if !self.table.is_mor() { + return Ok(true); + } + read_options + .is_read_optimized() + .map_err(|e| external_error("Invalid read-optimized option", e)) + } + /// Check if the given expression can be pushed down to the Hudi table. /// /// The expression can be pushed down if it is: @@ -218,7 +423,7 @@ impl HudiDataSource { /// - A NOT expression wrapping a pushable expression /// - An AND compound expression where at least one side can be pushed down /// - A BETWEEN expression with column and literals - fn can_push_down(&self, expr: &Expr) -> bool { + fn can_push_down_expr(schema: &Schema, expr: &Expr) -> bool { match expr { Expr::BinaryExpr(binary_expr) => { let left = &binary_expr.left; @@ -228,35 +433,44 @@ impl HudiDataSource { match op { Operator::And => { // AND is pushable if at least one side is pushable - self.can_push_down(left) || self.can_push_down(right) + Self::can_push_down_expr(schema, left) + || Self::can_push_down_expr(schema, right) } Operator::Or => { // OR cannot be pushed down with current filter model false } _ => { - self.is_supported_operator(op) - && self.is_supported_operand(left) - && self.is_supported_operand(right) + Self::is_supported_operator(op) + && Self::is_supported_operand(schema, left) + && Self::is_supported_operand(schema, right) } } } Expr::Not(inner_expr) => { // Recursively check if the inner expression can be pushed down - self.can_push_down(inner_expr) + Self::can_push_down_expr(schema, inner_expr) } Expr::Between(between) => { // BETWEEN can be pushed if expr is a column and bounds are literals !between.negated - && matches!(&*between.expr, Expr::Column(_)) + && matches!(&*between.expr, Expr::Column(col) if schema.column_with_name(&col.name).is_some()) && matches!(&*between.low, Expr::Literal(..)) && matches!(&*between.high, Expr::Literal(..)) } + Expr::InList(in_list) => { + !in_list.list.is_empty() + && matches!(in_list.expr.as_ref(), Expr::Column(col) if schema.column_with_name(&col.name).is_some()) + && in_list + .list + .iter() + .all(|expr| matches!(expr, Expr::Literal(..))) + } _ => false, } } - fn is_supported_operator(&self, op: &Operator) -> bool { + fn is_supported_operator(op: &Operator) -> bool { matches!( op, Operator::Eq @@ -268,9 +482,9 @@ impl HudiDataSource { ) } - fn is_supported_operand(&self, expr: &Expr) -> bool { + fn is_supported_operand(schema: &Schema, expr: &Expr) -> bool { match expr { - Expr::Column(col) => self.schema().column_with_name(&col.name).is_some(), + Expr::Column(col) => schema.column_with_name(&col.name).is_some(), Expr::Literal(..) => true, _ => false, } @@ -285,101 +499,139 @@ impl HudiDataSource { .collect() } - /// Checks if expression filters on a partition column. + /// Checks if expression filters only on a partition column. /// - /// Partition column filters can be marked as `Exact` because they are - /// fully handled by partition pruning and don't need post-filtering. + /// Partition filters are safe to push into Hudi file listing for every + /// scan path. For Parquet scans, non-partition predicates are left to + /// DataFusion's `ParquetSource` so Hudi doesn't read Parquet footers for + /// stats pruning before DataFusion reads the same files. fn is_partition_column_filter(expr: &Expr, partition_cols: &[String]) -> bool { match expr { - Expr::BinaryExpr(binary_expr) => { - match binary_expr.op { - Operator::And => { - // For AND, check if both sides are partition filters - Self::is_partition_column_filter(&binary_expr.left, partition_cols) - && Self::is_partition_column_filter(&binary_expr.right, partition_cols) - } - _ => { - // For partition filters, one side must be a partition column - // and the other side must be a literal value - match (&*binary_expr.left, &*binary_expr.right) { - (Expr::Column(col), Expr::Literal(..)) - if partition_cols.contains(&col.name) => - { - true - } - (Expr::Literal(..), Expr::Column(col)) - if partition_cols.contains(&col.name) => - { - true - } - _ => false, - } - } + Expr::BinaryExpr(binary_expr) => match binary_expr.op { + Operator::And => { + Self::is_partition_column_filter(&binary_expr.left, partition_cols) + && Self::is_partition_column_filter(&binary_expr.right, partition_cols) } - } + Operator::Or => false, + _ => match (&*binary_expr.left, &*binary_expr.right) { + (Expr::Column(col), Expr::Literal(..)) + | (Expr::Literal(..), Expr::Column(col)) => partition_cols.contains(&col.name), + _ => false, + }, + }, Expr::Not(inner) => Self::is_partition_column_filter(inner, partition_cols), Expr::Between(between) => { - matches!(&*between.expr, Expr::Column(col) if partition_cols.contains(&col.name)) + !between.negated + && matches!(&*between.expr, Expr::Column(col) if partition_cols.contains(&col.name)) + && matches!(&*between.low, Expr::Literal(..)) + && matches!(&*between.high, Expr::Literal(..)) + } + Expr::InList(in_list) => { + !in_list.list.is_empty() + && matches!(in_list.expr.as_ref(), Expr::Column(col) if partition_cols.contains(&col.name)) + && in_list + .list + .iter() + .all(|expr| matches!(expr, Expr::Literal(..))) } _ => false, } } -} -#[async_trait] -impl TableProvider for HudiDataSource { - fn as_any(&self) -> &dyn Any { - self + fn is_exact_partition_equality_filter(expr: &Expr, partition_cols: &[String]) -> bool { + match expr { + Expr::BinaryExpr(binary_expr) if binary_expr.op == Operator::Eq => { + match (&*binary_expr.left, &*binary_expr.right) { + (Expr::Column(col), Expr::Literal(..)) + | (Expr::Literal(..), Expr::Column(col)) => partition_cols.contains(&col.name), + _ => false, + } + } + _ => false, + } } - fn schema(&self) -> SchemaRef { - self.schema.clone() + fn filter_pushdown_support( + table_schema: &Schema, + partition_cols: &[String], + expr: &Expr, + ) -> TableProviderFilterPushDown { + let conjuncts = split_conjunction(expr); + let has_pushable_conjunct = conjuncts + .iter() + .any(|conjunct| Self::can_push_down_expr(table_schema, conjunct)); + + if !has_pushable_conjunct { + return TableProviderFilterPushDown::Unsupported; + } + + let all_conjuncts_are_exact_partition_eq = conjuncts.iter().all(|conjunct| { + Self::can_push_down_expr(table_schema, conjunct) + && Self::is_exact_partition_equality_filter(conjunct, partition_cols) + }); + + if all_conjuncts_are_exact_partition_eq { + TableProviderFilterPushDown::Exact + } else { + TableProviderFilterPushDown::Inexact + } } - fn table_type(&self) -> TableType { - TableType::Base + /// Returns `(partition_pushdown_exprs, all_pushdown_exprs)`. + fn split_scan_pushdown_exprs(&self, filters: &[Expr]) -> (Vec, Vec) { + let partition_cols = self.get_partition_columns(); + Self::split_scan_pushdown_exprs_for_schema(self.schema.as_ref(), &partition_cols, filters) } - fn statistics(&self) -> Option { - self.cached_stats.clone() + fn split_scan_pushdown_exprs_for_schema( + table_schema: &Schema, + partition_cols: &[String], + filters: &[Expr], + ) -> (Vec, Vec) { + let all_pushdown_exprs: Vec = filters + .iter() + .flat_map(|expr| split_conjunction(expr).into_iter()) + .filter(|expr| Self::can_push_down_expr(table_schema, expr)) + .cloned() + .collect(); + let partition_pushdown_exprs = all_pushdown_exprs + .iter() + .filter(|expr| Self::is_partition_column_filter(expr, partition_cols)) + .cloned() + .collect(); + + (partition_pushdown_exprs, all_pushdown_exprs) } - async fn scan( + fn use_parquet_source_without_file_slices( + &self, + read_options: &ReadOptions, + ) -> Result> { + if self.table.is_mor() + && !read_options + .is_read_optimized() + .map_err(|e| external_error("Invalid read-optimized option", e))? + { + return Ok(Some(false)); + } + + match &self.base_file_format { + Some(BaseFileFormatValue::Parquet) => Ok(Some(true)), + Some(_) => Ok(Some(false)), + None => Ok(None), + } + } + + #[allow(clippy::too_many_arguments)] + async fn scan_parquet( &self, state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, + flat_slices: Vec, ) -> Result> { - self.table.register_storage(state.runtime_env().clone()); - - // Resolve input partitions: use Hudi config if set, otherwise fall back - // to DataFusion's target_partitions (defaults to number of CPU cores). - let input_partitions = match self.get_input_partitions() { - 0 => state.config_options().execution.target_partitions, - n => n, - }; - - // Only push down partition column filters to Hudi's file listing layer. - // Non-partition filters would trigger Parquet footer reads for file-level - // stats pruning, which is redundant since DataFusion's ParquetSource already - // handles row-group/page-level pruning during the scan. - let partition_cols = self.get_partition_columns(); - let partition_filters: Vec = filters - .iter() - .filter(|expr| Self::is_partition_column_filter(expr, &partition_cols)) - .cloned() - .collect(); - let pushdown_filters = exprs_to_filters(&partition_filters); - let read_options = ReadOptions::new() - .with_filters(pushdown_filters) - .map_err(|e| Execution(format!("Invalid pushdown filter: {e}")))? - .with_hudi_option(UseReadOptimizedMode.as_ref(), "true"); - let flat_slices = self - .table - .get_file_slices(&read_options) - .await - .map_err(|e| Execution(format!("Failed to get file slices from Hudi table: {e}")))?; + let input_partitions = self.get_input_partitions_for_scan(state); let file_slices = hudi_core::util::collection::split_into_chunks(flat_slices, input_partitions); let base_url = self.table.base_url(); @@ -388,12 +640,13 @@ impl TableProvider for HudiDataSource { let mut parquet_file_group_vec = Vec::new(); for f in file_slice_vec { let relative_path = f.base_file_relative_path().map_err(|e| { - Execution(format!( - "Failed to get base file relative path for {f:?} due to {e:?}" - )) + external_error( + format!("Failed to get base file relative path for {f:?}"), + e, + ) })?; let url = join_url_segments(&base_url, &[relative_path.as_str()]) - .map_err(|e| Execution(format!("Failed to join URL segments: {e:?}")))?; + .map_err(|e| external_error("Failed to join URL segments", e))?; let size = f.base_file.file_metadata.as_ref().map_or(0, |m| m.size); let partitioned_file = PartitionedFile::new(url.path(), size); parquet_file_group_vec.push(partitioned_file); @@ -401,9 +654,7 @@ impl TableProvider for HudiDataSource { parquet_file_groups.push(parquet_file_group_vec) } - let base_url = self.table.base_url(); let url = ObjectStoreUrl::parse(get_scheme_authority(&base_url))?; - let parquet_opts = TableParquetOptions { global: state.config_options().execution.parquet.clone(), column_specific_options: Default::default(), @@ -433,17 +684,204 @@ impl TableProvider for HudiDataSource { .with_projection_indices(projection.cloned())? .with_limit(limit); - // Pass table statistics to the physical plan so DataFusion's join_selection - // optimizer can swap build/probe sides and choose broadcast joins. if let Some(stats) = &self.cached_stats { + // DataFusion's FileScanConfig stores unprojected table statistics + // and applies the source projection inside partition_statistics(). fsc_builder = fsc_builder.with_statistics(stats.clone()); } let fsc = fsc_builder.build(); - Ok(Arc::new(DataSourceExec::new(Arc::new(fsc)))) } + async fn scan_hudi( + &self, + projection: Option<&Vec>, + limit: Option, + input_partitions: usize, + flat_slices: Vec, + read_options: ReadOptions, + ) -> Result> { + let file_slices = + hudi_core::util::collection::split_into_chunks(flat_slices, input_partitions); + + // The reader is built with the caller's full options so table-level + // setup sees the same query context. Dropped partition filters are only + // stripped from the per-slice options passed into HudiScanExec below, + // before FileGroupReader validates base-file batch schemas. + let file_group_reader = Arc::new( + self.table + .create_file_group_reader_with_options(Some(&read_options), empty_options()) + .map_err(|e| external_error("Failed to create FileGroupReader", e))?, + ); + + let mut hudi_read_options = + Self::read_options_for_hudi_exec(&self.table.hudi_configs, &read_options); + if let Some(proj) = projection { + let col_names: Vec = proj + .iter() + .map(|&i| self.schema.field(i).name().clone()) + .collect(); + hudi_read_options = hudi_read_options.with_projection(col_names); + } + + Ok(Arc::new(HudiScanExec::new( + file_slices, + file_group_reader, + hudi_read_options, + input_partitions, + self.file_slice_read_concurrency, + self.schema.clone(), + projection.cloned(), + limit, + ))) + } + + fn get_input_partitions_for_scan(&self, state: &dyn Session) -> usize { + match self.get_input_partitions() { + 0 => state.config_options().execution.target_partitions, + n => n, + } + } +} + +#[async_trait] +impl TableProvider for HudiDataSource { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + fn statistics(&self) -> Option { + self.cached_stats.clone() + } + + /// Builds the scan `ExecutionPlan` for this Hudi table. + /// + /// Per DataFusion's [custom table provider guide], `scan()` should avoid + /// data-reading I/O (work proportional to row count). Catalog I/O — the + /// listing needed to decide what files exist and which ones the query + /// touches — is expected; it is what DataFusion's own `ListingTable` does. + /// + /// `Table::get_file_slices(&read_options)` is catalog I/O: it loads the + /// timeline, resolves file groups, and applies partition pruning. Its cost + /// is proportional to partition count, not row count. When the metadata + /// table is enabled, partition discovery reads a few HFile blocks instead + /// of doing recursive object-store listing, so MDT is strictly cheaper + /// than the generic `ListingTable` path. + /// + /// Callers that issue many `scan()` calls per session against the same + /// table (e.g. ad-hoc SQL across the same dataset) may still benefit from + /// caching the resolved file-slice set keyed by + /// `(filters, as_of_timestamp)`. Tracked as future work. + /// + /// We implement the legacy `scan()` rather than `scan_with_args()`. The + /// default `scan_with_args()` impl delegates to `scan()`, so today's + /// behavior is identical. Revisit if a future DataFusion version adds + /// optimization hints to `scan_with_args()` that we'd want to consume. + /// + /// [custom table provider guide]: https://datafusion.apache.org/library-user-guide/custom-table-providers.html + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result> { + // Idempotent: registers our object store with the session's runtime + // so the Parquet path (DataSourceExec) can resolve `s3://`/`gs://`/`az://` + // URIs. Per-scan invocation is intentional: `TableProvider` has no + // registration hook called by the SessionContext, and direct callers + // (`HudiDataSource::new`) wouldn't have a chance to call it otherwise. + self.table.register_storage(state.runtime_env().clone()); + + let input_partitions = self.get_input_partitions_for_scan(state); + + let (partition_pushdown_exprs, all_pushdown_exprs) = + self.split_scan_pushdown_exprs(filters); + let partition_pushdown_filters = exprs_to_filters(&partition_pushdown_exprs); + let all_pushdown_filters = exprs_to_filters(&all_pushdown_exprs); + let all_filters_are_partition_filters = all_pushdown_filters == partition_pushdown_filters; + + let read_optimized = self.effective_read_optimized(); + let partition_read_options = + self.scan_read_options(partition_pushdown_filters.clone(), read_optimized)?; + let all_read_options = if all_filters_are_partition_filters { + partition_read_options.clone() + } else { + self.scan_read_options(all_pushdown_filters, read_optimized)? + }; + + match self.use_parquet_source_without_file_slices(&partition_read_options)? { + Some(true) => { + let flat_slices = self + .table + .get_file_slices(&partition_read_options) + .await + .map_err(|e| external_error("Failed to get file slices from Hudi table", e))?; + self.scan_parquet(state, projection, filters, limit, flat_slices) + .await + } + Some(false) => { + let flat_slices = self + .table + .get_file_slices(&all_read_options) + .await + .map_err(|e| external_error("Failed to get file slices from Hudi table", e))?; + self.scan_hudi( + projection, + limit, + input_partitions, + flat_slices, + all_read_options, + ) + .await + } + None => { + let partition_flat_slices = self + .table + .get_file_slices(&partition_read_options) + .await + .map_err(|e| external_error("Failed to get file slices from Hudi table", e))?; + + if self.use_parquet_source(&partition_read_options, &partition_flat_slices)? { + self.scan_parquet(state, projection, filters, limit, partition_flat_slices) + .await + } else { + let flat_slices = if all_filters_are_partition_filters { + partition_flat_slices + } else { + self.table + .get_file_slices(&all_read_options) + .await + .map_err(|e| { + external_error("Failed to get file slices from Hudi table", e) + })? + }; + self.scan_hudi( + projection, + limit, + input_partitions, + flat_slices, + all_read_options, + ) + .await + } + } + } + } + + /// Reports partition equality predicates as `Exact`; all other pushed + /// filters are `Inexact` so DataFusion retains a residual `FilterExec`. + /// `scan()` splits conjunctions before converting them to Hudi filters, so + /// pushable atoms inside mixed `AND` predicates still help pruning. fn supports_filters_pushdown( &self, filters: &[&Expr], @@ -453,16 +891,11 @@ impl TableProvider for HudiDataSource { filters .iter() .map(|expr| { - if !self.can_push_down(expr) { - return Ok(TableProviderFilterPushDown::Unsupported); - } - - // Partition column filters are fully handled by partition pruning - if Self::is_partition_column_filter(expr, &partition_cols) { - Ok(TableProviderFilterPushDown::Exact) - } else { - Ok(TableProviderFilterPushDown::Inexact) - } + Ok(Self::filter_pushdown_support( + self.schema.as_ref(), + &partition_cols, + expr, + )) }) .collect() } @@ -551,6 +984,7 @@ impl TableProviderFactory for HudiTableFactory { #[cfg(test)] mod tests { use super::*; + use arrow_schema::{DataType, Field}; use datafusion_common::{Column, ScalarValue}; use hudi_core::config::internal::HudiInternalConfig; use hudi_core::config::table::{BaseFileFormatValue, HudiTableConfig}; @@ -559,7 +993,8 @@ mod tests { use url::Url; use datafusion::logical_expr::BinaryExpr; - use hudi_test::SampleTable::{V6Nonpartitioned, V6SimplekeygenNonhivestyle}; + use datafusion::prelude::SessionContext; + use hudi_test::SampleTable::{V6Nonpartitioned, V6SimplekeygenNonhivestyle, V9LanceTxnsSimple}; use crate::HudiDataSource; @@ -570,12 +1005,49 @@ mod tests { .unwrap(); let hudi = HudiDataSource::new(base_url.as_str()).await.unwrap(); assert_eq!(hudi.get_input_partitions(), 0); + assert_eq!( + hudi.get_file_slice_read_concurrency(), + default_file_slice_read_concurrency() + ); assert_eq!(hudi.table_type(), TableType::Base); assert_eq!(hudi.statistics(), None); } #[tokio::test] - async fn test_new_with_options_fails_fast_on_non_parquet_base_file_format() { + async fn test_new_with_options_sets_file_slice_read_concurrency() { + let hudi = HudiDataSource::new_with_options( + V6Nonpartitioned.path_to_cow().as_str(), + [(FileSliceReadConcurrency.as_ref(), "2")], + ) + .await + .unwrap(); + + assert_eq!(hudi.get_file_slice_read_concurrency(), 2); + } + + #[tokio::test] + async fn test_new_with_options_rejects_invalid_file_slice_read_concurrency() { + for invalid in ["0", "abc"] { + let result = HudiDataSource::new_with_options( + V6Nonpartitioned.path_to_cow().as_str(), + [(FileSliceReadConcurrency.as_ref(), invalid)], + ) + .await; + + assert!(result.is_err()); + let error = result.unwrap_err().to_string(); + assert!(error.contains(FileSliceReadConcurrency.as_ref())); + assert!(error.contains(invalid)); + } + } + + #[test] + fn test_file_slices_are_parquet_empty_is_false() { + assert!(!HudiDataSource::file_slices_are_parquet(&[]).unwrap()); + } + + #[tokio::test] + async fn test_new_with_options_rejects_hfile_format_for_regular_scan() { let result = HudiDataSource::new_with_options( V6Nonpartitioned.path_to_cow().as_str(), [ @@ -590,187 +1062,380 @@ mod tests { assert!( result.is_err(), - "Expected constructor fail-fast for unsupported base file format" - ); - let error_msg = result.unwrap_err().to_string(); - assert!( - error_msg.contains("Unsupported base file format 'hfile'"), - "Expected explicit base format error, got: {error_msg}" - ); - assert!( - error_msg.contains("only parquet is supported"), - "Expected parquet-only guidance, got: {error_msg}" + "HFile format should be rejected for regular DataFusion scans" ); + assert!(result.unwrap_err().to_string().contains("HFile")); } #[tokio::test] - async fn test_supports_filters_pushdown() { - let table_provider = HudiDataSource::new_with_options( + async fn test_new_with_options_rejects_invalid_base_file_format_config() { + let result = HudiDataSource::new_with_options( V6Nonpartitioned.path_to_cow().as_str(), - empty_options(), + [ + (HudiTableConfig::BaseFileFormat.as_ref(), "orc"), + (HudiInternalConfig::SkipConfigValidation.as_ref(), "true"), + ], ) - .await - .unwrap(); + .await; - let expr0 = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("name".to_string()))), - op: Operator::Eq, - right: Box::new(Expr::Literal( - ScalarValue::Utf8(Some("Alice".to_string())), - None, - )), - }); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("orc")); + } - let expr1 = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("intField".to_string()))), - op: Operator::Gt, - right: Box::new(Expr::Literal(ScalarValue::Int32(Some(20000)), None)), - }); + fn pushdown_test_schema() -> Schema { + Schema::new(vec![ + Field::new("byteField", DataType::Int8, false), + Field::new("name", DataType::Utf8, true), + Field::new("intField", DataType::Int32, true), + ]) + } - let expr2 = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name( - "nonexistent_column".to_string(), - ))), - op: Operator::Eq, - right: Box::new(Expr::Literal(ScalarValue::Int32(Some(1)), None)), - }); + fn partition_cols() -> Vec { + vec!["byteField".to_string()] + } + + fn col_lit(name: &str, op: Operator, lit: ScalarValue) -> Expr { + Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column::from_name(name.to_string()))), + op, + right: Box::new(Expr::Literal(lit, None)), + }) + } + + fn pushdown_support( + schema: &Schema, + partition_cols: &[String], + expr: &Expr, + ) -> TableProviderFilterPushDown { + HudiDataSource::filter_pushdown_support(schema, partition_cols, expr) + } + + async fn scan_hudi_exec_filter_triplets( + hudi: &HudiDataSource, + filters: &[Expr], + ) -> Vec<(String, String, Vec)> { + let ctx = SessionContext::new(); + let state = ctx.state(); + let plan = hudi.scan(&state, None, filters, None).await.unwrap(); + let exec = plan + .as_any() + .downcast_ref::() + .expect("scan should route to HudiScanExec"); + + exec.read_options() + .filters + .iter() + .map(|filter| { + ( + filter.field.clone(), + filter.operator.to_string(), + filter.values.clone(), + ) + }) + .collect() + } + + fn assert_scan_filter( + filters: &[(String, String, Vec)], + field: &str, + operator: &str, + values: &[&str], + ) { + assert!( + filters + .iter() + .any(|(actual_field, actual_operator, actual_values)| { + actual_field == field + && actual_operator == operator + && actual_values + .iter() + .map(String::as_str) + .eq(values.iter().copied()) + }), + "expected filter ({field}, {operator}, {values:?}) in {filters:?}" + ); + } - let expr3 = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("name".to_string()))), - op: Operator::NotEq, - right: Box::new(Expr::Literal( + #[test] + fn test_filter_pushdown_support_for_non_partitioned_schema() { + let schema = pushdown_test_schema(); + let partition_cols = vec![]; + let filters = [ + col_lit( + "name", + Operator::Eq, + ScalarValue::Utf8(Some("Alice".to_string())), + ), + col_lit("intField", Operator::Gt, ScalarValue::Int32(Some(20000))), + col_lit( + "nonexistent_column", + Operator::Eq, + ScalarValue::Int32(Some(1)), + ), + col_lit( + "name", + Operator::NotEq, ScalarValue::Utf8(Some("Diana".to_string())), - None, - )), - }); + ), + Expr::Literal(ScalarValue::Int32(Some(10)), None), + Expr::Not(Box::new(col_lit( + "intField", + Operator::Gt, + ScalarValue::Int32(Some(20000)), + ))), + ]; - let expr4 = Expr::Literal(ScalarValue::Int32(Some(10)), None); + let result = filters + .iter() + .map(|expr| pushdown_support(&schema, &partition_cols, expr)) + .collect::>(); + + assert_eq!( + result, + vec![ + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Unsupported, + TableProviderFilterPushDown::Inexact, + TableProviderFilterPushDown::Unsupported, + TableProviderFilterPushDown::Inexact, + ] + ); + } - let expr5 = Expr::Not(Box::new(Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("intField".to_string()))), - op: Operator::Gt, - right: Box::new(Expr::Literal(ScalarValue::Int32(Some(20000)), None)), - }))); + #[test] + fn test_filter_pushdown_exact_only_for_partition_equality() { + let schema = pushdown_test_schema(); + let partition_cols = partition_cols(); + + let partition_eq = col_lit("byteField", Operator::Eq, ScalarValue::Int8(Some(1))); + let partition_gt = col_lit("byteField", Operator::Gt, ScalarValue::Int8(Some(1))); + let non_partition_eq = col_lit( + "name", + Operator::Eq, + ScalarValue::Utf8(Some("Alice".to_string())), + ); - let filters = vec![&expr0, &expr1, &expr2, &expr3, &expr4, &expr5]; - let result = table_provider.supports_filters_pushdown(&filters).unwrap(); + assert_eq!( + pushdown_support(&schema, &partition_cols, &partition_eq), + TableProviderFilterPushDown::Exact + ); + assert_eq!( + pushdown_support(&schema, &partition_cols, &partition_gt), + TableProviderFilterPushDown::Inexact + ); + assert_eq!( + pushdown_support(&schema, &partition_cols, &non_partition_eq), + TableProviderFilterPushDown::Inexact + ); + } - assert_eq!(result.len(), 6); - // Non-partitioned table - all filters are Inexact (no partition columns) - assert_eq!(result[0], TableProviderFilterPushDown::Inexact); - assert_eq!(result[1], TableProviderFilterPushDown::Inexact); - assert_eq!(result[2], TableProviderFilterPushDown::Unsupported); - assert_eq!(result[3], TableProviderFilterPushDown::Inexact); - assert_eq!(result[4], TableProviderFilterPushDown::Unsupported); - assert_eq!(result[5], TableProviderFilterPushDown::Inexact); + #[test] + fn test_filter_pushdown_splits_conjunctions_for_classification() { + let schema = pushdown_test_schema(); + let partition_cols = partition_cols(); + + let partition_eq = col_lit("byteField", Operator::Eq, ScalarValue::Int8(Some(1))); + let second_partition_eq = col_lit("byteField", Operator::Eq, ScalarValue::Int8(Some(2))); + let non_partition_eq = col_lit( + "name", + Operator::Eq, + ScalarValue::Utf8(Some("Alice".to_string())), + ); + let unsupported = Expr::Literal(ScalarValue::Boolean(Some(true)), None); + + assert_eq!( + pushdown_support( + &schema, + &partition_cols, + &partition_eq.clone().and(second_partition_eq) + ), + TableProviderFilterPushDown::Exact + ); + assert_eq!( + pushdown_support( + &schema, + &partition_cols, + &partition_eq.clone().and(non_partition_eq) + ), + TableProviderFilterPushDown::Inexact + ); + assert_eq!( + pushdown_support(&schema, &partition_cols, &partition_eq.and(unsupported)), + TableProviderFilterPushDown::Inexact + ); } - #[tokio::test] - async fn test_supports_filters_pushdown_exact_for_partition_columns() { - // Use a partitioned table - byteField is the partition column - let table_provider = HudiDataSource::new_with_options( - V6SimplekeygenNonhivestyle.path_to_cow().as_str(), - empty_options(), + #[test] + fn test_filter_pushdown_between_in_list_and_or() { + let schema = pushdown_test_schema(); + let partition_cols = partition_cols(); + + let partition_between = Expr::Between(datafusion_expr::Between::new( + Box::new(Expr::Column(Column::from_name("byteField".to_string()))), + false, + Box::new(Expr::Literal(ScalarValue::Int8(Some(1)), None)), + Box::new(Expr::Literal(ScalarValue::Int8(Some(3)), None)), + )); + let partition_in = Expr::InList(datafusion_expr::expr::InList::new( + Box::new(Expr::Column(Column::from_name("byteField".to_string()))), + vec![ + Expr::Literal(ScalarValue::Int8(Some(1)), None), + Expr::Literal(ScalarValue::Int8(Some(2)), None), + ], + false, + )); + let or_expr = col_lit( + "name", + Operator::Eq, + ScalarValue::Utf8(Some("Alice".to_string())), ) - .await - .unwrap(); + .or(col_lit( + "name", + Operator::Eq, + ScalarValue::Utf8(Some("Bob".to_string())), + )); - // Filter on partition column (byteField) - should be Exact - let partition_filter = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("byteField".to_string()))), - op: Operator::Eq, - right: Box::new(Expr::Literal(ScalarValue::Int8(Some(1)), None)), - }); + assert_eq!( + pushdown_support(&schema, &partition_cols, &partition_between), + TableProviderFilterPushDown::Inexact + ); + assert_eq!( + pushdown_support(&schema, &partition_cols, &partition_in), + TableProviderFilterPushDown::Inexact + ); + assert_eq!( + pushdown_support(&schema, &partition_cols, &or_expr), + TableProviderFilterPushDown::Unsupported + ); + } - // Filter on non-partition column (name) - should be Inexact - let non_partition_filter = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("name".to_string()))), - op: Operator::Eq, - right: Box::new(Expr::Literal( - ScalarValue::Utf8(Some("Alice".to_string())), - None, - )), - }); + #[test] + fn test_scan_pushdown_exprs_separates_partition_filters_for_listing() { + let schema = pushdown_test_schema(); + let partition_cols = partition_cols(); + let partition_eq = col_lit("byteField", Operator::Eq, ScalarValue::Int8(Some(1))); + let non_partition_eq = col_lit( + "name", + Operator::Eq, + ScalarValue::Utf8(Some("Alice".to_string())), + ); + let partition_in = Expr::InList(datafusion_expr::expr::InList::new( + Box::new(Expr::Column(Column::from_name("byteField".to_string()))), + vec![ + Expr::Literal(ScalarValue::Int8(Some(1)), None), + Expr::Literal(ScalarValue::Int8(Some(2)), None), + ], + false, + )); + + let (partition_exprs, all_exprs) = HudiDataSource::split_scan_pushdown_exprs_for_schema( + &schema, + &partition_cols, + &[partition_eq.and(non_partition_eq), partition_in], + ); + let partition_filters = exprs_to_filters(&partition_exprs); + let all_filters = exprs_to_filters(&all_exprs); + + let partition_fields = partition_filters + .iter() + .map(|(field, _, _)| field.as_str()) + .collect::>(); + let all_fields = all_filters + .iter() + .map(|(field, _, _)| field.as_str()) + .collect::>(); - let filters = vec![&partition_filter, &non_partition_filter]; - let result = table_provider.supports_filters_pushdown(&filters).unwrap(); + assert_eq!(partition_fields, ["byteField", "byteField"]); + assert_eq!(all_fields, ["byteField", "name", "byteField"]); + } - assert_eq!(result.len(), 2); - // Partition column filter is Exact - assert_eq!(result[0], TableProviderFilterPushDown::Exact); - // Non-partition column filter is Inexact - assert_eq!(result[1], TableProviderFilterPushDown::Inexact); + #[test] + fn test_read_options_for_hudi_exec_strips_dropped_partition_filters() { + let hudi_configs = HudiConfigs::new([ + (HudiTableConfig::DropsPartitionFields, "true"), + (HudiTableConfig::PartitionFields, "region,country"), + ]); + let read_options = ReadOptions::new() + .with_filters([ + ("region", "=", "us"), + ("amount", ">", "10"), + ("txns.country", "=", "ca"), + ]) + .unwrap() + .with_projection(["txn_id", "amount"]); + + let actual = HudiDataSource::read_options_for_hudi_exec(&hudi_configs, &read_options); + + assert_eq!(actual.filters.len(), 1); + assert_eq!(actual.filters[0].field, "amount"); + assert_eq!(actual.projection, read_options.projection); + assert_eq!(actual.hudi_options, read_options.hudi_options); } #[tokio::test] - async fn test_supports_filters_pushdown_and_between() { - let table_provider = HudiDataSource::new_with_options( - V6Nonpartitioned.path_to_cow().as_str(), - empty_options(), - ) - .await - .unwrap(); + async fn test_scan_hudi_keeps_inexact_non_partition_filters() { + let hudi = HudiDataSource::new(V6SimplekeygenNonhivestyle.url_to_mor_parquet().as_str()) + .await + .unwrap(); + let read_options = hudi + .scan_read_options( + vec![("id".to_string(), ">".to_string(), "1".to_string())], + false, + ) + .unwrap(); + let flat_slices = hudi.table.get_file_slices(&read_options).await.unwrap(); - // AND expression: name = 'Alice' AND intField > 100 - let left = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("name".to_string()))), - op: Operator::Eq, - right: Box::new(Expr::Literal( - ScalarValue::Utf8(Some("Alice".to_string())), - None, - )), - }); - let right = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("intField".to_string()))), - op: Operator::Gt, - right: Box::new(Expr::Literal(ScalarValue::Int32(Some(100)), None)), - }); - let and_expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(left), - op: Operator::And, - right: Box::new(right), - }); + let plan = hudi + .scan_hudi(None, None, 1, flat_slices, read_options) + .await + .unwrap(); + let exec = plan + .as_any() + .downcast_ref::() + .expect("MOR snapshot scan should use HudiScanExec"); + + assert_eq!(exec.read_options().filters.len(), 1); + let filter = &exec.read_options().filters[0]; + assert_eq!(filter.field, "id"); + assert_eq!(filter.values, vec!["1".to_string()]); + } - // BETWEEN expression: intField BETWEEN 10 AND 100 - let between_expr = Expr::Between(datafusion_expr::Between::new( - Box::new(Expr::Column(Column::from_name("intField".to_string()))), - false, - Box::new(Expr::Literal(ScalarValue::Int32(Some(10)), None)), - Box::new(Expr::Literal(ScalarValue::Int32(Some(100)), None)), - )); + #[tokio::test] + async fn test_scan_mor_snapshot_keeps_partition_and_non_partition_filters_for_hudi_exec() { + let hudi = HudiDataSource::new(V6SimplekeygenNonhivestyle.url_to_mor_parquet().as_str()) + .await + .unwrap(); + let partition_filter = col_lit("byteField", Operator::Eq, ScalarValue::Int32(Some(10))); + let non_partition_filter = col_lit("id", Operator::Gt, ScalarValue::Int32(Some(1))); - // OR expression - should be unsupported - let or_left = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("name".to_string()))), - op: Operator::Eq, - right: Box::new(Expr::Literal( - ScalarValue::Utf8(Some("Alice".to_string())), - None, - )), - }); - let or_right = Expr::BinaryExpr(BinaryExpr { - left: Box::new(Expr::Column(Column::from_name("name".to_string()))), - op: Operator::Eq, - right: Box::new(Expr::Literal( - ScalarValue::Utf8(Some("Bob".to_string())), - None, - )), - }); - let or_expr = Expr::BinaryExpr(BinaryExpr { - left: Box::new(or_left), - op: Operator::Or, - right: Box::new(or_right), - }); + let filters = + scan_hudi_exec_filter_triplets(&hudi, &[partition_filter, non_partition_filter]).await; + + assert_scan_filter(&filters, "byteField", "=", &["10"]); + assert_scan_filter(&filters, "id", ">", &["1"]); + } + + #[tokio::test] + async fn test_scan_lance_keeps_partition_and_non_partition_filters_for_hudi_exec() { + let hudi = HudiDataSource::new(V9LanceTxnsSimple.url_to_cow().as_str()) + .await + .unwrap(); + let partition_filter = col_lit( + "region", + Operator::Eq, + ScalarValue::Utf8(Some("us".to_string())), + ); + let non_partition_filter = col_lit( + "txn_id", + Operator::Eq, + ScalarValue::Utf8(Some("TXN-001".to_string())), + ); - let filters = vec![&and_expr, &between_expr, &or_expr]; - let result = table_provider.supports_filters_pushdown(&filters).unwrap(); + let filters = + scan_hudi_exec_filter_triplets(&hudi, &[partition_filter, non_partition_filter]).await; - assert_eq!(result.len(), 3); - // AND expression is supported - assert_eq!(result[0], TableProviderFilterPushDown::Inexact); - // BETWEEN expression is supported - assert_eq!(result[1], TableProviderFilterPushDown::Inexact); - // OR expression is not supported - assert_eq!(result[2], TableProviderFilterPushDown::Unsupported); + assert_scan_filter(&filters, "region", "=", &["us"]); + assert_scan_filter(&filters, "txn_id", "=", &["TXN-001"]); } } diff --git a/crates/datafusion/src/util/expr.rs b/crates/datafusion/src/util/expr.rs index 306d5e45..9e10d41d 100644 --- a/crates/datafusion/src/util/expr.rs +++ b/crates/datafusion/src/util/expr.rs @@ -18,6 +18,7 @@ */ use datafusion::logical_expr::Operator; +use datafusion_common::ScalarValue; use datafusion_expr::expr::InList; use datafusion_expr::{Between, BinaryExpr, Expr}; use hudi_core::expr::filter::{Filter as HudiFilter, col}; @@ -108,7 +109,7 @@ fn binary_expr_to_filter(binary_expr: &BinaryExpr) -> Option { }; let field = col(column.name()); - let lit_str = literal.to_string(); + let lit_str = scalar_to_filter_value(literal); let filter = match binary_expr.op { Operator::Eq => field.eq(lit_str), @@ -153,7 +154,7 @@ fn between_to_filters(between: &Between) -> Vec { // Extract literal values from low and high bounds let low_str = match &*between.low { - Expr::Literal(lit, _) => lit.to_string(), + Expr::Literal(lit, _) => scalar_to_filter_value(lit), _ => { warn!( "BETWEEN low bound is not a literal for column '{column_name}', skipping pushdown" @@ -163,7 +164,7 @@ fn between_to_filters(between: &Between) -> Vec { }; let high_str = match &*between.high { - Expr::Literal(lit, _) => lit.to_string(), + Expr::Literal(lit, _) => scalar_to_filter_value(lit), _ => { warn!( "BETWEEN high bound is not a literal for column '{column_name}', skipping pushdown" @@ -201,7 +202,7 @@ fn inlist_expr_to_filter(in_list: &InList) -> Option { .list .iter() .filter_map(|expr| match expr { - Expr::Literal(lit, _) => Some(lit.to_string()), + Expr::Literal(lit, _) => Some(scalar_to_filter_value(lit)), _ => None, }) .collect(); @@ -219,6 +220,69 @@ fn inlist_expr_to_filter(in_list: &InList) -> Option { } } +/// Stringifies a DataFusion literal for Hudi's string-typed `Filter` API. +/// +/// `ScalarValue::Display` is lossy for decimals — it prints the unscaled +/// integer (e.g. `Decimal128(7500, 10, 2)` becomes `"7500"` rather than +/// `"75.00"`) — so those are formatted explicitly here. Other types fall +/// through to `to_string()`. +/// +/// Tracking the typed-value refactor that would eliminate this round trip: +/// . +fn scalar_to_filter_value(literal: &ScalarValue) -> String { + match literal { + ScalarValue::Decimal32(Some(value), _, scale) => { + format_decimal_value(*value as i128, *scale) + } + ScalarValue::Decimal64(Some(value), _, scale) => { + format_decimal_value(*value as i128, *scale) + } + ScalarValue::Decimal128(Some(value), _, scale) => format_decimal_value(*value, *scale), + ScalarValue::Decimal256(Some(value), _, scale) => { + format_decimal_digits(value.to_string(), *scale) + } + ScalarValue::Decimal32(None, _, _) + | ScalarValue::Decimal64(None, _, _) + | ScalarValue::Decimal128(None, _, _) + | ScalarValue::Decimal256(None, _, _) => "NULL".to_string(), + _ => literal.to_string(), + } +} + +fn format_decimal_value(value: i128, scale: i8) -> String { + format_decimal_digits(value.to_string(), scale) +} + +fn format_decimal_digits(mut digits: String, scale: i8) -> String { + let negative = digits.starts_with('-'); + if negative { + digits.remove(0); + } + + if scale <= 0 { + digits.push_str(&"0".repeat((-scale) as usize)); + return if negative { + format!("-{digits}") + } else { + digits + }; + } + + let scale = scale as usize; + if digits.len() <= scale { + let padding = "0".repeat(scale + 1 - digits.len()); + digits = format!("{padding}{digits}"); + } + + let split_at = digits.len() - scale; + let (whole, fractional) = digits.split_at(split_at); + if negative { + format!("-{whole}.{fractional}") + } else { + format!("{whole}.{fractional}") + } +} + #[cfg(test)] mod tests { use super::*; @@ -480,6 +544,25 @@ mod tests { assert_eq!(result[1].2, "20"); } + #[test] + fn test_convert_decimal_literal() { + let expr = Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("amount")), + Operator::Eq, + Box::new(Expr::Literal( + ScalarValue::Decimal128(Some(7500), 10, 2), + None, + )), + )); + + let result = exprs_to_filters(&[expr]); + + assert_eq!(result.len(), 1); + assert_eq!(result[0].0, "amount"); + assert_eq!(result[0].1, "="); + assert_eq!(result[0].2, "75.00"); + } + #[test] fn test_convert_not_between_returns_empty() { // NOT BETWEEN cannot be represented in current filter model diff --git a/crates/datafusion/tests/plan_tests.rs b/crates/datafusion/tests/plan_tests.rs index 47156edc..419d2ced 100644 --- a/crates/datafusion/tests/plan_tests.rs +++ b/crates/datafusion/tests/plan_tests.rs @@ -30,11 +30,11 @@ use datafusion::execution::session_state::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::{DataFusionError, ScalarValue}; -use hudi_core::config::read::HudiReadConfig::InputPartitions; +use hudi_core::config::read::HudiReadConfig::{InputPartitions, UseReadOptimizedMode}; use hudi_core::config::util::empty_options; use hudi_core::metadata::meta_field::MetaField; use hudi_datafusion::{HudiDataSource, HudiTableFactory}; -use hudi_test::util::{get_bool_column, get_i32_column, get_str_column}; +use hudi_test::util::{explain_physical_plan, get_bool_column, get_i32_column, get_str_column}; use hudi_test::{SampleTable, assert_arrow_field_names_eq}; // ============================================================================ @@ -100,6 +100,22 @@ where Ok(ctx) } +async fn register_uri_as_table( + table_name: &str, + base_uri: &str, + options: I, +) -> Result +where + I: IntoIterator, + K: AsRef, + V: Into, +{ + let ctx = create_test_session().await; + let hudi = HudiDataSource::new_with_options(base_uri, options).await?; + ctx.register_table(table_name, Arc::new(hudi))?; + Ok(ctx) +} + fn concat_as_sql_options(options: I) -> String where I: IntoIterator, @@ -431,3 +447,364 @@ mod v8_tests { assert_eq!(get_bool_column(rb, "isActive"), &[false, true]); } } + +mod dispatch_tests { + use super::*; + use hudi_test::QuickstartTripsTable; + use hudi_test::SampleTable::{V6Nonpartitioned, V6SimplekeygenNonhivestyle}; + + #[tokio::test] + async fn test_parquet_cow_uses_data_source_exec() { + let ctx = register_table_direct(&V6Nonpartitioned, empty_options()) + .await + .unwrap(); + + let sql = format!("SELECT id FROM {}", V6Nonpartitioned.as_ref()); + let plan = explain_physical_plan(&ctx, &sql).await; + + assert!( + plan.contains("DataSourceExec"), + "Parquet COW should use DataSourceExec. Plan: {plan}" + ); + assert!( + !plan.contains("HudiScanExec"), + "Parquet COW should not use HudiScanExec. Plan: {plan}" + ); + } + + #[tokio::test] + async fn test_parquet_mor_read_optimized_uses_data_source_exec() { + let base_url = V6Nonpartitioned.url_to_mor_parquet(); + let ctx = register_uri_as_table( + "mor_ro", + base_url.as_str(), + [(UseReadOptimizedMode.as_ref(), "true")], + ) + .await + .unwrap(); + + let plan = explain_physical_plan(&ctx, "SELECT id FROM mor_ro").await; + + assert!( + plan.contains("DataSourceExec"), + "Parquet MOR read-optimized should use DataSourceExec. Plan: {plan}" + ); + assert!( + !plan.contains("HudiScanExec"), + "Parquet MOR read-optimized should not use HudiScanExec. Plan: {plan}" + ); + } + + #[tokio::test] + async fn test_parquet_mor_snapshot_uses_hudi_scan_exec() { + let base_url = V6Nonpartitioned.url_to_mor_parquet(); + let ctx = register_uri_as_table("mor_snapshot", base_url.as_str(), empty_options()) + .await + .unwrap(); + + let plan = explain_physical_plan(&ctx, "SELECT id FROM mor_snapshot").await; + + assert!( + plan.contains("HudiScanExec"), + "Parquet MOR snapshot should use HudiScanExec. Plan: {plan}" + ); + } + + #[tokio::test] + async fn test_partitioned_parquet_mor_snapshot_prunes_file_slices() { + let base_url = V6SimplekeygenNonhivestyle.url_to_mor_parquet(); + let ctx = register_uri_as_table( + "partitioned_mor_snapshot", + base_url.as_str(), + [(InputPartitions.as_ref(), "2")], + ) + .await + .unwrap(); + + let no_filter_plan = + explain_physical_plan(&ctx, "SELECT id FROM partitioned_mor_snapshot").await; + assert!( + no_filter_plan.contains("HudiScanExec"), + "Partitioned Parquet MOR snapshot should use HudiScanExec. Plan: {no_filter_plan}" + ); + assert!( + !no_filter_plan.contains("DataSourceExec"), + "Partitioned Parquet MOR snapshot should not use DataSourceExec. Plan: {no_filter_plan}" + ); + assert!( + !no_filter_plan.contains("ParquetSource"), + "Partitioned Parquet MOR snapshot should not use ParquetSource. Plan: {no_filter_plan}" + ); + assert!( + no_filter_plan.contains("input_partitions=2"), + "Plan should preserve configured input_partitions=2. Plan: {no_filter_plan}" + ); + assert!( + no_filter_plan.contains("partitions=2"), + "Unfiltered scan should be split across 2 execution partitions. Plan: {no_filter_plan}" + ); + assert!( + no_filter_plan.contains("file_slices=3"), + "Unfiltered scan should include all 3 file slices. Plan: {no_filter_plan}" + ); + + let filtered_plan = explain_physical_plan( + &ctx, + "SELECT id FROM partitioned_mor_snapshot WHERE byteField = 10", + ) + .await; + assert!( + filtered_plan.contains("HudiScanExec"), + "Partition-filtered Parquet MOR snapshot should use HudiScanExec. Plan: {filtered_plan}" + ); + assert!( + filtered_plan.contains("input_partitions=2"), + "Partition-filtered plan should preserve configured input_partitions=2. Plan: {filtered_plan}" + ); + assert!( + filtered_plan.contains("partitions=1"), + "Partition filter should reduce the scan to 1 execution partition. Plan: {filtered_plan}" + ); + assert!( + filtered_plan.contains("file_slices=1"), + "Partition filter should reduce the scan to 1 file slice. Plan: {filtered_plan}" + ); + assert!( + !filtered_plan.contains("FilterExec"), + "Exact partition filter should not leave a residual FilterExec. Plan: {filtered_plan}" + ); + } + + #[tokio::test] + async fn test_lance_mor_snapshot_uses_hudi_scan_exec() { + let base_url = QuickstartTripsTable::V9TripsLance.url_to_mor_avro(); + let ctx = register_uri_as_table("lance_mor_snapshot", base_url.as_str(), empty_options()) + .await + .unwrap(); + + let plan = explain_physical_plan(&ctx, "SELECT uuid FROM lance_mor_snapshot").await; + + assert!( + plan.contains("HudiScanExec"), + "Lance MOR snapshot should use HudiScanExec. Plan: {plan}" + ); + assert!( + !plan.contains("DataSourceExec"), + "Lance MOR snapshot should not use DataSourceExec. Plan: {plan}" + ); + assert!( + !plan.contains("ParquetSource"), + "Lance MOR snapshot should not use ParquetSource. Plan: {plan}" + ); + } + + #[tokio::test] + async fn test_hudi_scan_exec_display_includes_limit() { + let base_url = V6Nonpartitioned.url_to_mor_parquet(); + let ctx = register_uri_as_table("mor_snapshot", base_url.as_str(), empty_options()) + .await + .unwrap(); + + let plan = explain_physical_plan(&ctx, "SELECT id FROM mor_snapshot LIMIT 1").await; + + assert!( + plan.contains("HudiScanExec"), + "Parquet MOR snapshot should use HudiScanExec. Plan: {plan}" + ); + assert!( + plan.contains("limit=Some(1)"), + "HudiScanExec display should advertise pushed limit. Plan: {plan}" + ); + } + + #[tokio::test] + async fn test_mor_snapshot_with_full_pruning_returns_empty() { + let ctx = register_table_direct(&V6SimplekeygenNonhivestyle, empty_options()) + .await + .unwrap(); + let sql = format!( + "SELECT id FROM {} WHERE byteField = 99", + V6SimplekeygenNonhivestyle.as_ref() + ); + + let plan = explain_physical_plan(&ctx, &sql).await; + assert!( + plan.contains("HudiScanExec"), + "MOR snapshot with aggressive filter pruning should use HudiScanExec. Plan: {plan}" + ); + + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let rows: usize = batches.iter().map(|batch| batch.num_rows()).sum(); + assert_eq!(rows, 0); + } +} + +// ============================================================================ +// Lance Table Tests +// ============================================================================ + +mod lance_tests { + use super::*; + use arrow_array::{Float64Array, Int32Array, Int64Array}; + use hudi_test::SampleTable::{V9LanceNonpartitioned, V9LanceTxnsSimple}; + use std::collections::HashMap; + + #[tokio::test] + async fn test_datafusion_read_lance_cow_table() { + let test_table = V9LanceNonpartitioned; + let ctx = register_table_direct(&test_table, empty_options()) + .await + .unwrap(); + + let sql = format!("SELECT * FROM {}", test_table.as_ref()); + let df = ctx.sql(&sql).await.unwrap(); + let batches = df.collect().await.unwrap(); + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert!(total_rows > 0, "Lance COW table should return rows"); + + let schema = batches[0].schema(); + let field_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert!(field_names.contains(&"id"), "Schema should contain 'id'"); + assert!( + field_names.contains(&"name"), + "Schema should contain 'name'" + ); + } + + #[tokio::test] + async fn test_datafusion_lance_table_has_meta_fields() { + let test_table = V9LanceNonpartitioned; + let ctx = register_table_direct(&test_table, empty_options()) + .await + .unwrap(); + + let sql = format!("SELECT * FROM {} LIMIT 1", test_table.as_ref()); + let df = ctx.sql(&sql).await.unwrap(); + let schema = df.schema(); + + assert!( + schema.field_with_name(None, "_hoodie_commit_time").is_ok(), + "Lance table schema should include _hoodie_commit_time meta field" + ); + assert!( + schema.field_with_name(None, "_hoodie_record_key").is_ok(), + "Lance table schema should include _hoodie_record_key meta field" + ); + } + + #[tokio::test] + async fn test_datafusion_lance_uses_hudi_scan_exec() { + let test_table = V9LanceNonpartitioned; + let ctx = register_table_direct(&test_table, empty_options()) + .await + .unwrap(); + + let sql = format!("SELECT id, name FROM {}", test_table.as_ref()); + let explaining_df = ctx.sql(&sql).await.unwrap().explain(false, true).unwrap(); + let explaining_rb = explaining_df.collect().await.unwrap(); + let explaining_rb = explaining_rb.first().unwrap(); + let plan = get_str_column(explaining_rb, "plan").join(""); + + assert!( + plan.contains("HudiScanExec"), + "Lance table should use HudiScanExec, not DataSourceExec. Plan: {plan}" + ); + } + + #[tokio::test] + async fn test_datafusion_lance_cow_count_star_projection_elimination() { + // Regression test: DataFusion forwards `Some(vec![])` as the + // projection for COUNT(*). lance-file rejects zero-column + // projections, so the Lance reader must emit zero-column + // batches with the row count derived from metadata. + let test_table = V9LanceNonpartitioned; + let ctx = register_table_direct(&test_table, empty_options()) + .await + .unwrap(); + + let count_sql = format!("SELECT COUNT(*) FROM {}", test_table.as_ref()); + let count_batches = ctx.sql(&count_sql).await.unwrap().collect().await.unwrap(); + let counted = count_batches[0] + .column(0) + .as_any() + .downcast_ref::() + .expect("COUNT(*) returns Int64") + .value(0); + + let rows_sql = format!("SELECT id FROM {}", test_table.as_ref()); + let row_batches = ctx.sql(&rows_sql).await.unwrap().collect().await.unwrap(); + let expected: i64 = row_batches.iter().map(|b| b.num_rows() as i64).sum(); + + assert_eq!(counted, expected); + } + + #[tokio::test] + async fn test_datafusion_lance_cow_query_applies_hudi_updates_deletes_and_inserts() { + let test_table = V9LanceNonpartitioned; + let ctx = register_table_direct(&test_table, empty_options()) + .await + .unwrap(); + + let sql = format!("SELECT id, score FROM {} ORDER BY id", test_table.as_ref()); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let mut rows = HashMap::new(); + for batch in batches { + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let scores = batch + .column_by_name("score") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + for row_idx in 0..batch.num_rows() { + rows.insert(ids.value(row_idx), scores.value(row_idx)); + } + } + + let mut ids = rows.keys().copied().collect::>(); + ids.sort_unstable(); + assert_eq!(ids, vec![1, 2, 3, 5, 6, 7, 8, 9, 10]); + assert!(!rows.contains_key(&4)); + assert!((rows[&1] - 0.96).abs() < 1e-9); + assert!((rows[&2] - 0.93).abs() < 1e-9); + assert!(rows.contains_key(&9)); + assert!(rows.contains_key(&10)); + } + + #[tokio::test] + async fn test_datafusion_partitioned_lance_filter_uses_hudi_scan_exec() { + let test_table = V9LanceTxnsSimple; + let ctx = register_table_direct(&test_table, empty_options()) + .await + .unwrap(); + + let sql = format!( + "SELECT txn_id, txn_type FROM {} WHERE region = 'us' ORDER BY txn_id", + test_table.as_ref() + ); + let plan = explain_physical_plan(&ctx, &sql).await; + assert!( + plan.contains("HudiScanExec"), + "Partitioned Lance table should use HudiScanExec. Plan: {plan}" + ); + + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + let txn_ids = batches + .iter() + .flat_map(|batch| get_str_column(batch, "txn_id")) + .collect::>(); + let txn_types = batches + .iter() + .flat_map(|batch| get_str_column(batch, "txn_type")) + .collect::>(); + + assert_eq!(txn_ids, ["TXN-001", "TXN-003", "TXN-013", "TXN-014"]); + assert_eq!(txn_types, ["reversal", "transfer", "debit", "debit"]); + } +} diff --git a/crates/datafusion/tests/query_tests.rs b/crates/datafusion/tests/query_tests.rs index 7039062d..e9c9089a 100644 --- a/crates/datafusion/tests/query_tests.rs +++ b/crates/datafusion/tests/query_tests.rs @@ -17,75 +17,497 @@ * under the License. */ -//! DataFusion read tests for v9 txns tables. +//! DataFusion read tests for Hudi sample tables. +use std::fs::{self, File}; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use arrow_array::{ArrayRef, Float64Array, RecordBatch, StringArray}; +use arrow_schema::{Field, Schema, SchemaRef}; +use datafusion::prelude::SessionContext; +use hudi_core::config::internal::HudiInternalConfig; +use hudi_core::config::read::HudiReadConfig::{InputPartitions, UseReadOptimizedMode}; +use hudi_core::config::table::HudiTableConfig; +use hudi_core::table::{ReadOptions, Table}; +use hudi_datafusion::HudiDataSource; +use hudi_test::QuickstartTripsTable; use hudi_test::SampleTable; -use hudi_test::v9_verification::verify_v9_txns_table; +use hudi_test::util::explain_physical_plan; +use hudi_test::v9_verification::{ + verify_partitioned_records, verify_v9_txns_table, verify_v9_txns_table_snapshot, +}; +use parquet::arrow::ArrowWriter; +use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + +fn rider_fare_rows(batches: &[RecordBatch]) -> Vec<(String, f64)> { + let mut rows = Vec::new(); + for batch in batches { + let riders = batch + .column_by_name("rider") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let fares = batch + .column_by_name("fare") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + for row_idx in 0..batch.num_rows() { + rows.push((riders.value(row_idx).to_string(), fares.value(row_idx))); + } + } + rows.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + rows +} + +fn uuid_rider_fare_rows(batches: &[RecordBatch]) -> Vec<(String, String, f64)> { + let mut rows = Vec::new(); + for batch in batches { + rows.extend(QuickstartTripsTable::uuid_rider_and_fare(batch)); + } + rows.sort_unstable_by(|left, right| left.0.cmp(&right.0)); + rows +} + +fn fares_for_rider(rows: &[(String, f64)], rider: &str) -> Vec { + let mut fares = rows + .iter() + .filter_map(|(row_rider, fare)| (row_rider == rider).then_some(*fare)) + .collect::>(); + fares.sort_unstable_by(f64::total_cmp); + fares +} + +fn uuid_fares_for_rider(rows: &[(String, String, f64)], rider: &str) -> Vec { + let mut fares = rows + .iter() + .filter_map(|(_, row_rider, fare)| (row_rider == rider).then_some(*fare)) + .collect::>(); + fares.sort_unstable_by(f64::total_cmp); + fares +} + +fn id_name_active_rows(batches: &[RecordBatch]) -> Vec<(i32, String, bool)> { + let mut rows = Vec::new(); + for batch in batches { + rows.extend( + SampleTable::sample_data_order_by_id(batch) + .into_iter() + .map(|(id, name, active)| (id, name.to_string(), active)), + ); + } + rows.sort_unstable_by_key(|(id, _, _)| *id); + rows +} + +fn collect_parquet_files(dir: &Path, files: &mut Vec) { + for entry in fs::read_dir(dir).unwrap() { + let path = entry.unwrap().path(); + if path + .file_name() + .is_some_and(|name| name.to_string_lossy() == ".hoodie") + { + continue; + } + if path.is_dir() { + collect_parquet_files(&path, files); + } else if path.extension().is_some_and(|ext| ext == "parquet") { + files.push(path); + } + } +} + +fn record_batch_without_column( + batch: &RecordBatch, + schema: SchemaRef, + indices: &[usize], +) -> RecordBatch { + let columns: Vec = indices + .iter() + .map(|index| batch.column(*index).clone()) + .collect(); + RecordBatch::try_new(schema, columns).unwrap() +} + +fn rewrite_parquet_without_column(path: &Path, column_name: &str) { + let input = File::open(path).unwrap(); + let builder = ParquetRecordBatchReaderBuilder::try_new(input).unwrap(); + let original_schema = builder.schema().clone(); + let indices: Vec = original_schema + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| (field.name() != column_name).then_some(index)) + .collect(); + assert!( + indices.len() < original_schema.fields().len(), + "test fixture parquet file should contain {column_name}: {path:?}" + ); + + let projected_fields: Vec = indices + .iter() + .map(|index| original_schema.field(*index).as_ref().clone()) + .collect(); + let projected_schema = SchemaRef::new(Schema::new(projected_fields)); + let temp_path = path.with_extension("parquet.tmp"); + let output = File::create(&temp_path).unwrap(); + let mut writer = ArrowWriter::try_new(output, projected_schema.clone(), None).unwrap(); + + for batch in builder.build().unwrap() { + let batch = batch.unwrap(); + let projected_batch = + record_batch_without_column(&batch, projected_schema.clone(), &indices); + writer.write(&projected_batch).unwrap(); + } + + writer.close().unwrap(); + fs::rename(temp_path, path).unwrap(); +} + +fn set_drop_partition_columns(table_path: &Path) { + let props_path = table_path.join(".hoodie/hoodie.properties"); + let props = fs::read_to_string(&props_path).unwrap(); + let key = HudiTableConfig::DropsPartitionFields.as_ref(); + let mut replaced = false; + let mut lines = Vec::new(); + + for line in props.lines() { + if line + .split_once('=') + .is_some_and(|(line_key, _)| line_key == key) + { + lines.push(format!("{key}=true")); + replaced = true; + } else { + lines.push(line.to_string()); + } + } + + if !replaced { + lines.push(format!("{key}=true")); + } + + fs::write(props_path, format!("{}\n", lines.join("\n"))).unwrap(); +} + +fn dropped_partition_column_mor_table() -> PathBuf { + let table_path = + PathBuf::from(SampleTable::V6SimplekeygenNonhivestyle.path_to_mor_parquet_fresh()); + set_drop_partition_columns(&table_path); + + let mut parquet_files = Vec::new(); + collect_parquet_files(&table_path, &mut parquet_files); + assert!(!parquet_files.is_empty()); + for parquet_file in parquet_files { + rewrite_parquet_without_column(&parquet_file, "byteField"); + } + + table_path +} // ============================================================================ // COW tests // ============================================================================ -#[tokio::test] -async fn test_v9_txns_cow_simple_nometa() { - verify_v9_txns_table(&SampleTable::V9TxnsSimpleNometa, true, true).await; +fn v9_txns_cases() -> [(SampleTable, bool); 6] { + [ + (SampleTable::V9TxnsSimpleNometa, true), + (SampleTable::V9TxnsSimpleMeta, true), + (SampleTable::V9TxnsComplexNometa, true), + (SampleTable::V9TxnsComplexMeta, true), + (SampleTable::V9TxnsNonpartNometa, false), + (SampleTable::V9TxnsNonpartMeta, false), + ] } #[tokio::test] -async fn test_v9_txns_cow_simple_meta() { - verify_v9_txns_table(&SampleTable::V9TxnsSimpleMeta, true, true).await; +async fn test_v9_txns_cow_tables() { + for (table, partitioned) in v9_txns_cases() { + verify_v9_txns_table(&table, true, partitioned).await; + } } +// ============================================================================ +// MOR tests (read-optimized mode, after compaction + clustering) +// ============================================================================ + #[tokio::test] -async fn test_v9_txns_cow_complex_nometa() { - verify_v9_txns_table(&SampleTable::V9TxnsComplexNometa, true, true).await; +async fn test_v9_txns_mor_read_optimized_tables() { + for (table, partitioned) in v9_txns_cases() { + verify_v9_txns_table(&table, false, partitioned).await; + } } +// ============================================================================ +// MOR tests (snapshot mode, after compaction + clustering) +// ============================================================================ + #[tokio::test] -async fn test_v9_txns_cow_complex_meta() { - verify_v9_txns_table(&SampleTable::V9TxnsComplexMeta, true, true).await; +async fn test_v9_txns_mor_snapshot_simple_nometa() { + let base_url = SampleTable::V9TxnsSimpleNometa.url_to_mor_avro(); + let ctx = SessionContext::new(); + let hudi = + HudiDataSource::new_with_options(base_url.as_str(), [(InputPartitions.as_ref(), "2")]) + .await + .unwrap(); + ctx.register_table("txns", Arc::new(hudi)).unwrap(); + + verify_partitioned_records(&ctx).await; + + let plan = explain_physical_plan(&ctx, "SELECT txn_id FROM txns WHERE region = 'us'").await; + assert!(plan.contains("HudiScanExec")); + assert!(plan.contains("input_partitions=2")); } #[tokio::test] -async fn test_v9_txns_cow_nonpart_nometa() { - verify_v9_txns_table(&SampleTable::V9TxnsNonpartNometa, true, false).await; +async fn test_v9_txns_mor_snapshot_simple_meta() { + verify_v9_txns_table_snapshot(&SampleTable::V9TxnsSimpleMeta, true).await; } #[tokio::test] -async fn test_v9_txns_cow_nonpart_meta() { - verify_v9_txns_table(&SampleTable::V9TxnsNonpartMeta, true, false).await; +async fn test_v9_txns_mor_snapshot_complex_nometa() { + verify_v9_txns_table_snapshot(&SampleTable::V9TxnsComplexNometa, true).await; } -// ============================================================================ -// MOR tests (read-optimized mode, after compaction + clustering) -// ============================================================================ +#[tokio::test] +async fn test_v9_txns_mor_snapshot_complex_meta() { + verify_v9_txns_table_snapshot(&SampleTable::V9TxnsComplexMeta, true).await; +} #[tokio::test] -async fn test_v9_txns_mor_simple_nometa() { - verify_v9_txns_table(&SampleTable::V9TxnsSimpleNometa, false, true).await; +async fn test_v9_txns_mor_snapshot_nonpart_nometa() { + verify_v9_txns_table_snapshot(&SampleTable::V9TxnsNonpartNometa, false).await; } #[tokio::test] -async fn test_v9_txns_mor_simple_meta() { - verify_v9_txns_table(&SampleTable::V9TxnsSimpleMeta, false, true).await; +async fn test_v9_txns_mor_snapshot_nonpart_meta() { + verify_v9_txns_table_snapshot(&SampleTable::V9TxnsNonpartMeta, false).await; } #[tokio::test] -async fn test_v9_txns_mor_complex_nometa() { - verify_v9_txns_table(&SampleTable::V9TxnsComplexNometa, false, true).await; +async fn test_mor_snapshot_query_matches_core_log_merged_read() { + let base_url = QuickstartTripsTable::V8Trips8I3U1D.url_to_mor_avro(); + let projection = ["rider", "fare"]; + + let table = Table::new(base_url.path()).await.unwrap(); + let expected_batches = table + .read(&ReadOptions::new().with_projection(projection)) + .await + .unwrap(); + let expected_rows = rider_fare_rows(&expected_batches); + let read_optimized_batches = table + .read( + &ReadOptions::new() + .with_projection(projection) + .with_hudi_option(UseReadOptimizedMode.as_ref(), "true"), + ) + .await + .unwrap(); + let read_optimized_rows = rider_fare_rows(&read_optimized_batches); + + let ctx = SessionContext::new(); + let hudi = HudiDataSource::new(base_url.as_str()).await.unwrap(); + ctx.register_table("trips", Arc::new(hudi)).unwrap(); + let actual_batches = ctx + .sql("SELECT rider, fare FROM trips") + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual_rows = rider_fare_rows(&actual_batches); + + assert_eq!(actual_rows, expected_rows); + assert_eq!(actual_rows.len(), 6); + assert!(actual_rows.iter().all(|(rider, _)| rider != "rider-F")); + assert!(actual_rows.iter().all(|(rider, _)| rider != "rider-J")); + assert_ne!( + actual_rows, read_optimized_rows, + "snapshot read should merge Parquet MOR log files" + ); + assert_eq!(fares_for_rider(&actual_rows, "rider-A"), vec![0.0]); + assert_eq!(fares_for_rider(&actual_rows, "rider-G"), vec![0.0]); } #[tokio::test] -async fn test_v9_txns_mor_complex_meta() { - verify_v9_txns_table(&SampleTable::V9TxnsComplexMeta, false, true).await; +async fn test_lance_mor_snapshot_query_matches_core_log_merged_read() { + let base_url = QuickstartTripsTable::V9TripsLance.url_to_mor_avro(); + + let table = Table::new(base_url.path()).await.unwrap(); + let expected_batches = table + .read(&ReadOptions::new().with_projection(["uuid", "rider", "fare"])) + .await + .unwrap(); + let expected_rows = uuid_rider_fare_rows(&expected_batches); + let read_optimized_batches = table + .read( + &ReadOptions::new() + .with_projection(["uuid", "rider", "fare"]) + .with_hudi_option(UseReadOptimizedMode.as_ref(), "true"), + ) + .await + .unwrap(); + let read_optimized_rows = uuid_rider_fare_rows(&read_optimized_batches); + + let ctx = SessionContext::new(); + let hudi = HudiDataSource::new(base_url.as_str()).await.unwrap(); + ctx.register_table("trips", Arc::new(hudi)).unwrap(); + let actual_batches = ctx + .sql("SELECT uuid, rider, fare FROM trips") + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual_rows = uuid_rider_fare_rows(&actual_batches); + + assert_eq!(actual_rows, expected_rows); + assert_eq!(actual_rows.len(), 12); + assert_ne!( + actual_rows, read_optimized_rows, + "snapshot read should merge the active Lance MOR data log" + ); + + // The archived fixture has one active data log, for rider-A. The later + // SQL delete operations are not present as active data-log delete blocks, + // so the merge-specific assertion here is the rider-A fare update. + assert_eq!(uuid_fares_for_rider(&actual_rows, "rider-A"), vec![0.0]); + assert_eq!( + uuid_fares_for_rider(&read_optimized_rows, "rider-A"), + vec![19.1] + ); } #[tokio::test] -async fn test_v9_txns_mor_nonpart_nometa() { - verify_v9_txns_table(&SampleTable::V9TxnsNonpartNometa, false, false).await; +async fn test_partitioned_parquet_mor_snapshot_query_matches_core_after_partition_pruning() { + let base_url = SampleTable::V6SimplekeygenNonhivestyle.url_to_mor_parquet(); + let projection = ["id", "name", "isActive"]; + + let table = Table::new(base_url.path()).await.unwrap(); + let expected_all_batches = table + .read(&ReadOptions::new().with_projection(projection)) + .await + .unwrap(); + let expected_all_rows = id_name_active_rows(&expected_all_batches); + let expected_filtered_batches = table + .read( + &ReadOptions::new() + .with_filters([("byteField", "=", "10")]) + .unwrap() + .with_projection(projection), + ) + .await + .unwrap(); + let expected_filtered_rows = id_name_active_rows(&expected_filtered_batches); + let read_optimized_filtered_batches = table + .read( + &ReadOptions::new() + .with_filters([("byteField", "=", "10")]) + .unwrap() + .with_projection(projection) + .with_hudi_option(UseReadOptimizedMode.as_ref(), "true"), + ) + .await + .unwrap(); + let read_optimized_filtered_rows = id_name_active_rows(&read_optimized_filtered_batches); + + let ctx = SessionContext::new(); + let hudi = + HudiDataSource::new_with_options(base_url.as_str(), [(InputPartitions.as_ref(), "2")]) + .await + .unwrap(); + ctx.register_table("sample", Arc::new(hudi)).unwrap(); + + let actual_all_batches = ctx + .sql(r#"SELECT id, name, "isActive" FROM sample"#) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual_all_rows = id_name_active_rows(&actual_all_batches); + let actual_filtered_batches = ctx + .sql(r#"SELECT id, name, "isActive" FROM sample WHERE "byteField" = 10"#) + .await + .unwrap() + .collect() + .await + .unwrap(); + let actual_filtered_rows = id_name_active_rows(&actual_filtered_batches); + + assert_eq!(actual_all_rows, expected_all_rows); + assert_eq!(actual_all_rows.len(), 4); + assert_eq!(actual_filtered_rows, expected_filtered_rows); + assert_eq!( + actual_filtered_rows, + [ + (1, "Alice".to_string(), false), + (3, "Carol".to_string(), true) + ] + ); + assert_ne!( + actual_filtered_rows, read_optimized_filtered_rows, + "partition-filtered snapshot should merge the active MOR log in byteField=10" + ); + assert_eq!( + read_optimized_filtered_rows, + [ + (1, "Alice".to_string(), true), + (3, "Carol".to_string(), true) + ] + ); } #[tokio::test] -async fn test_v9_txns_mor_nonpart_meta() { - verify_v9_txns_table(&SampleTable::V9TxnsNonpartMeta, false, false).await; +async fn test_hudi_scan_exec_strips_filters_on_dropped_partition_columns() { + let table_path = dropped_partition_column_mor_table(); + let ctx = SessionContext::new(); + let hudi = HudiDataSource::new_with_options( + table_path.to_str().unwrap(), + [ + (InputPartitions.as_ref(), "2"), + (HudiInternalConfig::SkipConfigValidation.as_ref(), "true"), + ], + ) + .await + .unwrap(); + ctx.register_table("sample", Arc::new(hudi)).unwrap(); + + let sql = r#"SELECT id, name, "isActive" FROM sample WHERE "byteField" = 10 ORDER BY id"#; + let plan = explain_physical_plan(&ctx, sql).await; + assert!( + plan.contains("HudiScanExec"), + "dropped-partition MOR snapshot should use HudiScanExec. Plan: {plan}" + ); + + let batches = ctx.sql(sql).await.unwrap().collect().await.unwrap(); + let rows = id_name_active_rows(&batches); + + assert_eq!( + rows, + [ + (1, "Alice".to_string(), false), + (3, "Carol".to_string(), true) + ] + ); + + let err = match ctx + .sql(r#"SELECT "byteField" FROM sample WHERE "byteField" = 10"#) + .await + .unwrap() + .collect() + .await + { + Ok(_) => panic!("projecting a dropped partition column should fail"), + Err(err) => err, + }; + let err_message = err.to_string(); + assert!( + err_message.contains("byteField"), + "dropped partition projection error should name byteField, got: {err_message}" + ); } diff --git a/crates/test/Cargo.toml b/crates/test/Cargo.toml index d1e5caa5..4b9a98cf 100644 --- a/crates/test/Cargo.toml +++ b/crates/test/Cargo.toml @@ -17,6 +17,7 @@ [package] name = "hudi-test" +publish = false version.workspace = true edition.workspace = true license.workspace = true diff --git a/crates/test/src/lib.rs b/crates/test/src/lib.rs index f018b898..6d02ccc9 100644 --- a/crates/test/src/lib.rs +++ b/crates/test/src/lib.rs @@ -18,11 +18,15 @@ */ use arrow_array::{BooleanArray, Float64Array, Int32Array, RecordBatch, StringArray}; +use std::collections::hash_map::DefaultHasher; use std::fs; +use std::hash::{Hash, Hasher}; use std::io::Cursor; use std::path::{Path, PathBuf}; +use std::sync::{Mutex, OnceLock}; +use std::time::UNIX_EPOCH; use strum_macros::{AsRefStr, EnumIter, EnumString}; -use tempfile::tempdir; +use tempfile::{Builder as TempDirBuilder, tempdir}; use url::Url; use zip::ZipArchive; @@ -31,16 +35,123 @@ pub mod util; #[cfg(feature = "datafusion")] pub mod v9_verification; +static EXTRACT_LOCK: OnceLock> = OnceLock::new(); + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)] +pub enum TableFormat { + Cow, + MorParquet, + MorAvro, +} + +impl TableFormat { + fn table_type(self) -> &'static str { + match self { + Self::Cow => "cow", + Self::MorParquet | Self::MorAvro => "mor", + } + } + + fn log_format(self) -> Option<&'static str> { + match self { + Self::Cow => None, + Self::MorParquet => Some("parquet"), + Self::MorAvro => Some("avro"), + } + } +} + +const COW: &[TableFormat] = &[TableFormat::Cow]; +const MOR_AVRO: &[TableFormat] = &[TableFormat::MorAvro]; +const MOR_PARQUET: &[TableFormat] = &[TableFormat::MorParquet]; +const COW_AND_MOR_AVRO: &[TableFormat] = &[TableFormat::Cow, TableFormat::MorAvro]; +const COW_AND_MOR_PARQUET: &[TableFormat] = &[TableFormat::Cow, TableFormat::MorParquet]; + pub fn extract_test_table(zip_path: &Path) -> PathBuf { - let target_dir = tempdir().unwrap().path().to_path_buf(); - let archive = fs::read(zip_path).unwrap(); - ZipArchive::new(Cursor::new(archive)) - .unwrap() - .extract(&target_dir) - .unwrap(); + let _lock = EXTRACT_LOCK + .get_or_init(|| Mutex::new(())) + .lock() + .expect("fixture extraction lock should not be poisoned"); + + let target_dir = cached_extract_dir(zip_path); + if target_dir.exists() { + return target_dir; + } + + let cache_root = target_dir + .parent() + .expect("fixture cache path should have a parent"); + fs::create_dir_all(cache_root) + .unwrap_or_else(|e| panic!("create fixture cache {}: {e}", cache_root.display())); + let temp_dir = TempDirBuilder::new() + .prefix("extract-") + .tempdir_in(cache_root) + .unwrap_or_else(|e| panic!("create temp fixture dir in {}: {e}", cache_root.display())); + + extract_zip_to_dir(zip_path, temp_dir.path()); + + match fs::rename(temp_dir.path(), &target_dir) { + Ok(()) => target_dir, + Err(_) if target_dir.exists() => target_dir, + Err(e) => panic!( + "move extracted fixture {} to {}: {e}", + temp_dir.path().display(), + target_dir.display() + ), + } +} + +/// Extracts a fixture into a unique temp directory. +/// +/// Use this only for tests that mutate the extracted table in place. Normal +/// readers should use [`extract_test_table`], which reuses a bounded fixture +/// cache keyed by zip path and metadata. +pub fn extract_test_table_fresh(zip_path: &Path) -> PathBuf { + let temp_dir = tempdir().expect("create temp fixture dir"); + let target_dir = temp_dir.path().to_path_buf(); + extract_zip_to_dir(zip_path, &target_dir); + let kept_dir = temp_dir.keep(); + debug_assert_eq!(kept_dir, target_dir); target_dir } +fn cached_extract_dir(zip_path: &Path) -> PathBuf { + let zip_path = zip_path + .canonicalize() + .unwrap_or_else(|e| panic!("canonicalize fixture {}: {e}", zip_path.display())); + let metadata = fs::metadata(&zip_path) + .unwrap_or_else(|e| panic!("stat fixture {}: {e}", zip_path.display())); + let modified = metadata + .modified() + .ok() + .and_then(|modified| modified.duration_since(UNIX_EPOCH).ok()) + .map(|duration| duration.as_nanos()) + .unwrap_or_default(); + + let mut hasher = DefaultHasher::new(); + zip_path.hash(&mut hasher); + metadata.len().hash(&mut hasher); + modified.hash(&mut hasher); + + std::env::temp_dir() + .join("hudi-rs-test-fixtures") + .join(format!("{:016x}", hasher.finish())) +} + +fn extract_zip_to_dir(zip_path: &Path, target_dir: &Path) { + let archive = + fs::read(zip_path).unwrap_or_else(|e| panic!("read fixture {}: {e}", zip_path.display())); + let mut zip = ZipArchive::new(Cursor::new(archive)) + .unwrap_or_else(|e| panic!("open fixture zip {}: {e}", zip_path.display())); + zip.extract(target_dir).unwrap_or_else(|e| { + panic!( + "extract fixture {} to {}: {e}", + zip_path.display(), + target_dir.display() + ) + }); +} + #[allow(dead_code)] #[derive(Debug, EnumString, AsRefStr, EnumIter)] pub enum QuickstartTripsTable { @@ -98,21 +209,39 @@ impl QuickstartTripsTable { data_path.into_boxed_path() } - pub fn path_to_cow(&self) -> String { - let zip_path = self.zip_path("cow", None); + fn zip_path_for(&self, format: TableFormat) -> Box { + self.zip_path(format.table_type(), format.log_format()) + } + + pub fn available_formats(&self) -> &'static [TableFormat] { + match self { + Self::V6Trips8I1U | Self::V6Trips8I3D | Self::V8Trips8I3U1D => MOR_AVRO, + Self::V9TripsLance => COW_AND_MOR_AVRO, + } + } + + pub fn path(&self, format: TableFormat) -> String { + let zip_path = self.zip_path_for(format); let path_buf = extract_test_table(zip_path.as_ref()).join(self.as_ref()); path_buf.to_str().unwrap().to_string() } + pub fn path_to_cow(&self) -> String { + self.path(TableFormat::Cow) + } + pub fn url_to_cow(&self) -> Url { let path = self.path_to_cow(); Url::from_file_path(path).unwrap() } pub fn path_to_mor_avro(&self) -> String { - let zip_path = self.zip_path("mor", Some("avro")); - let path_buf = extract_test_table(zip_path.as_ref()).join(self.as_ref()); - path_buf.to_str().unwrap().to_string() + self.path(TableFormat::MorAvro) + } + + pub fn url(&self, format: TableFormat) -> Url { + let path = self.path(format); + Url::from_file_path(path).unwrap() } pub fn url_to_mor_avro(&self) -> Url { @@ -197,27 +326,86 @@ impl SampleTable { data_path.into_boxed_path() } - pub fn path_to_cow(&self) -> String { - let zip_path = self.zip_path("cow", None); + fn zip_path_for(&self, format: TableFormat) -> Box { + self.zip_path(format.table_type(), format.log_format()) + } + + pub fn available_formats(&self) -> &'static [TableFormat] { + match self { + Self::V6TimebasedkeygenNonhivestyle + | Self::V8ComplexkeygenHivestyle + | Self::V8Empty + | Self::V8Nonpartitioned + | Self::V8SimplekeygenHivestyleNoMetafields + | Self::V8SimplekeygenNonhivestyle + | Self::V9TimebasedkeygenEpochmillis + | Self::V9TimebasedkeygenUnixtimestamp + | Self::V9LanceNonpartitioned + | Self::V9LanceTxnsNonpart + | Self::V9LanceTxnsSimple => COW, + + Self::V6NonpartitionedRollback => MOR_PARQUET, + + Self::V9NonpartitionedRollback | Self::V9LanceNonhivestyle => MOR_AVRO, + + Self::V9TimebasedkeygenNonhivestyle + | Self::V9TxnsComplexMeta + | Self::V9TxnsComplexNometa + | Self::V9TxnsNonpartMeta + | Self::V9TxnsNonpartNometa + | Self::V9TxnsSimpleMeta + | Self::V9TxnsSimpleNometa + | Self::V9TxnsSimpleOverwrite => COW_AND_MOR_AVRO, + + Self::V6ComplexkeygenHivestyle + | Self::V6Empty + | Self::V6Nonpartitioned + | Self::V6SimplekeygenHivestyleNoMetafields + | Self::V6SimplekeygenNonhivestyle + | Self::V6SimplekeygenNonhivestyleOverwritetable => COW_AND_MOR_PARQUET, + } + } + + pub fn path(&self, format: TableFormat) -> String { + let zip_path = self.zip_path_for(format); let path_buf = extract_test_table(zip_path.as_ref()).join(self.as_ref()); path_buf.to_str().unwrap().to_string() } - pub fn path_to_mor_parquet(&self) -> String { - let zip_path = self.zip_path("mor", Some("parquet")); - let path_buf = extract_test_table(zip_path.as_ref()).join(self.as_ref()); + pub fn path_fresh(&self, format: TableFormat) -> String { + let zip_path = self.zip_path_for(format); + let path_buf = extract_test_table_fresh(zip_path.as_ref()).join(self.as_ref()); path_buf.to_str().unwrap().to_string() } + pub fn path_to_cow(&self) -> String { + self.path(TableFormat::Cow) + } + + pub fn path_to_cow_fresh(&self) -> String { + self.path_fresh(TableFormat::Cow) + } + + pub fn path_to_mor_parquet(&self) -> String { + self.path(TableFormat::MorParquet) + } + + pub fn path_to_mor_parquet_fresh(&self) -> String { + self.path_fresh(TableFormat::MorParquet) + } + pub fn url_to_cow(&self) -> Url { let path = self.path_to_cow(); Url::from_file_path(path).unwrap() } pub fn path_to_mor_avro(&self) -> String { - let zip_path = self.zip_path("mor", Some("avro")); - let path_buf = extract_test_table(zip_path.as_ref()).join(self.as_ref()); - path_buf.to_str().unwrap().to_string() + self.path(TableFormat::MorAvro) + } + + pub fn url(&self, format: TableFormat) -> Url { + let path = self.path(format); + Url::from_file_path(path).unwrap() } pub fn url_to_mor_parquet(&self) -> Url { @@ -231,7 +419,10 @@ impl SampleTable { } pub fn urls(&self) -> Vec { - vec![self.url_to_cow(), self.url_to_mor_parquet()] + self.available_formats() + .iter() + .map(|format| self.url(*format)) + .collect() } } @@ -255,72 +446,25 @@ mod tests { #[test] fn quickstart_trips_table_zip_file_should_exist() { for t in QuickstartTripsTable::iter() { - match t { - QuickstartTripsTable::V6Trips8I1U => { - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } - QuickstartTripsTable::V6Trips8I3D => { - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } - QuickstartTripsTable::V8Trips8I3U1D => { - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } - QuickstartTripsTable::V9TripsLance => { - let path = t.zip_path("cow", None); - assert!(path.exists()); - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } + for format in t.available_formats() { + let path = t.zip_path_for(*format); + assert!( + path.exists(), + "missing fixture {path:?} for {t:?} {format:?}" + ); } } } + #[test] fn sample_table_zip_file_should_exist() { for t in SampleTable::iter() { - match t { - SampleTable::V6TimebasedkeygenNonhivestyle => { - let path = t.zip_path("cow", None); - assert!(path.exists()); - } - SampleTable::V6NonpartitionedRollback => { - let path = t.zip_path("mor", Some("parquet")); - assert!(path.exists()); - } - SampleTable::V9NonpartitionedRollback => { - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } - SampleTable::V9TimebasedkeygenEpochmillis - | SampleTable::V9TimebasedkeygenUnixtimestamp - | SampleTable::V9LanceNonpartitioned - | SampleTable::V9LanceTxnsNonpart - | SampleTable::V9LanceTxnsSimple => { - let path = t.zip_path("cow", None); - assert!(path.exists()); - } - SampleTable::V9LanceNonhivestyle => { - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } - ref table if table.as_ref().starts_with("v9") => { - let path = t.zip_path("cow", None); - assert!(path.exists()); - let path = t.zip_path("mor", Some("avro")); - assert!(path.exists()); - } - ref table if table.as_ref().starts_with("v8") => { - let path = t.zip_path("cow", None); - assert!(path.exists()); - } - _ => { - let path = t.zip_path("cow", None); - assert!(path.exists()); - let path = t.zip_path("mor", Some("parquet")); - assert!(path.exists()); - } + for format in t.available_formats() { + let path = t.zip_path_for(*format); + assert!( + path.exists(), + "missing fixture {path:?} for {t:?} {format:?}" + ); } } } diff --git a/crates/test/src/util.rs b/crates/test/src/util.rs index 047aff54..56acaaa0 100644 --- a/crates/test/src/util.rs +++ b/crates/test/src/util.rs @@ -18,6 +18,8 @@ */ use arrow::record_batch::RecordBatch; use arrow_array::{Array, BooleanArray, Int32Array, StringArray}; +#[cfg(feature = "datafusion")] +use datafusion::prelude::SessionContext; use std::env; pub fn get_str_column<'a>(record_batch: &'a RecordBatch, name: &str) -> Vec<&'a str> { @@ -56,6 +58,17 @@ pub fn get_bool_column(record_batch: &RecordBatch, name: &str) -> Vec { .collect::>() } +#[cfg(feature = "datafusion")] +pub async fn explain_physical_plan(ctx: &SessionContext, sql: &str) -> String { + let explaining_df = ctx.sql(sql).await.unwrap().explain(false, true).unwrap(); + let explaining_rb = explaining_df.collect().await.unwrap(); + explaining_rb + .iter() + .flat_map(|batch| get_str_column(batch, "plan")) + .collect::>() + .join("\n") +} + /// Sets a fixed timezone by setting the TZ environment variable. pub fn set_fixed_timezone(tz: &str) { // SAFETY: Only used in serial tests diff --git a/crates/test/src/v9_verification.rs b/crates/test/src/v9_verification.rs index 1f1b4731..0ac99801 100644 --- a/crates/test/src/v9_verification.rs +++ b/crates/test/src/v9_verification.rs @@ -27,6 +27,7 @@ use datafusion::prelude::SessionContext; use hudi_datafusion::HudiDataSource; use crate::SampleTable; +use crate::util::explain_physical_plan; pub const EXPECTED_PARTITIONED_TXN_IDS: &[&str] = &[ "TXN-001", "TXN-003", "TXN-007", "TXN-008", "TXN-011", "TXN-012", "TXN-013", "TXN-014", @@ -144,14 +145,19 @@ pub async fn verify_nonpart_records(ctx: &SessionContext) { } /// Register a v9 txns table with a DataFusion session context. -async fn register_v9_table(ctx: &SessionContext, table: &SampleTable, cow: bool) { +async fn register_v9_table( + ctx: &SessionContext, + table: &SampleTable, + cow: bool, + mor_read_optimized: bool, +) { let url = if cow { table.url_to_cow() } else { table.url_to_mor_avro() }; let mut opts: Vec<(&str, &str)> = vec![]; - if !cow { + if !cow && mor_read_optimized { opts.push(("hoodie.read.use.read_optimized.mode", "true")); } let hudi = HudiDataSource::new_with_options(url.as_str(), opts) @@ -168,7 +174,36 @@ async fn register_v9_table(ctx: &SessionContext, table: &SampleTable, cow: bool) /// Entry point for v9 txns table verification, used by both Rust and Python tests. pub async fn verify_v9_txns_table(table: &SampleTable, cow: bool, partitioned: bool) { let ctx = SessionContext::new(); - register_v9_table(&ctx, table, cow).await; + register_v9_table(&ctx, table, cow, !cow).await; + if partitioned { + verify_partitioned_records(&ctx).await; + } else { + verify_nonpart_records(&ctx).await; + } +} + +/// Verify v9 txns MOR tables in snapshot mode. +/// +/// These fixtures currently have the same surviving rows in snapshot and +/// read-optimized mode: compaction and clustering materialize the update/delete +/// history into base files, and the post-clustering inserts are base-file +/// records rather than log-only records. +pub async fn verify_v9_txns_table_snapshot(table: &SampleTable, partitioned: bool) { + let ctx = SessionContext::new(); + register_v9_table(&ctx, table, false, false).await; + let plan = explain_physical_plan(&ctx, "SELECT txn_id FROM txns").await; + assert!( + plan.contains("HudiScanExec"), + "MOR snapshot should use HudiScanExec. Plan: {plan}" + ); + assert!( + !plan.contains("DataSourceExec"), + "MOR snapshot should not use DataSourceExec. Plan: {plan}" + ); + assert!( + !plan.contains("ParquetSource"), + "MOR snapshot should not use ParquetSource. Plan: {plan}" + ); if partitioned { verify_partitioned_records(&ctx).await; } else {