`RecordBatch` normalization (flattening) #6758

ngli-me · 2024-11-20T03:10:29Z

Which issue does this PR close?

Closes #6369.

Rationale for this change

Adds normalization (flattening) for RecordBatch, with normalization via Schema. Based on pandas/pola-rs.

What changes are included in this PR?

Are there any user-facing changes?

…on pola-rs.

… iterative function for `RecordBatch`. Not sure which one is better currently.

ngli-me

I had some questions regarding the implementation of this, since the one example from PyArrow doesn't seem to clarify on the edge cases here. Normalizing the Schema seems fairly straight forward to me, I'm just not sure on

Whether the iterative or recursive approach is better (or something I missed)
If DataType::Struct is the only DataType that requires flattening. To me, it looks like that's the only one that can contained nested Fields.

(I'm also not sure if I'm missing something with unwrapping like a List<Struct>)

Any feedback/help would be appreciated!

arrow-array/src/record_batch.rs

arrow-schema/src/schema.rs

…ch the example from PyArrow.

arrow-array/src/record_batch.rs

…h-flatten

alamb · 2024-12-18T13:04:18Z

@kszlim can you please help review this PR ? You requested the feature and we are currently quite short on review capacity in arrow-rs

alamb

Thank you for this contribution @ngli-me and I apologize for the delay in reviewing.

Hopefully @kszlim can give this a look and help us review / get it moving too.

alamb · 2024-12-18T13:05:38Z

arrow-array/src/record_batch.rs

@@ -394,6 +396,56 @@ impl RecordBatch {
        )
    }

+    /// Normalize a semi-structured [`RecordBatch`] into a flat table.
+    ///
+    /// If max_level is 0, normalizes all levels.


Can you please improve this documentation (maybe copy from the pyarrow version)?

Doucment what max_level means (in addition to that 0)

Document what separator does

provide an example of flatteing a record batch as a doc example?

For example like https://docs.rs/arrow/latest/arrow/index.html#columnar-format

Ah, missed doing this, will do!

alamb · 2024-12-18T13:09:16Z

arrow-schema/src/schema.rs

@@ -413,6 +413,81 @@ impl Schema {
        &self.metadata
    }

+    /// Returns a new schema, normalized based on the max_level
+    /// This carries metadata from the parent schema over as well


Likewise, please document the parametrs to this function and add a documentation example

Sounds good, thanks!

kszlim · 2024-12-19T06:10:58Z

I'll take a look, though please feel free to disregard anything I say and especially defer to the maintainers.

kszlim · 2024-12-19T06:11:59Z

arrow-array/src/record_batch.rs

