etseidl commented on code in PR #11475: URL: https://github.com/apache/arrow-rs/pull/11475#discussion_r4552196816
########## parquet/src/file/metadata/column_chunk.rs: ########## @@ -0,0 +1,380 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Sparse storage keyed by row-group and column coordinates. + +use crate::errors::{ParquetError, Result}; +use crate::file::metadata::memory::HeapSize; +use std::collections::HashMap; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub(crate) struct ColumnChunkKey(u64); + +impl ColumnChunkKey { + pub(crate) fn try_new(row_group_idx: usize, column_idx: usize) -> Result<Self> { + let row_group_idx = u32::try_from(row_group_idx).map_err(|_| { + ParquetError::General(format!( + "row group index exceeds packed coordinate limit: {row_group_idx}" + )) + })?; + let column_idx = u32::try_from(column_idx).map_err(|_| { + ParquetError::General(format!( + "column index exceeds packed coordinate limit: {column_idx}" + )) + })?; + Ok(Self( + (u64::from(row_group_idx) << u32::BITS) | u64::from(column_idx), + )) + } + + pub(crate) fn coordinates(self) -> (usize, usize) { + ((self.0 >> u32::BITS) as usize, (self.0 as u32) as usize) + } +} + +impl HeapSize for ColumnChunkKey { + fn heap_size(&self) -> usize { + 0 + } +} + +fn validate_dimensions(num_row_groups: usize, num_columns: usize) -> Result<()> { + // Parquet metadata collection lengths are represented by signed 32-bit Thrift values. + i32::try_from(num_row_groups).map_err(|_| { + ParquetError::General(format!( + "row group count exceeds i32::MAX: {num_row_groups}" + )) + })?; + i32::try_from(num_columns).map_err(|_| { + ParquetError::General(format!("column count exceeds i32::MAX: {num_columns}")) + })?; + Ok(()) +} + +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct ColumnChunkMap<T> { + num_row_groups: usize, + num_columns: usize, + entries: HashMap<ColumnChunkKey, T>, +} + +impl<T> ColumnChunkMap<T> { + pub(crate) fn try_new(num_row_groups: usize, num_columns: usize) -> Result<Self> { + validate_dimensions(num_row_groups, num_columns)?; + Ok(Self { + num_row_groups, + num_columns, + entries: HashMap::new(), + }) + } + + pub(crate) fn num_row_groups(&self) -> usize { + self.num_row_groups + } + + pub(crate) fn num_columns(&self) -> usize { + self.num_columns + } + + pub(crate) fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + fn key(&self, row_group_idx: usize, column_idx: usize) -> Result<ColumnChunkKey> { + if row_group_idx >= self.num_row_groups || column_idx >= self.num_columns { + return Err(ParquetError::General(format!( + "column chunk coordinate out of bounds: row group {row_group_idx}, column \ + {column_idx}; dimensions are {} by {}", + self.num_row_groups, self.num_columns + ))); + } + ColumnChunkKey::try_new(row_group_idx, column_idx) + } + + /// Inserts or replaces an entry. + pub(crate) fn try_insert( + &mut self, + row_group_idx: usize, + column_idx: usize, + value: T, + ) -> Result<Option<T>> { + let key = self.key(row_group_idx, column_idx)?; + if let Some(existing) = self.entries.get_mut(&key) { + return Ok(Some(std::mem::replace(existing, value))); + } + + self.entries.try_reserve(1).map_err(|err| { + ParquetError::General(format!( + "failed to reserve storage for column chunk at row group {row_group_idx}, \ + column {column_idx}: {err}" + )) + })?; + Ok(self.entries.insert(key, value)) + } + + pub(crate) fn try_freeze(self) -> Result<ColumnChunkEntries<T>> { + validate_dimensions(self.num_row_groups, self.num_columns)?; + let num_entries = self.entries.len(); + let mut entries = Vec::new(); + entries.try_reserve_exact(num_entries).map_err(|err| { Review Comment: Yeah, it's a mania my agent has picked up regarding fallibility. I'm inclined to keep it as a defense against future issues claiming some malformed input can overflow blah blah blah, but if you'd prefer we keep allocations consistent I can remove it (and the one right above it). -- 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]
