Skip to content

Instantly share code, notes, and snippets.

@alamb
Created October 16, 2022 10:08
Show Gist options
  • Save alamb/acd653c49e318ff70672b61325ba3443 to your computer and use it in GitHub Desktop.
Save alamb/acd653c49e318ff70672b61325ba3443 to your computer and use it in GitHub Desktop.
Nested structure parquet encoding example
#[test]
fn test_foo() {
let mut a = ListBuilder::new(Int32Builder::new());
a.values().append_value(1);
a.append(true);
a.append(false);
a.values().append_null();
a.values().append_null();
a.append(true);
let values = Arc::new(a.finish()) as ArrayRef;
let mut builder = LevelInfoBuilder::try_new(
&Field::new("test", values.data_type().clone(), true),
Default::default(),
)
.unwrap();
builder.write(&values, 0..3);
let levels = builder.finish();
assert_eq!(levels.len(), 1);
let list_level = levels.get(0).unwrap();
let expected_level = LevelInfo {
def_levels: Some(vec![3, 0, 2, 2]),
rep_levels: Some(vec![0, 0, 0, 1]),
non_null_indices: vec![0],
max_def_level: 3,
max_rep_level: 1,
};
assert_eq!(list_level, &expected_level);
let b2 = DataType::List(Box::new(Field::new("element", DataType::Int32, false)));
let s = StructBuilder::new(
vec![
Field::new("b1", DataType::Int32, false),
Field::new("b2", b2, true),
],
vec![
Box::new(Int32Builder::new()),
Box::new(ListBuilder::new(Int32Builder::new())),
],
);
let mut b = ListBuilder::new(s);
b.values()
.field_builder::<Int32Builder>(0)
.unwrap()
.append_value(1);
b.values()
.field_builder::<ListBuilder<Int32Builder>>(1)
.unwrap()
.append(false);
b.values().append(true);
b.values()
.field_builder::<Int32Builder>(0)
.unwrap()
.append_value(1);
b.values()
.field_builder::<ListBuilder<Int32Builder>>(1)
.unwrap()
.values()
.append_value(3);
b.values()
.field_builder::<ListBuilder<Int32Builder>>(1)
.unwrap()
.values()
.append_value(4);
b.values()
.field_builder::<ListBuilder<Int32Builder>>(1)
.unwrap()
.append(true);
b.values().append(true);
b.append(true);
b.values()
.field_builder::<Int32Builder>(0)
.unwrap()
.append_value(2);
b.values()
.field_builder::<ListBuilder<Int32Builder>>(1)
.unwrap()
.append(false);
b.values().append(true);
b.append(true);
b.values()
.field_builder::<Int32Builder>(0)
.unwrap()
.append_value(0);
b.values()
.field_builder::<ListBuilder<Int32Builder>>(1)
.unwrap()
.append(false);
b.values().append_null();
b.append(true);
let values = Arc::new(b.finish()) as ArrayRef;
let mut builder = LevelInfoBuilder::try_new(
&Field::new("test", values.data_type().clone(), false),
Default::default(),
)
.unwrap();
builder.write(&values, 0..3);
let levels = builder.finish();
assert_eq!(levels.len(), 2);
let list_level = levels.get(0).unwrap();
let expected_level = LevelInfo {
def_levels: Some(vec![2, 2, 2, 1]),
rep_levels: Some(vec![0, 1, 0, 0]),
non_null_indices: vec![0, 1, 2],
max_def_level: 2,
max_rep_level: 1,
};
assert_eq!(list_level, &expected_level);
let list_level = levels.get(1).unwrap();
let expected_level = LevelInfo {
def_levels: Some(vec![2, 4, 4, 2, 1]),
rep_levels: Some(vec![0, 1, 2, 0, 0]),
non_null_indices: vec![0, 1],
max_def_level: 4,
max_rep_level: 2,
};
assert_eq!(list_level, &expected_level);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment