use std::collections::HashMap;
use std::sync::Arc;
use arrow::datatypes::{DataType, Field, FieldRef, Fields};
fn nested_child(dt: &DataType) -> Option<&DataType> {
match dt {
DataType::List(f)
| DataType::LargeList(f)
| DataType::ListView(f)
| DataType::LargeListView(f)
| DataType::FixedSizeList(f, _)
| DataType::Map(f, _) => Some(f.data_type()),
DataType::Dictionary(_, value) => Some(value),
DataType::RunEndEncoded(_, value) => Some(value.data_type()),
_ => None,
}
}
pub(crate) fn clip_for_cast(
physical: &DataType,
cast_target: &DataType,
) -> Option<(Vec<usize>, DataType)> {
let total = count_leaves(physical);
let mut kept = Vec::new();
let mut next_leaf = 0;
let mut unclippable = false;
let pruned_type = clip_type(
physical,
cast_target,
&mut next_leaf,
&mut kept,
&mut unclippable,
);
debug_assert_eq!(next_leaf, total, "leaf accounting must cover the type");
if unclippable || kept.is_empty() || kept.len() >= total {
return None;
}
Some((kept, pruned_type))
}
pub(crate) fn count_leaves(dt: &DataType) -> usize {
match dt {
DataType::Struct(fields) => {
fields.iter().map(|f| count_leaves(f.data_type())).sum()
}
_ => nested_child(dt).map_or(1, count_leaves),
}
}
pub(crate) fn contains_struct(dt: &DataType) -> bool {
matches!(dt, DataType::Struct(_)) || nested_child(dt).is_some_and(contains_struct)
}
const LINEAR_FIELD_SCAN_MAX: usize = 8;
fn lookup_field<'a>(
fields: &'a Fields,
by_name: &Option<HashMap<&'a str, &'a FieldRef>>,
name: &str,
) -> Option<&'a FieldRef> {
match by_name {
Some(map) => map.get(name).copied(),
None => fields.iter().find(|f| f.name() == name),
}
}
fn clip_type(
physical: &DataType,
target: &DataType,
next_leaf: &mut usize,
kept: &mut Vec<usize>,
unclippable: &mut bool,
) -> DataType {
match (physical, target) {
(DataType::Struct(p_children), DataType::Struct(t_children)) => {
let t_by_name = (t_children.len() > LINEAR_FIELD_SCAN_MAX).then(|| {
let mut map = HashMap::with_capacity(t_children.len());
for tc in t_children.iter() {
map.entry(tc.name().as_str()).or_insert(tc);
}
map
});
let kept_children: Fields = p_children
.iter()
.filter_map(|pc| {
let Some(tc) = lookup_field(t_children, &t_by_name, pc.name()) else {
skip_leaves(pc.data_type(), next_leaf);
return None;
};
let before = kept.len();
let pruned = clip_type(
pc.data_type(),
tc.data_type(),
next_leaf,
kept,
unclippable,
);
if kept.len() == before {
*unclippable = true;
}
Some(field_with_type(pc, pruned))
})
.collect();
DataType::Struct(kept_children)
}
(DataType::List(p_item), DataType::List(t_item)) => {
let pruned = clip_type(
p_item.data_type(),
t_item.data_type(),
next_leaf,
kept,
unclippable,
);
DataType::List(field_with_type(p_item, pruned))
}
(DataType::LargeList(p_item), DataType::LargeList(t_item)) => {
let pruned = clip_type(
p_item.data_type(),
t_item.data_type(),
next_leaf,
kept,
unclippable,
);
DataType::LargeList(field_with_type(p_item, pruned))
}
_ => keep_all_leaves(physical, next_leaf, kept),
}
}
fn keep_all_leaves(
dt: &DataType,
next_leaf: &mut usize,
kept: &mut Vec<usize>,
) -> DataType {
let n = count_leaves(dt);
kept.extend(*next_leaf..*next_leaf + n);
*next_leaf += n;
dt.clone()
}
fn skip_leaves(dt: &DataType, next_leaf: &mut usize) {
*next_leaf += count_leaves(dt);
}
#[derive(Debug, Clone)]
pub(crate) struct CastColumnAccess {
pub(crate) root_index: usize,
pub(crate) target_type: DataType,
}
pub(crate) fn field_with_type(field: &Field, data_type: DataType) -> FieldRef {
Arc::new(field.clone().with_data_type(data_type))
}
#[cfg(test)]
mod tests {
use super::*;
fn utf8(name: &str) -> Field {
Field::new(name, DataType::Utf8, true)
}
fn int64(name: &str) -> Field {
Field::new(name, DataType::Int64, true)
}
fn struct_of(fields: Vec<Field>) -> DataType {
DataType::Struct(Fields::from(fields))
}
fn list_of(item: DataType) -> DataType {
DataType::List(Arc::new(Field::new("item", item, true)))
}
#[test]
fn count_leaves_shapes() {
assert_eq!(count_leaves(&DataType::Int32), 1);
assert_eq!(count_leaves(&struct_of(vec![utf8("a"), int64("b")])), 2);
assert_eq!(
count_leaves(&list_of(struct_of(vec![
utf8("a"),
struct_of(vec![int64("x"), int64("y")]).into_field("s")
]))),
3
);
let map = DataType::Map(
Arc::new(Field::new(
"entries",
struct_of(vec![utf8("key"), int64("value")]),
false,
)),
false,
);
assert_eq!(count_leaves(&map), 2);
let dict =
DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
assert_eq!(count_leaves(&dict), 1);
assert_eq!(
count_leaves(&DataType::Dictionary(
Box::new(DataType::Int32),
Box::new(struct_of(vec![utf8("a"), int64("b")]))
)),
2
);
assert_eq!(
count_leaves(&DataType::RunEndEncoded(
Arc::new(Field::new("run_ends", DataType::Int32, false)),
Arc::new(Field::new(
"values",
struct_of(vec![utf8("a"), int64("b")]),
true
))
)),
2
);
}
#[test]
fn contains_struct_shapes() {
assert!(!contains_struct(&DataType::Int32));
assert!(!contains_struct(&list_of(DataType::Int32)));
assert!(contains_struct(&struct_of(vec![int64("a")])));
assert!(contains_struct(&list_of(struct_of(vec![int64("a")]))));
assert!(contains_struct(&DataType::LargeList(Arc::new(Field::new(
"item",
struct_of(vec![int64("a")]),
true
)))));
assert!(contains_struct(&DataType::Dictionary(
Box::new(DataType::Int32),
Box::new(struct_of(vec![int64("a")]))
)));
assert!(!contains_struct(&DataType::Dictionary(
Box::new(DataType::Int32),
Box::new(DataType::Utf8)
)));
assert!(contains_struct(&DataType::Map(
Arc::new(Field::new(
"entries",
struct_of(vec![utf8("key"), int64("value")]),
false
)),
false
)));
}
#[test]
fn clip_struct_subset() {
let physical = struct_of(vec![utf8("a"), int64("b"), utf8("c")]);
let target = struct_of(vec![int64("b")]);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![1]);
assert_eq!(emitted, struct_of(vec![int64("b")]));
}
#[test]
fn clip_struct_reordered_target() {
let physical = struct_of(vec![utf8("a"), int64("b"), utf8("c")]);
let target = struct_of(vec![utf8("c"), utf8("a")]);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0, 2]);
assert_eq!(emitted, struct_of(vec![utf8("a"), utf8("c")]));
}
#[test]
fn clip_struct_target_field_missing_from_physical() {
let physical = struct_of(vec![utf8("a"), int64("b")]);
let target = struct_of(vec![utf8("a"), int64("z")]);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0]);
assert_eq!(emitted, struct_of(vec![utf8("a")]));
}
#[test]
fn clip_keeps_physical_leaf_types() {
let physical =
struct_of(vec![Field::new("x", DataType::Int32, true), utf8("pad")]);
let target = struct_of(vec![int64("x")]);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0]);
assert_eq!(
emitted,
struct_of(vec![Field::new("x", DataType::Int32, true)])
);
}
#[test]
fn clip_nested_struct() {
let inner_physical = struct_of(vec![int64("x"), utf8("pad_inner")]);
let physical = struct_of(vec![
inner_physical.clone().into_field("inner"),
utf8("pad_outer"),
]);
let target = struct_of(vec![struct_of(vec![int64("x")]).into_field("inner")]);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0]);
assert_eq!(
emitted,
struct_of(vec![struct_of(vec![int64("x")]).into_field("inner")])
);
}
#[test]
fn clip_list_of_struct() {
let physical = list_of(struct_of(vec![int64("x"), utf8("y"), utf8("pad")]));
let target = list_of(struct_of(vec![int64("x"), utf8("y")]));
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0, 1]);
assert_eq!(emitted, list_of(struct_of(vec![int64("x"), utf8("y")])));
}
#[test]
fn clip_two_level_nested_list_of_struct() {
let physical = list_of(struct_of(vec![
int64("a"),
utf8("pad"),
struct_of(vec![int64("x"), utf8("y")]).into_field("aux"),
list_of(struct_of(vec![int64("g"), utf8("pad2")])).into_field("items"),
]));
let target = list_of(struct_of(vec![
int64("a"),
list_of(struct_of(vec![int64("g")])).into_field("items"),
]));
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0, 4]);
assert_eq!(
emitted,
list_of(struct_of(vec![
int64("a"),
list_of(struct_of(vec![int64("g")])).into_field("items"),
]))
);
}
#[test]
fn clip_large_list_of_struct() {
let item = |fields| Arc::new(Field::new("item", struct_of(fields), true));
let physical = DataType::LargeList(item(vec![int64("x"), utf8("pad")]));
let target = DataType::LargeList(item(vec![int64("x")]));
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0]);
assert_eq!(emitted, DataType::LargeList(item(vec![int64("x")])));
}
#[test]
fn no_clip_on_wrapper_mismatch() {
let physical = list_of(struct_of(vec![int64("x"), utf8("pad")]));
let target = DataType::LargeList(Arc::new(Field::new(
"item",
struct_of(vec![int64("x")]),
true,
)));
assert!(clip_for_cast(&physical, &target).is_none());
}
#[test]
fn no_clip_on_map() {
let entries = |fields| Arc::new(Field::new("entries", struct_of(fields), false));
let physical =
DataType::Map(entries(vec![utf8("key"), int64("a"), int64("b")]), false);
let target = DataType::Map(entries(vec![utf8("key"), int64("a")]), false);
assert!(clip_for_cast(&physical, &target).is_none());
}
#[test]
fn no_clip_when_identical() {
let t = struct_of(vec![utf8("a"), int64("b")]);
assert!(clip_for_cast(&t, &t).is_none());
}
#[test]
fn no_clip_on_primitives() {
assert!(clip_for_cast(&DataType::Int32, &DataType::Int64).is_none());
}
#[test]
fn no_clip_on_zero_overlap() {
let physical = struct_of(vec![utf8("a"), int64("b")]);
let target = struct_of(vec![utf8("z")]);
assert!(clip_for_cast(&physical, &target).is_none());
}
#[test]
fn no_clip_when_nested_struct_level_has_no_overlap() {
let physical = struct_of(vec![
struct_of(vec![int64("a"), int64("b")]).into_field("inner"),
int64("c"),
]);
let target = struct_of(vec![
struct_of(vec![int64("z")]).into_field("inner"),
int64("c"),
]);
assert!(clip_for_cast(&physical, &target).is_none());
}
#[test]
fn no_clip_when_nested_list_struct_level_has_no_overlap() {
let physical = struct_of(vec![
list_of(struct_of(vec![int64("a"), int64("b")])).into_field("items"),
int64("c"),
]);
let target = struct_of(vec![
list_of(struct_of(vec![int64("z")])).into_field("items"),
int64("c"),
]);
assert!(clip_for_cast(&physical, &target).is_none());
}
#[test]
fn clip_wide_struct_matches_by_name() {
let width = LINEAR_FIELD_SCAN_MAX * 4;
let physical = struct_of((0..width).map(|i| int64(&format!("f{i}"))).collect());
let target = struct_of(
(0..width)
.rev()
.filter(|i| i % 2 == 0)
.map(|i| int64(&format!("f{i}")))
.collect(),
);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, (0..width).filter(|i| i % 2 == 0).collect::<Vec<_>>());
assert_eq!(
emitted,
struct_of(
(0..width)
.filter(|i| i % 2 == 0)
.map(|i| int64(&format!("f{i}")))
.collect()
)
);
}
#[test]
fn clip_keeps_duplicate_physical_field_names() {
let physical = struct_of(vec![int64("a"), utf8("pad"), int64("a")]);
let target = struct_of(vec![int64("a")]);
let (kept, emitted) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0, 2]);
assert_eq!(emitted, struct_of(vec![int64("a"), int64("a")]));
}
#[test]
fn reader_drops_struct_child_with_no_selected_leaves() {
use arrow::array::{ArrayRef, Int64Array, StructArray};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::arrow::{ArrowWriter, ProjectionMask};
let inner_fields = Fields::from(vec![int64("a"), int64("b")]);
let outer_fields = Fields::from(vec![
Field::new("inner", DataType::Struct(inner_fields.clone()), true),
int64("c"),
]);
let inner: ArrayRef = Arc::new(StructArray::new(
inner_fields,
vec![
Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
Arc::new(Int64Array::from(vec![3, 4])) as ArrayRef,
],
None,
));
let outer = StructArray::new(
outer_fields.clone(),
vec![inner, Arc::new(Int64Array::from(vec![5, 6])) as ArrayRef],
None,
);
let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new(
"s",
DataType::Struct(outer_fields),
true,
)]));
let batch =
RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(outer)]).unwrap();
let file = tempfile::NamedTempFile::new().unwrap();
let mut writer =
ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
let builder =
ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap();
assert_eq!(builder.parquet_schema().num_columns(), 3);
let mask = ProjectionMask::leaves(builder.parquet_schema(), [2usize]);
let reader = builder.with_projection(mask).build().unwrap();
let out: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
assert_eq!(
out[0].schema().field(0).data_type(),
&struct_of(vec![int64("c")]),
"the fully masked `inner` child is dropped, not emitted as an empty struct"
);
}
#[test]
fn arrow_reader_emits_clipped_type_for_masked_list_struct() {
use arrow::array::{
Array, ArrayRef, Int64Array, ListArray, StringArray, StructArray,
};
use arrow::buffer::{NullBuffer, OffsetBuffer};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
use parquet::arrow::{ArrowWriter, ProjectionMask};
let item_fields = Fields::from(vec![int64("x"), utf8("y"), utf8("pad")]);
let item_field = Arc::new(Field::new(
"item",
DataType::Struct(item_fields.clone()),
true,
));
let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new(
"events",
DataType::List(Arc::clone(&item_field)),
true,
)]));
let columns: Vec<ArrayRef> = vec![
Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])),
Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])),
Arc::new(StringArray::from(vec![Some("p0"), None, Some("p2")])),
];
let struct_validity = NullBuffer::from(vec![true, false, true]);
let values = StructArray::new(item_fields, columns, Some(struct_validity));
let list_validity = NullBuffer::from(vec![true, false, true]);
let events = ListArray::new(
item_field,
OffsetBuffer::from_lengths([2, 0, 1]),
Arc::new(values),
Some(list_validity),
);
let batch =
RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(events)]).unwrap();
let file = tempfile::NamedTempFile::new().unwrap();
let mut writer =
ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap();
writer.write(&batch).unwrap();
writer.close().unwrap();
let physical = batch.schema().field(0).data_type().clone();
let target = list_of(struct_of(vec![int64("x"), utf8("y")]));
let (kept, predicted_type) = clip_for_cast(&physical, &target).unwrap();
assert_eq!(kept, vec![0, 1]);
let builder =
ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap();
let mask = ProjectionMask::leaves(builder.parquet_schema(), kept.iter().copied());
let reader = builder.with_projection(mask).build().unwrap();
let out: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
assert_eq!(out.len(), 1);
let out = &out[0];
assert_eq!(out.schema().field(0).data_type(), &predicted_type);
let events = out.column(0).as_any().downcast_ref::<ListArray>().unwrap();
assert!(events.is_valid(0));
assert!(events.is_null(1));
assert!(events.is_valid(2));
let structs = events
.values()
.as_any()
.downcast_ref::<StructArray>()
.unwrap();
assert_eq!(structs.len(), 3);
assert!(structs.is_valid(0));
assert!(structs.is_null(1));
assert!(structs.is_valid(2));
let x = structs
.column(0)
.as_any()
.downcast_ref::<Int64Array>()
.unwrap();
assert_eq!(x.value(0), 1);
assert_eq!(x.value(2), 3);
}
trait IntoField {
fn into_field(self, name: &str) -> Field;
}
impl IntoField for DataType {
fn into_field(self, name: &str) -> Field {
Field::new(name, self, true)
}
}
}