Support arbitrary user defined partition column in ListingTable (rather than assuming they are always Dictionary encoded) (#5545)

* refactor: user may choose to dict-encode partition values

Let the user decide if they may want to encode partition values for
file-based data sources. Dictionary encoding makes sense for string
values but is probably pointless or even counterproductive for integer
types.

* refactor: improve transition

* fix: remove leftover dbg

* refactor: `partition_{type,value}_wrap` -> `wrap_partition_{type,value}_in_dict`

---------

Co-authored-by: Andrew Lamb <[email protected]>

Author: crepererum

datafusion/core/src/datasource/listing/mod.rs

@@ -55,7 +55,16 @@ pub struct FileRange {
 pub struct PartitionedFile {
     /// Path for the file (e.g. URL, filesystem path, etc)
     pub object_meta: ObjectMeta,
-    /// Values of partition columns to be appended to each row
+    /// Values of partition columns to be appended to each row.
+    ///
+    /// These MUST have the same count, order, and type than the [`table_partition_cols`].
+    ///
+    /// You may use [`wrap_partition_value_in_dict`] to wrap them if you have used [`wrap_partition_type_in_dict`] to wrap the column type.
+    ///
+    ///
+    /// [`wrap_partition_type_in_dict`]: crate::physical_plan::file_format::wrap_partition_type_in_dict
+    /// [`wrap_partition_value_in_dict`]: crate::physical_plan::file_format::wrap_partition_value_in_dict
+    /// [`table_partition_cols`]: table::ListingOptions::table_partition_cols
     pub partition_values: Vec<ScalarValue>,
     /// An optional file range for a more fine-grained parallel execution
     pub range: Option<FileRange>,

datafusion/core/src/datasource/listing/table.rs

@@ -44,7 +44,6 @@ use crate::datasource::{
 };
 use crate::logical_expr::TableProviderFilterPushDown;
 use crate::physical_plan;
-use crate::physical_plan::file_format::partition_type_wrap;
 use crate::{
     error::{DataFusionError, Result},
     execution::context::SessionState,
@@ -300,6 +299,8 @@ impl ListingOptions {
 
     /// Set table partition column names on [`ListingOptions`] and returns self.
     ///
+    /// You may use [`wrap_partition_type_in_dict`] to request a dictionary-encoded type.
+    ///
     /// ```
     /// # use std::sync::Arc;
     /// # use arrow::datatypes::DataType;
@@ -315,6 +316,9 @@ impl ListingOptions {
     /// assert_eq!(listing_options.table_partition_cols, vec![("col_a".to_string(), DataType::Utf8),
     ///     ("col_b".to_string(), DataType::Utf8)]);
     /// ```
+    ///
+    ///
+    /// [`wrap_partition_type_in_dict`]: crate::physical_plan::file_format::wrap_partition_type_in_dict
     pub fn with_table_partition_cols(
         mut self,
         table_partition_cols: Vec<(String, DataType)>,
@@ -538,11 +542,7 @@ impl ListingTable {
         // Add the partition columns to the file schema
         let mut table_fields = file_schema.fields().clone();
         for (part_col_name, part_col_type) in &options.table_partition_cols {
-            table_fields.push(Field::new(
-                part_col_name,
-                partition_type_wrap(part_col_type.clone()),
-                false,
-            ));
+            table_fields.push(Field::new(part_col_name, part_col_type.clone(), false));
         }
         let infinite_source = options.infinite_source;
 
@@ -1012,10 +1012,7 @@ mod tests {
 
         let opt = ListingOptions::new(Arc::new(AvroFormat {}))
             .with_file_extension(FileType::AVRO.get_ext())
-            .with_table_partition_cols(vec![(
-                String::from("p1"),
-                partition_type_wrap(DataType::Utf8),
-            )])
+            .with_table_partition_cols(vec![(String::from("p1"), DataType::Utf8)])
             .with_target_partitions(4);
 
         let table_path = ListingTableUrl::parse("test:///table/").unwrap();

datafusion/core/src/datasource/listing_table_factory.rs

@@ -86,7 +86,15 @@ impl TableProviderFactory for ListingTableFactory {
                 None,
                 cmd.table_partition_cols
                     .iter()
-                    .map(|x| (x.clone(), DataType::Utf8))
+                    .map(|x| {
+                        (
+                            x.clone(),
+                            DataType::Dictionary(
+                                Box::new(DataType::UInt16),
+                                Box::new(DataType::Utf8),
+                            ),
+                        )
+                    })
                     .collect::<Vec<_>>(),
             )
         } else {

datafusion/core/src/physical_plan/file_format/avro.rs

@@ -213,7 +213,6 @@ mod tests {
     use crate::datasource::listing::PartitionedFile;
     use crate::datasource::object_store::ObjectStoreUrl;
     use crate::physical_plan::file_format::chunked_store::ChunkedStore;
-    use crate::physical_plan::file_format::partition_type_wrap;
     use crate::prelude::SessionContext;
     use crate::scalar::ScalarValue;
     use crate::test::object_store::local_unpartitioned_file;
@@ -409,10 +408,7 @@ mod tests {
             file_schema,
             statistics: Statistics::default(),
             limit: None,
-            table_partition_cols: vec![(
-                "date".to_owned(),
-                partition_type_wrap(DataType::Utf8),
-            )],
+            table_partition_cols: vec![("date".to_owned(), DataType::Utf8)],
             output_ordering: None,
             infinite_source: false,
         });

datafusion/core/src/physical_plan/file_format/csv.rs

@@ -326,7 +326,6 @@ mod tests {
     use super::*;
     use crate::datasource::file_format::file_type::FileType;
     use crate::physical_plan::file_format::chunked_store::ChunkedStore;
-    use crate::physical_plan::file_format::partition_type_wrap;
     use crate::prelude::*;
     use crate::test::{partitioned_csv_config, partitioned_file_groups};
     use crate::test_util::{aggr_test_schema_with_missing_col, arrow_test_data};
@@ -580,8 +579,7 @@ mod tests {
         let mut config = partitioned_csv_config(file_schema, file_groups)?;
 
         // Add partition columns
-        config.table_partition_cols =
-            vec![("date".to_owned(), partition_type_wrap(DataType::Utf8))];
+        config.table_partition_cols = vec![("date".to_owned(), DataType::Utf8)];
         config.file_groups[0][0].partition_values =
             vec![ScalarValue::Utf8(Some("2021-10-26".to_owned()))];