Skip to content

Commit 717d31a

Browse files
Rich-T-kidRich-T-kid
committed
trim test
1 parent 17480b8 commit 717d31a

2 files changed

Lines changed: 35 additions & 244 deletions

File tree

‎datafusion/core/tests/parquet/schema_coercion.rs‎

Lines changed: 32 additions & 185 deletions
Original file line numberDiff line numberDiff line change
@@ -31,14 +31,8 @@ use datafusion_common::ScalarValue;
3131
use datafusion_common::test_util::batches_to_sort_string;
3232
use datafusion_execution::object_store::ObjectStoreUrl;
3333

34-
use datafusion_common::config::TableParquetOptions;
3534
use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
3635
use datafusion_datasource::source::DataSourceExec;
37-
use datafusion_physical_expr::expressions::{Column, IsNotNullExpr};
38-
use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
39-
use datafusion_physical_plan::ExecutionPlan;
40-
use datafusion_physical_plan::filter::FilterExec;
41-
use datafusion_physical_plan::projection::ProjectionExec;
4236
use insta::assert_snapshot;
4337
use object_store::ObjectMeta;
4438
use parquet::arrow::ArrowWriter;
@@ -152,9 +146,6 @@ async fn multi_parquet_coercion_projection() {
152146
");
153147
}
154148

