Add explicit column mask for selection in parquet: `ProjectionMask` (#1701) by tustvold · Pull Request #1716 · apache/arrow-rs

tustvold · 2022-05-21T09:57:54Z

Which issue does this PR close?

Closes #1701.
Closes #1653

Rationale for this change

The current API is confusing, surfacing errors at runtime, and is liable to accidental misuse - apache/datafusion#2453 and apache/datafusion#2543.

What changes are included in this PR?

This adds an explicit ColumnMask that replaces the iterators of indices. This makes for a cleaner API that should hopefully make it harder to accidentally misuse.

In particular it also adds the ability to construct a RecordReader based on a mask of root columns, as opposed to leaf columns. This is the core behind apache/datafusion#2543

Are there any user-facing changes?

Yes, this changes the arrow projection API

tustvold · 2022-05-21T09:59:03Z

parquet/src/schema/types.rs

This change to store the root index, as opposed to a copy of the root ptr makes it easier to convert from a root mask to a leaf mask

tustvold · 2022-05-21T10:04:57Z

parquet/src/arrow/mod.rs

+    ///
+    /// i.e. `[0, 1, 2]` will construct the same mask as `[1, 0, 0, 2]`
+    pub fn leaves(
+        schema: &SchemaDescriptor,


The mask could theoretically carry this along and use it as a sanity check, but I'm inclined to think if a user constructs a mask and then uses it on a different schema, it's not something we can reasonably be expected to handle sensibly

As long as a useful error is produced, I agree this is fine behavior

tustvold · 2022-05-21T10:05:24Z

parquet/src/arrow/schema.rs

                    .header_as_schema()
                    .map(arrow::ipc::convert::fb_to_schema)
-                    .ok_or(ArrowError("the message is not Arrow Schema".to_string())),
+                    .ok_or(arrow_err!("the message is not Arrow Schema")),


Drive by cleanup to move to arrow_err! macro

Using arrow_err is not obviously better to me: it is the same verbosity but now requires looking / knowing what the arrow_err! macro does

;0 But it is not worse either

It's consistent with how errors are handled elsewhere in the crate, with arrow_err!, general_err!, etc...

codecov-commenter · 2022-05-21T10:22:31Z

Codecov Report

Merging #1716 (2b327ef) into master (a30e787) will decrease coverage by 0.00%.
The diff coverage is 79.02%.

❗ Current head 2b327ef differs from pull request most recent head 637cb2e. Consider uploading reports for the commit 637cb2e to get more accurate results

@@            Coverage Diff             @@
##           master    #1716      +/-   ##
==========================================
- Coverage   83.32%   83.31%   -0.01%     
==========================================
  Files         195      196       +1     
  Lines       56023    55961      -62     
==========================================
- Hits        46681    46625      -56     
+ Misses       9342     9336       -6

Impacted Files	Coverage Δ
arrow/src/array/transform/mod.rs	`86.96% <ø> (+0.16%)`	⬆️
arrow/src/compute/util.rs	`98.90% <ø> (-0.01%)`	⬇️
parquet/src/arrow/async_reader.rs	`0.00% <0.00%> (ø)`
parquet/src/arrow/mod.rs	`44.44% <44.44%> (ø)`
parquet/src/arrow/schema.rs	`96.81% <75.00%> (+0.02%)`	⬆️
arrow/src/ipc/reader.rs	`88.72% <80.00%> (-0.08%)`	⬇️
parquet/src/schema/types.rs	`83.73% <80.00%> (-0.10%)`	⬇️
arrow/src/array/data.rs	`83.65% <85.71%> (-0.86%)`	⬇️
arrow/src/ipc/writer.rs	`80.37% <100.00%> (+0.06%)`	⬆️
parquet/src/arrow/array_reader/builder.rs	`93.50% <100.00%> (-0.05%)`	⬇️
... and 7 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update a30e787...637cb2e. Read the comment docs.

tustvold · 2022-05-21T10:27:07Z

As a happy side-effect this actually found and fixed a bug in the handling of nested projection pushdown in ParquetRecordBatchStream

tustvold · 2022-05-21T10:28:51Z

parquet/src/arrow/async_reader.rs

-            Some(projection) => {
-                if let Some(col) = projection.iter().find(|x| **x >= num_columns) {
-                    return Err(general_err!(
-                        "column projection {} outside bounds of schema 0..{}",


This check was actually incorrect as it was checking against the arrow schema not the parquet schema. I think this demonstrates the footgun prone nature of the old API

alamb

I think it looks like a great improvement. Thanks @tustvold

cc @viirya @sunchao @nevi-me

alamb · 2022-05-21T10:49:55Z

parquet/src/arrow/mod.rs

+    ///
+    /// i.e. `[0, 1, 2]` will construct the same mask as `[1, 0, 0, 2]`
+    pub fn leaves(
+        schema: &SchemaDescriptor,


As long as a useful error is produced, I agree this is fine behavior

alamb · 2022-05-21T10:50:33Z

parquet/src/arrow/mod.rs

+/// A [`ProjectionMask`] identifies a set of columns within a potentially nested schema to project
+#[derive(Debug, Clone)]
+pub struct ProjectionMask {
+    /// A mask of


The comment seems to be truncated.

Also, since we have Bitmap and all the associated handling in Arrow, I wonder if it is worth using that (though a Vec<bool> is nice and simple

Describing in the docstring that a mask of None means All is probably also a good idea as well as which schema a ProjectionMasks indexes refer to (I think Parquet)

Also, since we have Bitmap and all the associated handling in Arrow, I wonder if it is worth using that (though a Vec is nice and simple

Let's stick with simple, and maybe if/when we promote this construct to arrow-rs we can switch to using Bitmap

alamb · 2022-05-21T10:51:55Z

parquet/src/arrow/mod.rs

+
+    /// Create a [`ProjectionMask`] which selects only the specified leaf columns
+    ///
+    /// Note: repeated or out of order indices will not impact the final mask


Nice -- so the idea is that you enforce masks having in-order indicies by wrapping them in a ProjectionMask which enforces this invariant during construction 👍

alamb · 2022-05-21T10:52:56Z

parquet/src/arrow/mod.rs

+        Self { mask: None }
+    }
+
+    /// Create a [`ProjectionMask`] which selects only the specified leaf columns


Can you please explain (or provide a link to something that explains) leaves and roots and what "order" they are in. I think it refers to the parquet schema (or maybe the arrow schema and types within Structs / LIsts / others nested types?)

alamb · 2022-05-21T10:55:11Z

parquet/src/arrow/schema.rs

    parquet_to_arrow_schema_by_columns(
        parquet_schema,
-        0..parquet_schema.columns().len(),
+        ProjectionMask::all(),


alamb · 2022-05-21T10:55:47Z

parquet/src/arrow/schema.rs

+pub fn parquet_to_arrow_schema_by_columns(
    parquet_schema: &SchemaDescriptor,
-    column_indices: T,
+    mask: ProjectionMask,


I wonder if this needs an owned mask or if it could be taken by reference

Suggested change

mask: ProjectionMask,

mask: &ProjectionMask,

Currently yes, it gets moved into the Visitor. Theoretically it could borrow and have lifetimes, but in most cases I suspect we have the mask by value anyway.

Edit: This might be an argument to move to arrow Bitmap, as that is internally refcounted... Future PR me thinks

alamb · 2022-05-21T10:58:21Z

parquet/src/arrow/schema.rs

                    .header_as_schema()
                    .map(arrow::ipc::convert::fb_to_schema)
-                    .ok_or(ArrowError("the message is not Arrow Schema".to_string())),
+                    .ok_or(arrow_err!("the message is not Arrow Schema")),


Using arrow_err is not obviously better to me: it is the same verbosity but now requires looking / knowing what the arrow_err! macro does

;0 But it is not worse either

alamb · 2022-05-21T10:58:47Z

parquet/src/arrow/schema.rs

        // required int64 leaf5;

        let parquet_schema = SchemaDescriptor::new(Arc::new(parquet_group_type));
+        let mask = ProjectionMask::leaves(&parquet_schema, [3, 0, 4]);


this is much nicer to read / reason about

tustvold added the api-change Changes to the arrow API label May 21, 2022

github-actions bot added the parquet Changes to the parquet crate label May 21, 2022

Add explicit column mask construction (apache#1701)

489d54a

tustvold force-pushed the explicit-column-mask branch from 284724e to 489d54a Compare May 21, 2022 10:02

tustvold commented May 21, 2022

View reviewed changes

This was referenced May 21, 2022

Incorrect Parquet Projection For Nested Types apache/datafusion#2453

Closed

Introduce ProjectionMask To Allow Nested Projection Pushdown apache/datafusion#2581

Open

Fix ParquetRecordBatchStream

a7f349c

tustvold commented May 21, 2022

View reviewed changes

Fix docs

4a58fd3

alamb changed the title ~~Add explicit column mask construction (#1701)~~ Add explicit column mask construction in parquet: ProjectionMask (#1701) May 21, 2022

alamb approved these changes May 21, 2022

View reviewed changes

tustvold added 2 commits May 21, 2022 12:14

Fix async_reader test

24904a8

Review feedback

637cb2e

tustvold merged commit ca1d85f into apache:master May 24, 2022

alamb changed the title ~~Add explicit column mask construction in parquet: ProjectionMask (#1701)~~ Add explicit column mask for selection in parquet: ProjectionMask (#1701) May 26, 2022

tustvold mentioned this pull request Jun 1, 2022

Simplify ParquetFileArrowReader Metadata API #1773

Merged

Comments

Conversation

tustvold commented May 21, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented May 21, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

tustvold commented May 21, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alamb left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

tustvold May 23, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tustvold commented May 21, 2022 •

edited

Loading

codecov-commenter commented May 21, 2022 •

edited

Loading

tustvold May 23, 2022 •

edited

Loading