add helper methods

Author: adriangb

datafusion-examples/examples/csv_json_opener.rs

@@ -69,7 +69,7 @@ async fn csv_opener() -> Result<()> {
 
     let config = CsvSource::new(true, b',', b'"')
         .with_comment(Some(b'#'))
-        .with_schema(TableSchema::new(schema, vec![]))
+        .with_schema(TableSchema::new_from_file_schema(schema))
         .with_batch_size(8192)
         .with_projection(&scan_config);
 

datafusion/datasource/src/table_schema.rs

@@ -85,6 +85,11 @@ impl TableSchema {
     /// The table schema is automatically computed by appending the partition columns
     /// to the file schema.
     ///
+    /// You should prefer calling this method over
+    /// chaining [`TableSchema::new_from_file_schema`] and [`TableSchema::with_partition_cols`]
+    /// if you have both the file schema and partition columns available at construction time
+    /// since it avoids re-computing the table schema.
+    ///
     /// # Arguments
     ///
     /// * `file_schema` - Schema of the data files (without partition columns)
@@ -121,6 +126,27 @@ impl TableSchema {
         }
     }
 
+    /// Create a new TableSchema with no partition columns.
+    ///
+    /// You should prefer calling [`TableSchema::new`] if you have partition columns at
+    /// construction time since it avoids re-computing the table schema.
+    pub fn new_from_file_schema(file_schema: SchemaRef) -> Self {
+        Self::new(file_schema, vec![])
+    }
+
+    /// Add partition columns to an existing TableSchema, returning a new instance.
+    ///
+    /// You should prefer calling [`TableSchema::new`] instead of chaining [`TableSchema::new_from_file_schema`]
+    /// into [`TableSchema::with_partition_cols`] if you have partition columns at construction time
+    /// since it avoids re-computing the table schema.
+    pub fn with_partition_cols(mut self, partition_cols: Vec<FieldRef>) -> Self {
+        self.table_partition_cols = partition_cols;
+        let mut builder = SchemaBuilder::from(self.file_schema.as_ref());
+        builder.extend(self.table_partition_cols.iter().cloned());
+        self.table_schema = Arc::new(builder.finish());
+        self
+    }
+
     /// Get the file schema (without partition columns).
     ///
     /// This is the schema of the actual data files on disk.