+        if max_level == 0 {
+            max_level = usize::MAX;
+        }
+        if self.num_rows() == 0 {


Is this worth keeping? I could be reading this wrong, but it looks like there's a lot of code strictly to support normalization for the 0 row case (which is likely very rare)?

Yeah, I'm not sure the best case to handle this.

I think the secondary case can be removed, as I missed that even with a new_empty RecordBatch, it will still have an (empty) columns field, good catch!

For the usize::MAX setting, this was because polars/pandas had this , with a default value of 0. However, since Rust does not have default parameters, I wasn't sure the best way to adapt this.

A possible idea could be to set up like

enum Depth { Default, // All levels Value(usize) }

then do some matching? Might be overkill though.

And of course, there's the option of removing it altogether, although then a value of 0 would mean no change?

kszlim · 2024-12-19T06:12:30Z

arrow-schema/src/schema.rs

@@ -413,6 +413,81 @@ impl Schema {
        &self.metadata
    }

+    /// Returns a new schema, normalized based on the max_level
+    /// This carries metadata from the parent schema over as well
+    pub fn normalize(&self, separator: &str, mut max_level: usize) -> Result<Self, ArrowError> {


I mean this is unused outside of the 0 row case right?

This was to help with the recursion, since using just the helper function would result in the separator as a prefix in the field.name(). I do agree that this is not the best option though, maybe I can count this as a vote against the recursion approach? 😆

kszlim · 2024-12-19T06:17:31Z

arrow-array/src/record_batch.rs

+                    DataType::Struct(ff) => {
+                        // Need to zip these in reverse to maintain original order
+                        for (cff, fff) in c.as_struct().columns().iter().zip(ff.into_iter()).rev() {
+                            let new_key = format!("{}{}{}", f.name(), separator, fff.name());


Not sure if there's a better way to structure it, but is there a way to keep the field name parts in a Vec and create the flattened fields at the end? That allows you to avoid the repeated format! in a deeply nested schema.

Might not be worth the trouble though.

I think this is a good point, this is definitely not my favorite way to do this. I'll have to do some testing and think about it some more, but it may be better to construct the queue with the components of the Field, then go through and construct all of the Fields at the very end.

ngli-me · 2024-12-19T13:33:59Z

Thank you for this contribution @ngli-me and I apologize for the delay in reviewing.

Hopefully @kszlim can give this a look and help us review / get it moving too.

No problem at all, it's the holiday season! Hope everyone's taking a good break.

Appreciate the feedback though! I'll get to work on it :)

nglime added 2 commits November 18, 2024 14:11

Added set up for the example of flattening from pyarrow.

bbd7c8b

Logic for recursive normalizer with a base normalize function, based …

8abcd25

…on pola-rs.

ngli-me changed the title ~~Feature/record batch flatten~~ RecordBatch normalization (flattening) Nov 20, 2024

ngli-me changed the title ~~RecordBatch normalization (flattening)~~ RecordBatch normalization (flattening) Nov 20, 2024

Added recursive normalize function for Schema, and started building…

6bba7d3

… iterative function for `RecordBatch`. Not sure which one is better currently.

github-actions bot added the arrow Changes to the arrow crate label Nov 23, 2024

Built out a bit more of the iterative normalize.

55eb953

ngli-me commented Nov 23, 2024

View reviewed changes

arrow-array/src/record_batch.rs Outdated Show resolved Hide resolved

arrow-array/src/record_batch.rs Outdated Show resolved Hide resolved

arrow-schema/src/schema.rs Show resolved Hide resolved

ngli-me marked this pull request as ready for review November 23, 2024 19:03

ngli-me marked this pull request as draft November 23, 2024 23:30

nglime added 2 commits November 23, 2024 21:03

Fixed normalize function for RecordBatch. Adjusted test case to mat…

30d6294

…ch the example from PyArrow.

Added tests for Schema normalization. Partial tests for RecordBatch.

0ed979d

ngli-me commented Nov 25, 2024

View reviewed changes

arrow-array/src/record_batch.rs Outdated Show resolved Hide resolved

nglime added 2 commits November 24, 2024 21:54

Removed stray comments.

d9d08cd

Commenting out exclamation field.

d1b3260

ngli-me marked this pull request as ready for review November 25, 2024 04:02

nglime added 3 commits December 4, 2024 22:04

Merge remote-tracking branch 'upstream/main' into feature/record-batc…

a12082c

…h-flatten

Fixed test for RecordBatch.

7adda58

Formatting.

9c9c699

alamb reviewed Dec 18, 2024

View reviewed changes

kszlim reviewed Dec 19, 2024

View reviewed changes

ngli-me marked this pull request as draft December 20, 2024 12:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

`RecordBatch` normalization (flattening) #6758

`RecordBatch` normalization (flattening) #6758

ngli-me commented Nov 20, 2024 •

edited

Loading

ngli-me left a comment •

edited

Loading

alamb commented Dec 18, 2024

alamb left a comment

alamb Dec 18, 2024

ngli-me Dec 19, 2024

alamb Dec 18, 2024

ngli-me Dec 19, 2024

kszlim commented Dec 19, 2024

kszlim Dec 19, 2024

ngli-me Dec 19, 2024

kszlim Dec 19, 2024

ngli-me Dec 19, 2024

kszlim Dec 19, 2024

ngli-me Dec 19, 2024

ngli-me commented Dec 19, 2024 •

edited

Loading

RecordBatch normalization (flattening) #6758

Are you sure you want to change the base?

RecordBatch normalization (flattening) #6758

Conversation

ngli-me commented Nov 20, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

ngli-me left a comment • edited Loading

Choose a reason for hiding this comment

alamb commented Dec 18, 2024

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kszlim commented Dec 19, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ngli-me commented Dec 19, 2024 • edited Loading

`RecordBatch` normalization (flattening) #6758

`RecordBatch` normalization (flattening) #6758

ngli-me commented Nov 20, 2024 •

edited

Loading

ngli-me left a comment •

edited

Loading

ngli-me commented Dec 19, 2024 •

edited

Loading