Merge pull request #17 from datafusion-contrib/feature/partition-cols

Feature: partition cols
datafusion-contrib · Apr 23, 2024 · 099c4b8 · 099c4b8
2 parents 85e4481 + b53d2ea
commit 099c4b8
Showing 1 changed file with 17 additions and 4 deletions.
diff --git a/src/catalog_provider/glue.rs b/src/catalog_provider/glue.rs
@@ -185,7 +185,8 @@ impl GlueCatalogProvider {
         let sd = Self::get_storage_descriptor(glue_table)?;
         let storage_location_uri = Self::get_storage_location(&sd)?;
 
-        let listing_options = Self::get_listing_options(database_name, table_name, &sd)?;
+        let listing_options =
+            Self::get_listing_options(database_name, table_name, &sd, glue_table)?;
 
         let schema_provider_for_database = self.ensure_schema_provider_for_database(database_name);
 
@@ -213,8 +214,9 @@ impl GlueCatalogProvider {
         database_name: &str,
         table_name: &str,
         sd: &StorageDescriptor,
+        glue_table: &Table,
     ) -> Result<ListingOptions> {
-        Self::calculate_options(sd)
+        Self::calculate_options(glue_table, sd)
             .map_err(|e| Self::wrap_error_with_table_info(database_name, table_name, e))
     }
 
@@ -280,7 +282,7 @@ impl GlueCatalogProvider {
         }
     }
 
-    fn calculate_options(sd: &StorageDescriptor) -> Result<ListingOptions> {
+    fn calculate_options(glue_table: &Table, sd: &StorageDescriptor) -> Result<ListingOptions> {
         let empty_str = String::from("");
         let input_format = sd.input_format.as_ref().unwrap_or(&empty_str);
         let output_format = sd.output_format.as_ref().unwrap_or(&empty_str);
@@ -365,10 +367,21 @@ impl GlueCatalogProvider {
         };
         let format = format_result?;
 
+        let partition_cols = glue_table
+            .partition_keys()
+            .iter()
+            .map(|c| {
+                (
+                    c.name().to_owned(),
+                    Self::map_glue_data_type(&c.r#type.clone().unwrap()).unwrap(),
+                )
+            })
+            .collect::<Vec<(String, DataType)>>();
+
         let listing_options = ListingOptions {
             file_extension: String::new(),
             format: Arc::from(format),
-            table_partition_cols: vec![],
+            table_partition_cols: partition_cols,
             collect_stat: true,
             target_partitions: 1,
             file_sort_order: vec![],