155-
/// Writes a parquet file with mixed column encoding: `rle_cols` use RLE_DICTIONARY,
156-
/// all other string columns use plain encoding. Returns the file so the caller
157-
/// can keep it alive.
158149
fn store_mixed_encoding_parquet(
159150
batch: &RecordBatch,
160151
rle_cols: &[&str],
@@ -178,9 +169,9 @@ fn store_mixed_encoding_parquet(
178169
(meta, output)
179170
}
180171

181-
/// Logical planning: after `register_parquet` with the flag on, the table schema
182-
/// already has RLE-encoded columns as `Dictionary(Int32, Utf8)` before any query
183-
/// executes. Plain string columns remain `Utf8View`.
172+
/// Promotion happens at logical planning time: after `register_parquet` with the
173+
/// flag on, `ctx.table()` already shows RLE columns as `Dictionary(Int32, Utf8)`
174+
/// before any query executes.
184175
#[tokio::test]
185176
async fn rle_dictionary_logical_schema_promotion() {
186177
let schema = Arc::new(Schema::new(vec![
@@ -195,7 +186,6 @@ async fn rle_dictionary_logical_schema_promotion() {
195186
],
196187
)
197188
.unwrap();
198-
// env → RLE_DICTIONARY, version → plain
199189
let (_meta, file) = store_mixed_encoding_parquet(&batch, &["env"]);
200190

201191
let config = SessionConfig::new().set(
@@ -211,9 +201,7 @@ async fn rle_dictionary_logical_schema_promotion() {
211201
.await
212202
.unwrap();
213203

214-
// Inspect the logical schema — no query execution.
215-
let df = ctx.table("t").await.unwrap();
216-
let logical_schema = df.schema();
204+
let logical_schema = ctx.table("t").await.unwrap().schema().clone();
217205
let dict_type =
218206
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
219207
assert_eq!(
@@ -233,10 +221,10 @@ async fn rle_dictionary_logical_schema_promotion() {
233221
}
234222

235223
/// Five string columns, three RLE_DICTIONARY encoded and two plain.
236-
/// With the flag on DataFusion should surface only the three RLE columns as
237-
/// Dictionary(Int32, Utf8) and leave the plain columns as Utf8View.
224+
/// Flag on: only the three RLE columns become `Dictionary(Int32, Utf8)`; plain
225+
/// columns stay `Utf8View`. Flag off: all columns stay `Utf8View`.
238226
#[tokio::test]
239-
async fn rle_dictionary_selective_promotion_only_rle_columns_become_dict() {
227+
async fn rle_dictionary_selective_promotion() {
240228
let schema = Arc::new(Schema::new(vec![
241229
Field::new("env", DataType::Utf8, true),
242230
Field::new("region", DataType::Utf8, true),
@@ -262,202 +250,61 @@ async fn rle_dictionary_selective_promotion_only_rle_columns_become_dict() {
262250
.unwrap();
263251

264252
let (_meta, file) = store_mixed_encoding_parquet(&batch, &["env", "region", "tier"]);
265-
266-
let config = SessionConfig::new().set(
267-
"datafusion.execution.parquet.enable_rle_to_dictionary",
268-
&ScalarValue::Boolean(Some(true)),
269-
);
270-
let ctx = SessionContext::new_with_config(config);
271-
ctx.register_parquet(
272-
"t",
273-
file.path().to_str().unwrap(),
274-
ParquetReadOptions::default(),
275-
)
276-
.await
277-
.unwrap();
278-
279-
let result = ctx
280-
.sql("SELECT env, region, tier, version, build FROM t LIMIT 1")
281-
.await
282-
.unwrap()
283-
.collect()
284-
.await
285-
.unwrap();
286-
287-
assert!(!result.is_empty());
288-
let out_schema = result[0].schema();
289-
let dict_type =
290-
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
291-
assert_eq!(
292-
out_schema.field_with_name("env").unwrap().data_type(),
293-
&dict_type
294-
);
295-
assert_eq!(
296-
out_schema.field_with_name("region").unwrap().data_type(),
297-
&dict_type
298-
);
299-
assert_eq!(
300-
out_schema.field_with_name("tier").unwrap().data_type(),
301-
&dict_type
302-
);
303-
assert_eq!(
304-
out_schema.field_with_name("version").unwrap().data_type(),
305-
&DataType::Utf8View
306-
);
307-
assert_eq!(
308-
out_schema.field_with_name("build").unwrap().data_type(),
309-
&DataType::Utf8View
310-
);
311-
}
312-
313-
fn store_rle_dict_parquet(batch: &RecordBatch) -> (ObjectMeta, NamedTempFile) {
314-
let mut output = tempfile::Builder::new()
315-
.suffix(".parquet")
316-
.tempfile()
317-
.expect("creating temp file");
318-
let props = WriterProperties::builder()
319-
.set_dictionary_enabled(true)
320-
.build();
321-
let mut writer = ArrowWriter::try_new(&mut output, batch.schema(), Some(props))
322-
.expect("creating writer");
323-
writer.write(batch).expect("Writing batch");
324-
writer.close().unwrap();
325-
let meta = local_unpartitioned_file(&output);
326-
(meta, output)
327-
}
328-
329-
/// Runs scan → filter (IS NOT NULL on `env`) → project through a physical plan
330-
/// and returns the collected batches. Exercises planning-time schema consistency:
331-
/// FilterExec and ProjectionExec reference columns by the type the scan declares,
332-
/// so a mismatch between the plan's declared type and the arrays the scan produces
333-
/// will surface here.
334-
async fn scan_filter_project_env(
335-
meta: ObjectMeta,
336-
source: Arc<ParquetSource>,
337-
) -> Vec<RecordBatch> {
338-
let task_ctx = SessionContext::new().task_ctx();
339-
let conf = FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source)
340-
.with_file_group(vec![meta.into()].into())
341-
.build();
342-
let scan: Arc<dyn ExecutionPlan> = DataSourceExec::from_data_source(conf);
343-
344-
let env_col = Arc::new(Column::new("env", 0)) as Arc<dyn PhysicalExpr>;
345-
let not_null = Arc::new(IsNotNullExpr::new(Arc::clone(&env_col)));
346-
let filter: Arc<dyn ExecutionPlan> =
347-
Arc::new(FilterExec::try_new(not_null, Arc::clone(&scan)).unwrap());
348-
349-
let projection_exprs = vec![(env_col, "env".to_string())];
350-
let projection = Arc::new(ProjectionExec::try_new(projection_exprs, filter).unwrap())
351-
as Arc<dyn ExecutionPlan>;
352-
353-
collect(projection, task_ctx).await.unwrap()
354-
}
355-
356-
fn env_batch() -> RecordBatch {
357-
let schema = Arc::new(Schema::new(vec![Field::new("env", DataType::Utf8, true)]));
358-
let envs = Arc::new(StringArray::from(vec![
359-
Some("prod"),
360-
Some("staging"),
361-
Some("prod"),
362-
Some("prod"),
363-
Some("canary"),
364-
Some("staging"),
365-
Some("prod"),
366-
Some("canary"),
367-
]));
368-
RecordBatch::try_new(schema, vec![envs as _]).unwrap()
369-
}
370-
371-
/// Direct ParquetSource (no infer_schema): flag off + Utf8 schema stays Utf8;
372-
/// flag on + pre-promoted Dict schema (as infer_schema would produce) gives Dict.
373-
/// FilterExec + ProjectionExec catch any schema mismatch at runtime.
374-
#[tokio::test]
375-
async fn rle_dictionary_direct_source() {
376-
let batch = env_batch();
377-
let utf8_schema = batch.schema();
378-
let (meta, _file) = store_rle_dict_parquet(&batch);
379-
let dict_type =
380-
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
381-
382-
let read = scan_filter_project_env(
383-
meta.clone(),
384-
Arc::new(ParquetSource::new(Arc::clone(&utf8_schema))),
385-
)
386-
.await;
387-
assert_eq!(read[0].schema().field(0).data_type(), &DataType::Utf8);
388-
389-
let dict_schema = Arc::new(Schema::new(vec![Field::new(
390-
"env",
391-
dict_type.clone(),
392-
true,
393-
)]));
394-
let mut options = TableParquetOptions::default();
395-
options.global.enable_rle_to_dictionary = true;
396-
let source =
397-
Arc::new(ParquetSource::new(dict_schema).with_table_parquet_options(options));
398-
let read = scan_filter_project_env(meta, source).await;
399-
assert_eq!(read[0].schema().field(0).data_type(), &dict_type);
400-
}
401-
402-
/// E2E scan → filter → aggregate: flag off gives Utf8View, flag on gives Dict.
403-
#[tokio::test]
404-
async fn rle_dictionary_e2e_aggregate() {
405-
let (_meta, file) = store_rle_dict_parquet(&env_batch());
406253
let path = file.path().to_str().unwrap();
407-
let sql = "SELECT env, COUNT(*) AS cnt FROM envs WHERE env IS NOT NULL GROUP BY env";
254+
let sql = "SELECT env, region, tier, version, build FROM t LIMIT 1";
408255
let dict_type =
409256
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
410257

258+
// Flag off: all string columns stay Utf8View.
411259
let ctx = SessionContext::new();
412-
ctx.register_parquet("envs", path, ParquetReadOptions::default())
260+
ctx.register_parquet("t", path, ParquetReadOptions::default())
413261
.await
414262
.unwrap();
415263
let result = ctx.sql(sql).await.unwrap().collect().await.unwrap();
416-
assert_eq!(
417-
result[0]
418-
.schema()
419-
.field_with_name("env")
420-
.unwrap()
421-
.data_type(),
422-
&DataType::Utf8View,
423-
);
264+
assert!(!result.is_empty());
265+
let s = result[0].schema();
266+
for col in ["env", "region", "tier", "version", "build"] {
267+
assert_eq!(
268+
s.field_with_name(col).unwrap().data_type(),
269+
&DataType::Utf8View
270+
);
271+
}
424272

273+
// Flag on: only RLE-encoded columns become Dict; plain columns stay Utf8View.
425274
let config = SessionConfig::new().set(
426275
"datafusion.execution.parquet.enable_rle_to_dictionary",
427276
&ScalarValue::Boolean(Some(true)),
428277
);
429278
let ctx = SessionContext::new_with_config(config);
430-
ctx.register_parquet("envs", path, ParquetReadOptions::default())
279+
ctx.register_parquet("t", path, ParquetReadOptions::default())
431280
.await
432281
.unwrap();
433282
let result = ctx.sql(sql).await.unwrap().collect().await.unwrap();
434-
assert_eq!(
435-
result[0]
436-
.schema()
437-
.field_with_name("env")
438-
.unwrap()
439-
.data_type(),
440-
&dict_type,
441-
);
283+
assert!(!result.is_empty());
284+
let s = result[0].schema();
285+
for col in ["env", "region", "tier"] {
286+
assert_eq!(s.field_with_name(col).unwrap().data_type(), &dict_type);
287+
}
288+
for col in ["version", "build"] {
289+
assert_eq!(
290+
s.field_with_name(col).unwrap().data_type(),
291+
&DataType::Utf8View
292+
);
293+
}
442294
}
443295

444296
/// Writes `batches` to a temporary parquet file
445297
pub fn store_parquet(
446298
batches: Vec<RecordBatch>,
447299
) -> Result<(Vec<ObjectMeta>, Vec<NamedTempFile>)> {
448-
// Each batch writes to their own file
449300
let files: Vec<_> = batches
450301
.into_iter()
451302
.map(|batch| {
452303
let mut output = NamedTempFile::new().expect("creating temp file");
453-
454-
let builder = WriterProperties::builder();
455-
let props = builder.build();
456-
304+
let props = WriterProperties::builder().build();
457305
let mut writer =
458306
ArrowWriter::try_new(&mut output, batch.schema(), Some(props))
459307
.expect("creating writer");
460-
461308
writer.write(&batch).expect("Writing batch");
462309
writer.close().unwrap();
463310
output

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

Lines changed: 3 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -3916,28 +3916,6 @@ mod test {
39163916
}
39173917
}
39183918

3919-
async fn write_rle_string_parquet(
3920-
store: Arc<dyn ObjectStore>,
3921-
filename: &str,
3922-
) -> usize {
3923-
let props = WriterProperties::builder()
3924-
.set_dictionary_enabled(true)
3925-
.build();
3926-
let batch = record_batch!((
3927-
"category",
3928-
Utf8,
3929-
vec![
3930-
Some("apple"),
3931-
Some("banana"),
3932-
Some("apple"),
3933-
Some("cherry"),
3934-
Some("banana")
3935-
]
3936-
))
3937-
.unwrap();
3938-
write_parquet_batches(store, filename, vec![batch], Some(props)).await
3939-
}
3940-
39413919
async fn collect_batches(
39423920
morselizer: &ParquetMorselizer,
39433921
file: PartitionedFile,
@@ -3950,45 +3928,11 @@ mod test {
39503928
batches
39513929
}
39523930

3953-
// Covers three opener-level promotion scenarios:
3954-
// 1. flag off: RLE string column → Utf8 (no promotion)
3955-
// 2. flag on + Dict schema (pre-promoted by infer_schema): string → Dict(Int32, Utf8)
3956-
// 3. flag on + Dict schema: binary → Dict(Int32, Binary)
3931+
// Binary columns with RLE_DICTIONARY encoding are promoted to Dict(Int32, Binary)
3932+
// when the logical schema declares them as Dictionary and the flag is on.
39573933
#[tokio::test]
3958-
async fn test_rle_column_promotion() {
3934+
async fn test_rle_binary_column_promotion() {
39593935
let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
3960-
let file_size = write_rle_string_parquet(Arc::clone(&store), "rle.parquet").await;
3961-
let file = || PartitionedFile::new("rle.parquet".to_string(), file_size as u64);
3962-
3963-
let plain_schema = Arc::new(Schema::new(vec![Field::new(
3964-
"category",
3965-
DataType::Utf8,
3966-
true,
3967-
)]));
3968-
let morselizer = ParquetMorselizerBuilder::new()
3969-
.with_store(Arc::clone(&store))
3970-
.with_schema(Arc::clone(&plain_schema))
3971-
.with_enable_rle_to_dictionary(false)
3972-
.build();
3973-
let batches = collect_batches(&morselizer, file()).await;
3974-
assert_eq!(batches[0].schema().field(0).data_type(), &DataType::Utf8);
3975-
3976-
let dict_str_schema = Arc::new(Schema::new(vec![Field::new(
3977-
"category",
3978-
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3979-
true,
3980-
)]));
3981-
let morselizer = ParquetMorselizerBuilder::new()
3982-
.with_store(Arc::clone(&store))
3983-
.with_schema(Arc::clone(&dict_str_schema))
3984-
.with_enable_rle_to_dictionary(true)
3985-
.build();
3986-
let batches = collect_batches(&morselizer, file()).await;
3987-
assert_eq!(
3988-
batches[0].schema().field(0).data_type(),
3989-
&DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8))
3990-
);
3991-
39923936
let bin_schema = Arc::new(Schema::new(vec![Field::new(
39933937
"payload",
39943938
DataType::Binary,

0 commit comments

Comments
 (0)