liurenjie1024 commented on code in PR #245:
URL: https://github.com/apache/iceberg-rust/pull/245#discussion_r1537479886
##########
crates/iceberg/src/arrow.rs:
##########
@@ -49,10 +54,17 @@ impl ArrowReaderBuilder {
self
}
+ /// Sets the desired column projection with a list of field ids.
+ pub fn with_field_ids(mut self, field_ids: Vec<usize>) -> Self {
Review Comment:
```suggestion
pub fn with_field_ids(mut self, field_ids: impl
IntoIterator<Item=usize>) -> Self {
```
##########
crates/iceberg/src/scan.rs:
##########
@@ -187,6 +190,22 @@ impl TableScan {
let mut arrow_reader_builder =
ArrowReaderBuilder::new(self.file_io.clone(), self.schema.clone());
+ let mut field_ids = vec![];
+ for column_name in &self.column_names {
+ let field_id =
self.schema.field_id_by_name(column_name).ok_or_else(|| {
Review Comment:
As discussed in #244, we need to do two checks here to ensure that it's
valid:
1. The field is a direct child of schema, e.g. not a nested field. We can do
this by calling `Schema::as_struct::field_by_id::is_some`
2. Ensure that this field is primitive type.
##########
crates/iceberg/src/arrow.rs:
##########
@@ -101,8 +114,53 @@ impl ArrowReader {
.boxed())
}
- fn get_arrow_projection_mask(&self, _task: &FileScanTask) ->
ProjectionMask {
- // TODO: full implementation
- ProjectionMask::all()
+ fn get_arrow_projection_mask(
+ &self,
+ parquet_schema: &SchemaDescriptor,
+ ) -> crate::Result<ProjectionMask> {
+ if self.field_ids.is_empty() {
+ Ok(ProjectionMask::all())
+ } else {
+ let mut column_map = HashMap::new();
+ for (idx, field) in parquet_schema.columns().iter().enumerate() {
+ let field_type = field.self_type();
+ match field_type {
+ Type::PrimitiveType { basic_info, .. } => {
+ if !basic_info.has_id() {
+ return Err(Error::new(
+ ErrorKind::DataInvalid,
+ format!(
+ "Leave column {:?} in schema doesn't have
field id",
+ field_type
+ ),
+ ));
+ }
+ column_map.insert(basic_info.id(), idx);
Review Comment:
We need also a check that their types are matched. How about converting
parquet schema to arrow schema, and uses
[`filter_leaves`](https://arrow.apache.org/rust/arrow_schema/fields/struct.Fields.html#method.filter_leaves)
to do this match check? This way we only need to deal with iceberg schema and
arrow schema.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
