[GitHub] [arrow-rs] hntd187 commented on a change in pull request #1033: Projection on Schema and RecordBatch

[GitBox <gi...@apache.org>]

hntd187 commented on a change in pull request #1033:
URL: https://github.com/apache/arrow-rs/pull/1033#discussion_r768427611



##########
File path: arrow/src/datatypes/schema.rs
##########
@@ -87,6 +87,18 @@ impl Schema {
         Self { fields, metadata }
     }
 
+
+    /// Returns a new schema with only the specified columns in the new schema
+    /// This carries metadata from the parent schema over as well
+    pub fn project(&self, indices: impl IntoIterator<Item=usize>) -> Result<Schema> {
+        let mut new_fields = vec![];
+        for i in indices {
+            let f = self.fields[i].clone();
+            new_fields.push(f);
+        }

Review comment:
       Yea, that seems good to me, the for loop was the first thing that popped into my head, but I can't think of any reason it's better than yours.

##########
File path: arrow/src/datatypes/schema.rs
##########
@@ -369,4 +381,23 @@ mod tests {
 
         assert_eq!(schema, de_schema);
     }
+
+    #[test]
+    fn test_project() {
+        let mut metadata = HashMap::new();
+        metadata.insert("meta".to_string(), "data".to_string());
+
+        let schema = Schema::new_with_metadata(vec![
+            Field::new("name", DataType::Utf8, false),
+            Field::new("address", DataType::Utf8, false),
+            Field::new("priority", DataType::UInt8, false),
+        ], metadata);
+
+        let projected: Schema = schema.project(vec![0, 2]).unwrap();
+
+        assert_eq!(projected.fields().len(), 2);
+        assert_eq!(projected.fields()[0].name(), "name");
+        assert_eq!(projected.fields()[1].name(), "priority");
+        assert_eq!(projected.metadata.get("meta").unwrap(), "data")
+    }

Review comment:
       Sure, will do




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org