leaves12138 commented on code in PR #536: URL: https://github.com/apache/paimon-rust/pull/536#discussion_r3607266534
########## crates/paimon/src/arrow/shredding/map.rs: ########## @@ -0,0 +1,2762 @@ +// 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. + +//! MAP shared-shredding (PIP-43), compatible with Java's +//! `org.apache.paimon.data.shredding.MapSharedShredding*`. +//! +//! A logical `MAP<STRING, T>` field is stored physically as +//! `ROW<__field_mapping: ARRAY<INT>, __col_0..__col_{K-1}: T, __overflow: MAP<INT, T>>`: +//! the first `K` entries of each row (in row order) go to the shared columns, +//! the rest go to the overflow map keyed by a file-local field id. The field +//! dictionary and column statistics are committed into the `ARROW:schema` +//! footer metadata at close time so readers can rebuild the logical maps. + +use super::{option_usize, FieldMetadata, ShreddingReadPlan, ShreddingWritePlan}; +use crate::arrow::{build_target_arrow_schema, paimon_type_to_arrow}; +use crate::spec::{ArrayType, DataField, DataType, IntType, MapType, RowType}; +use crate::{Error, Result}; +use arrow_array::{ + Array, ArrayRef, Int32Array, ListArray, MapArray, RecordBatch, StringArray, StructArray, + UInt32Array, +}; +use arrow_buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow_schema::{DataType as ArrowDataType, Fields, Schema as ArrowSchema}; +use arrow_select::interleave::interleave; +use arrow_select::take::take; +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::sync::Arc; + +// --------------------------------------------------------------------------- +// Metadata keys (mirroring Java's MapShreddingDefine / MapSharedShreddingDefine) +// --------------------------------------------------------------------------- + +pub(crate) const MAP_STORAGE_LAYOUT_KEY: &str = "paimon.map.storage-layout"; +pub(crate) const MAP_STORAGE_LAYOUT_SHARED_SHREDDING: &str = "shared-shredding"; +const VERSION_KEY: &str = "paimon.map.shared-shredding.version"; +const CURRENT_VERSION: i32 = 1; +const FIELD_DICT_KEY: &str = "paimon.map.shared-shredding.field-dict"; +const FIELD_DICT_COMPRESSION_KEY: &str = "paimon.map.shared-shredding.field-dict-compression"; +const FIELD_DICT_ORIGINAL_SIZE_KEY: &str = "paimon.map.shared-shredding.field-dict-original-size"; +const FIELD_COLUMNS_KEY: &str = "paimon.map.shared-shredding.field-columns"; +const OVERFLOW_SET_KEY: &str = "paimon.map.shared-shredding.overflow-set"; +const NUM_COLUMNS_KEY: &str = "paimon.map.shared-shredding.num-columns"; +const MAX_ROW_WIDTH_KEY: &str = "paimon.map.shared-shredding.max-row-width"; + +const FIELD_MAPPING_NAME: &str = "__field_mapping"; +const OVERFLOW_NAME: &str = "__overflow"; + +fn physical_column_name(index: usize) -> String { + format!("__col_{index}") +} + +// --------------------------------------------------------------------------- +// Options (mirroring Java's CoreOptions MAP_STORAGE_LAYOUT / +// MAP_SHARED_SHREDDING_MAX_COLUMNS accessed as `fields.<name>.<key>`) +// --------------------------------------------------------------------------- + +const MAP_STORAGE_LAYOUT_OPTION_SUFFIX: &str = "map.storage-layout"; +const MAP_SHARED_SHREDDING_MAX_COLUMNS_OPTION_SUFFIX: &str = "map.shared-shredding.max-columns"; +const DEFAULT_MAP_SHARED_SHREDDING_MAX_COLUMNS: usize = 256; + +/// Mirrors Java's `MapSharedShreddingWritePlanFactory.INFER_BUFFER_ROW_COUNT`. +pub(crate) const MAP_SHREDDING_INFER_BUFFER_ROW_COUNT: usize = 1; + +/// One top-level field configured for shared-shredding. +pub(crate) struct MapShreddingFieldConfig { + /// Index of the field in the logical write fields. + pub(crate) field_index: usize, + pub(crate) field_name: String, + pub(crate) max_columns: usize, +} + +/// Whether the type is a MAP with a VARCHAR key, mirroring Java's +/// `MapSharedShreddingUtils.isShreddingKeyMap`. +pub(crate) fn is_shredding_key_map(data_type: &DataType) -> bool { + matches!(data_type, DataType::Map(map_type) if matches!(map_type.key_type(), DataType::VarChar(_))) +} + +/// Detect top-level fields configured with +/// `fields.<name>.map.storage-layout=shared-shredding`, mirroring Java's +/// `MapSharedShreddingUtils.detectShreddingColumns` + `buildColumnToNumColumns`. +pub(crate) fn detect_map_shredding_fields( + fields: &[DataField], + options: &HashMap<String, String>, +) -> Result<Vec<MapShreddingFieldConfig>> { + let mut configs = Vec::new(); + for (field_index, field) in fields.iter().enumerate() { + if !is_shredding_key_map(field.data_type()) { + continue; + } + let layout_key = format!( + "fields.{}.{}", + field.name(), + MAP_STORAGE_LAYOUT_OPTION_SUFFIX + ); + let Some(layout) = options.get(&layout_key) else { + continue; + }; + if layout.eq_ignore_ascii_case("default") { + continue; + } + if !layout.eq_ignore_ascii_case(MAP_STORAGE_LAYOUT_SHARED_SHREDDING) { + return Err(Error::DataInvalid { + message: format!( + "Invalid value '{layout}' for option '{layout_key}': expected 'default' or 'shared-shredding'" + ), + source: None, + }); + } + let max_columns_key = format!( + "fields.{}.{}", + field.name(), + MAP_SHARED_SHREDDING_MAX_COLUMNS_OPTION_SUFFIX + ); + let max_columns = option_usize( + options, + &max_columns_key, + DEFAULT_MAP_SHARED_SHREDDING_MAX_COLUMNS, + )?; + if max_columns == 0 { + return Err(Error::DataInvalid { + message: format!( + "options {MAP_SHARED_SHREDDING_MAX_COLUMNS_OPTION_SUFFIX} must > 0" + ), + source: None, + }); + } + configs.push(MapShreddingFieldConfig { + field_index, + field_name: field.name().to_string(), + max_columns, + }); + } + Ok(configs) +} + +// --------------------------------------------------------------------------- +// Field dictionary (mirroring Java's MapSharedShreddingFieldDict) +// --------------------------------------------------------------------------- + +/// File-local field name -> field id dictionary for one shared-shredding MAP column. +struct FieldDict { + name_to_id: BTreeMap<String, i32>, + next_id: i32, +} + +impl FieldDict { + fn new() -> Self { + Self { + name_to_id: BTreeMap::new(), + next_id: 0, + } + } + + fn get_or_assign(&mut self, name: &str) -> i32 { + if let Some(&id) = self.name_to_id.get(name) { + return id; + } + let new_id = self.next_id; + self.next_id += 1; + self.name_to_id.insert(name.to_string(), new_id); + new_id + } +} + +// --------------------------------------------------------------------------- +// Column allocator (mirroring Java's MapSharedShreddingColumnAllocator) +// --------------------------------------------------------------------------- + +/// Per-row physical column allocation for one row. +struct RowAllocation { + /// `col_to_field[i]` = field id stored in physical column `i`, -1 for empty. + col_to_field: Vec<i32>, + /// Field ids stored in the overflow map, in row order. + overflow_fields: Vec<i32>, +} + +/// Per-row physical column allocator for one shared-shredding MAP column. +/// +/// Assigns fields to physical columns by row order, mirroring the (temporary) +/// simple Java implementation. +struct ColumnAllocator { + num_columns: usize, + field_to_columns: BTreeMap<i32, BTreeSet<usize>>, + overflow_field_set: BTreeSet<i32>, + max_row_width: usize, +} + +impl ColumnAllocator { + fn new(num_columns: usize) -> Self { + Self { + num_columns, + field_to_columns: BTreeMap::new(), + overflow_field_set: BTreeSet::new(), + max_row_width: 0, + } + } + + fn allocate_row(&mut self, field_ids: &[i32]) -> RowAllocation { + self.max_row_width = self.max_row_width.max(field_ids.len()); + + let mut col_to_field = vec![-1; self.num_columns]; + let assign_limit = field_ids.len().min(self.num_columns); + for (i, &field_id) in field_ids.iter().take(assign_limit).enumerate() { + col_to_field[i] = field_id; + self.field_to_columns.entry(field_id).or_default().insert(i); + } + + let mut overflow_fields = Vec::new(); + for &field_id in field_ids.iter().skip(assign_limit) { + overflow_fields.push(field_id); + self.overflow_field_set.insert(field_id); + } + + RowAllocation { + col_to_field, + overflow_fields, + } + } +} + +// --------------------------------------------------------------------------- +// Field metadata (mirroring Java's MapSharedShreddingFieldMeta) +// --------------------------------------------------------------------------- + +/// File-level shredding metadata for one MAP column. +#[derive(Debug, PartialEq, Eq)] +pub(crate) struct MapSharedShreddingFieldMeta { + /// Field name -> field id, sorted by name. + name_to_id: BTreeMap<String, i32>, + /// Field id -> sorted physical column indices. + field_to_columns: BTreeMap<i32, Vec<usize>>, + /// Field ids that ever went to the overflow map, sorted. + overflow_set: BTreeSet<i32>, + num_columns: usize, + max_row_width: usize, +} + +// --------------------------------------------------------------------------- +// Physical schema (mirroring Java's MapSharedShreddingUtils.build*StructType) +// --------------------------------------------------------------------------- + +/// Build the physical `ROW<__field_mapping, __col_0.., [__overflow]>` type for +/// a MAP value type. The writer always includes the overflow column; readers +/// include it only when the overflow set is non-empty (mirroring Java's +/// `buildPhysicalStructType` / `buildSpecificPhysicalStructType`). +fn build_physical_struct_type( + value_type: &DataType, + num_columns: usize, + include_overflow: bool, +) -> DataType { + let mut fields = Vec::with_capacity(num_columns + 2); + fields.push(DataField::new( + 0, + FIELD_MAPPING_NAME.to_string(), + DataType::Array(ArrayType::new(DataType::Int(IntType::new()))), + )); + for i in 0..num_columns { + fields.push(DataField::new( + (i + 1) as i32, + physical_column_name(i), + value_type.clone(), Review Comment: Fixed in `0e3d0eb`. I re-ran the original sparse-row reproduction, the MAP shredding test suite, and an additional end-to-end Parquet roundtrip with `MAP<STRING, BIGINT NOT NULL>`; all pass. -- 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]
