trim test

Author: Rich-T-kid

datafusion/core/tests/parquet/schema_coercion.rs

@@ -31,14 +31,8 @@ use datafusion_common::ScalarValue;
 use datafusion_common::test_util::batches_to_sort_string;
 use datafusion_execution::object_store::ObjectStoreUrl;
 
-use datafusion_common::config::TableParquetOptions;
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_datasource::source::DataSourceExec;
-use datafusion_physical_expr::expressions::{Column, IsNotNullExpr};
-use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
-use datafusion_physical_plan::ExecutionPlan;
-use datafusion_physical_plan::filter::FilterExec;
-use datafusion_physical_plan::projection::ProjectionExec;
 use insta::assert_snapshot;
 use object_store::ObjectMeta;
 use parquet::arrow::ArrowWriter;
@@ -152,9 +146,7 @@ async fn multi_parquet_coercion_projection() {
     ");
 }
 
-/// Writes a parquet file with mixed column encoding: `rle_cols` use RLE_DICTIONARY,
-/// all other string columns use plain encoding.  Returns the file so the caller
-/// can keep it alive.
+/// Writes `batch` to a temp parquet file where `rle_cols` use RLE_DICTIONARY encoding and all other columns use plain.
 fn store_mixed_encoding_parquet(
     batch: &RecordBatch,
     rle_cols: &[&str],
@@ -178,9 +170,9 @@ fn store_mixed_encoding_parquet(
     (meta, output)
 }
 
-/// Logical planning: after `register_parquet` with the flag on, the table schema
-/// already has RLE-encoded columns as `Dictionary(Int32, Utf8)` before any query
-/// executes.  Plain string columns remain `Utf8View`.
+/// Promotion happens at logical planning time: after `register_parquet` with the
+/// flag on, `ctx.table()` already shows RLE columns as `Dictionary(Int32, Utf8)`
+/// before any query executes.
 #[tokio::test]
 async fn rle_dictionary_logical_schema_promotion() {
     let schema = Arc::new(Schema::new(vec![
@@ -195,7 +187,6 @@ async fn rle_dictionary_logical_schema_promotion() {
         ],
     )
     .unwrap();
-    // env → RLE_DICTIONARY, version → plain
     let (_meta, file) = store_mixed_encoding_parquet(&batch, &["env"]);
 
     let config = SessionConfig::new().set(
@@ -211,9 +202,7 @@ async fn rle_dictionary_logical_schema_promotion() {
     .await
     .unwrap();
 
-    // Inspect the logical schema — no query execution.
-    let df = ctx.table("t").await.unwrap();
-    let logical_schema = df.schema();
+    let logical_schema = ctx.table("t").await.unwrap().schema().clone();
     let dict_type =
         DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
     assert_eq!(
@@ -233,10 +222,10 @@ async fn rle_dictionary_logical_schema_promotion() {
 }
 
 /// Five string columns, three RLE_DICTIONARY encoded and two plain.
-/// With the flag on DataFusion should surface only the three RLE columns as
-/// Dictionary(Int32, Utf8) and leave the plain columns as Utf8View.
+/// Flag on: only the three RLE columns become `Dictionary(Int32, Utf8)`; plain
+/// columns stay `Utf8View`.  Flag off: all columns stay `Utf8View`.
 #[tokio::test]
-async fn rle_dictionary_selective_promotion_only_rle_columns_become_dict() {
+async fn rle_dictionary_selective_promotion() {
     let schema = Arc::new(Schema::new(vec![
         Field::new("env", DataType::Utf8, true),
         Field::new("region", DataType::Utf8, true),
@@ -262,183 +251,47 @@ async fn rle_dictionary_selective_promotion_only_rle_columns_become_dict() {
     .unwrap();
 
     let (_meta, file) = store_mixed_encoding_parquet(&batch, &["env", "region", "tier"]);
-
-    let config = SessionConfig::new().set(
-        "datafusion.execution.parquet.enable_rle_to_dictionary",
-        &ScalarValue::Boolean(Some(true)),
-    );
-    let ctx = SessionContext::new_with_config(config);
-    ctx.register_parquet(
-        "t",
-        file.path().to_str().unwrap(),
-        ParquetReadOptions::default(),
-    )
-    .await
-    .unwrap();
-
-    let result = ctx
-        .sql("SELECT env, region, tier, version, build FROM t LIMIT 1")
-        .await
-        .unwrap()
-        .collect()
-        .await
-        .unwrap();
-
-    assert!(!result.is_empty());
-    let out_schema = result[0].schema();
-    let dict_type =
-        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
-    assert_eq!(
-        out_schema.field_with_name("env").unwrap().data_type(),
-        &dict_type
-    );
-    assert_eq!(
-        out_schema.field_with_name("region").unwrap().data_type(),
-        &dict_type
-    );
-    assert_eq!(
-        out_schema.field_with_name("tier").unwrap().data_type(),
-        &dict_type
-    );
-    assert_eq!(
-        out_schema.field_with_name("version").unwrap().data_type(),
-        &DataType::Utf8View
-    );
-    assert_eq!(
-        out_schema.field_with_name("build").unwrap().data_type(),
-        &DataType::Utf8View
-    );
-}
-
-fn store_rle_dict_parquet(batch: &RecordBatch) -> (ObjectMeta, NamedTempFile) {
-    let mut output = tempfile::Builder::new()
-        .suffix(".parquet")
-        .tempfile()
-        .expect("creating temp file");
-    let props = WriterProperties::builder()
-        .set_dictionary_enabled(true)
-        .build();
-    let mut writer = ArrowWriter::try_new(&mut output, batch.schema(), Some(props))
-        .expect("creating writer");
-    writer.write(batch).expect("Writing batch");
-    writer.close().unwrap();
-    let meta = local_unpartitioned_file(&output);
-    (meta, output)
-}
-
-/// Runs scan → filter (IS NOT NULL on `env`) → project through a physical plan
-/// and returns the collected batches. Exercises planning-time schema consistency:
-/// FilterExec and ProjectionExec reference columns by the type the scan declares,
-/// so a mismatch between the plan's declared type and the arrays the scan produces
-/// will surface here.
-async fn scan_filter_project_env(
-    meta: ObjectMeta,
-    source: Arc<ParquetSource>,
-) -> Vec<RecordBatch> {
-    let task_ctx = SessionContext::new().task_ctx();
-    let conf = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
-        .with_file_group(vec![meta.into()].into())
-        .build();
-    let scan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(conf);
-
-    let env_col = Arc::new(Column::new("env", 0)) as Arc<dyn PhysicalExpr>;
-    let not_null = Arc::new(IsNotNullExpr::new(Arc::clone(&env_col)));
-    let filter: Arc<dyn ExecutionPlan> =
-        Arc::new(FilterExec::try_new(not_null, Arc::clone(&scan)).unwrap());
-
-    let projection_exprs = vec![(env_col, "env".to_string())];
-    let projection = Arc::new(ProjectionExec::try_new(projection_exprs, filter).unwrap())
-        as Arc<dyn ExecutionPlan>;
-
-    collect(projection, task_ctx).await.unwrap()
-}
-
-fn env_batch() -> RecordBatch {
-    let schema = Arc::new(Schema::new(vec![Field::new("env", DataType::Utf8, true)]));
-    let envs = Arc::new(StringArray::from(vec![
-        Some("prod"),
-        Some("staging"),
-        Some("prod"),
-        Some("prod"),
-        Some("canary"),
-        Some("staging"),
-        Some("prod"),
-        Some("canary"),
-    ]));
-    RecordBatch::try_new(schema, vec![envs as _]).unwrap()
-}
-
-/// Direct ParquetSource (no infer_schema): flag off + Utf8 schema stays Utf8;
-/// flag on + pre-promoted Dict schema (as infer_schema would produce) gives Dict.
-/// FilterExec + ProjectionExec catch any schema mismatch at runtime.
-#[tokio::test]
-async fn rle_dictionary_direct_source() {
-    let batch = env_batch();
-    let utf8_schema = batch.schema();
-    let (meta, _file) = store_rle_dict_parquet(&batch);
-    let dict_type =
-        DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
-
-    let read = scan_filter_project_env(
-        meta.clone(),
-        Arc::new(ParquetSource::new(Arc::clone(&utf8_schema))),
-    )
-    .await;
-    assert_eq!(read[0].schema().field(0).data_type(), &DataType::Utf8);
-
-    let dict_schema = Arc::new(Schema::new(vec![Field::new(
-        "env",
-        dict_type.clone(),
-        true,
-    )]));
-    let mut options = TableParquetOptions::default();
-    options.global.enable_rle_to_dictionary = true;
-    let source =
-        Arc::new(ParquetSource::new(dict_schema).with_table_parquet_options(options));
-    let read = scan_filter_project_env(meta, source).await;
-    assert_eq!(read[0].schema().field(0).data_type(), &dict_type);
-}
-
-/// E2E scan → filter → aggregate: flag off gives Utf8View, flag on gives Dict.
-#[tokio::test]
-async fn rle_dictionary_e2e_aggregate() {
-    let (_meta, file) = store_rle_dict_parquet(&env_batch());
     let path = file.path().to_str().unwrap();
-    let sql = "SELECT env, COUNT(*) AS cnt FROM envs WHERE env IS NOT NULL GROUP BY env";
+    let sql = "SELECT env, region, tier, version, build FROM t LIMIT 1";
     let dict_type =
         DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
 
+    // Flag off: all string columns stay Utf8View.
     let ctx = SessionContext::new();
-    ctx.register_parquet("envs", path, ParquetReadOptions::default())
+    ctx.register_parquet("t", path, ParquetReadOptions::default())
         .await
         .unwrap();
     let result = ctx.sql(sql).await.unwrap().collect().await.unwrap();
-    assert_eq!(
-        result[0]
-            .schema()
-            .field_with_name("env")
-            .unwrap()
-            .data_type(),
-        &DataType::Utf8View,
-    );
+    assert!(!result.is_empty());
+    let s = result[0].schema();
+    for col in ["env", "region", "tier", "version", "build"] {
+        assert_eq!(
+            s.field_with_name(col).unwrap().data_type(),
+            &DataType::Utf8View
+        );
+    }
 
+    // Flag on: only RLE-encoded columns become Dict; plain columns stay Utf8View.
     let config = SessionConfig::new().set(
         "datafusion.execution.parquet.enable_rle_to_dictionary",
         &ScalarValue::Boolean(Some(true)),
     );
     let ctx = SessionContext::new_with_config(config);
-    ctx.register_parquet("envs", path, ParquetReadOptions::default())
+    ctx.register_parquet("t", path, ParquetReadOptions::default())
         .await
         .unwrap();
     let result = ctx.sql(sql).await.unwrap().collect().await.unwrap();
-    assert_eq!(
-        result[0]
-            .schema()
-            .field_with_name("env")
-            .unwrap()
-            .data_type(),
-        &dict_type,
-    );
+    assert!(!result.is_empty());
+    let s = result[0].schema();
+    for col in ["env", "region", "tier"] {
+        assert_eq!(s.field_with_name(col).unwrap().data_type(), &dict_type);
+    }
+    for col in ["version", "build"] {
+        assert_eq!(
+            s.field_with_name(col).unwrap().data_type(),
+            &DataType::Utf8View
+        );
+    }
 }
 
 /// Writes `batches` to a temporary parquet file
@@ -450,14 +303,10 @@ pub fn store_parquet(
         .into_iter()
         .map(|batch| {
             let mut output = NamedTempFile::new().expect("creating temp file");
-
-            let builder = WriterProperties::builder();
-            let props = builder.build();
-
+            let props = WriterProperties::builder().build();
             let mut writer =
                 ArrowWriter::try_new(&mut output, batch.schema(), Some(props))
                     .expect("creating writer");
-
             writer.write(&batch).expect("Writing batch");
             writer.close().unwrap();
             output

datafusion/datasource-parquet/src/opener/mod.rs

@@ -3916,28 +3916,6 @@ mod test {
         }
     }
 
-    async fn write_rle_string_parquet(
-        store: Arc<dyn ObjectStore>,
-        filename: &str,
-    ) -> usize {
-        let props = WriterProperties::builder()
-            .set_dictionary_enabled(true)
-            .build();
-        let batch = record_batch!((
-            "category",
-            Utf8,
-            vec![
-                Some("apple"),
-                Some("banana"),
-                Some("apple"),
-                Some("cherry"),
-                Some("banana")
-            ]
-        ))
-        .unwrap();
-        write_parquet_batches(store, filename, vec![batch], Some(props)).await
-    }
-
     async fn collect_batches(
         morselizer: &ParquetMorselizer,
         file: PartitionedFile,
@@ -3950,45 +3928,11 @@ mod test {
         batches
     }
 
-    // Covers three opener-level promotion scenarios:
-    //   1. flag off: RLE string column → Utf8 (no promotion)
-    //   2. flag on + Dict schema (pre-promoted by infer_schema): string → Dict(Int32, Utf8)
-    //   3. flag on + Dict schema: binary → Dict(Int32, Binary)
+    // Binary columns with RLE_DICTIONARY encoding are promoted to Dict(Int32, Binary)
+    // when the logical schema declares them as Dictionary and the flag is on.
     #[tokio::test]
-    async fn test_rle_column_promotion() {
+    async fn test_rle_binary_column_promotion() {
         let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
-        let file_size = write_rle_string_parquet(Arc::clone(&store), "rle.parquet").await;
-        let file = || PartitionedFile::new("rle.parquet".to_string(), file_size as u64);
-
-        let plain_schema = Arc::new(Schema::new(vec![Field::new(
-            "category",
-            DataType::Utf8,
-            true,
-        )]));
-        let morselizer = ParquetMorselizerBuilder::new()
-            .with_store(Arc::clone(&store))
-            .with_schema(Arc::clone(&plain_schema))
-            .with_enable_rle_to_dictionary(false)
-            .build();
-        let batches = collect_batches(&morselizer, file()).await;
-        assert_eq!(batches[0].schema().field(0).data_type(), &DataType::Utf8);
-
-        let dict_str_schema = Arc::new(Schema::new(vec![Field::new(
-            "category",
-            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
-            true,
-        )]));
-        let morselizer = ParquetMorselizerBuilder::new()
-            .with_store(Arc::clone(&store))
-            .with_schema(Arc::clone(&dict_str_schema))
-            .with_enable_rle_to_dictionary(true)
-            .build();
-        let batches = collect_batches(&morselizer, file()).await;
-        assert_eq!(
-            batches[0].schema().field(0).data_type(),
-            &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
-        );
-
         let bin_schema = Arc::new(Schema::new(vec![Field::new(
             "payload",
             DataType::Binary,

datafusion/datasource-parquet/src/schema_coercion.rs

@@ -126,8 +126,6 @@ pub fn apply_file_schema_type_coercions(
                     (&DataType::BinaryView, DataType::Binary | DataType::LargeBinary) => {
                         return field_with_new_type(field, DataType::BinaryView);
                     }
-                    // When the logical schema has a Dictionary type (RLE promotion),
-                    // tell the parquet reader to decode this column as dictionary arrays.
                     (
                         DataType::Dictionary(_, _),
                         DataType::Utf8