dwsmith1983 commented on code in PR #5654:
URL: https://github.com/apache/datafusion-comet/pull/5654#discussion_r4100235146
##########
native/core/src/parquet/parquet_support.rs:
##########
@@ -162,43 +169,330 @@ impl SparkParquetOptions {
/// Spark-compatible cast implementation. Defers to DataFusion's cast where
that is known
/// to be compatible, and returns an error when a not supported and not
DF-compatible cast
-/// is requested.
+/// is requested. Resolves the nested field mapping for this one value; a
per-file caller
+/// resolves once and uses [`spark_parquet_convert_with_mapping`] for every
batch.
pub fn spark_parquet_convert(
arg: ColumnarValue,
data_type: &DataType,
parquet_options: &SparkParquetOptions,
+) -> DataFusionResult<ColumnarValue> {
+ let mapping =
+ resolve_field_mapping(&arg.data_type(), data_type,
parquet_options).map_err(spark_error)?;
+ spark_parquet_convert_with_mapping(arg, data_type, &mapping,
parquet_options)
+}
+
+/// [`spark_parquet_convert`] with a mapping already resolved for the value's
type.
+pub(crate) fn spark_parquet_convert_with_mapping(
+ arg: ColumnarValue,
+ data_type: &DataType,
+ mapping: &FieldMapping,
+ parquet_options: &SparkParquetOptions,
) -> DataFusionResult<ColumnarValue> {
match arg {
- ColumnarValue::Array(array) =>
Ok(ColumnarValue::Array(parquet_convert_array(
+ ColumnarValue::Array(array) => Ok(ColumnarValue::Array(convert_array(
array,
data_type,
+ mapping,
parquet_options,
+ None,
)?)),
ColumnarValue::Scalar(scalar) => {
// Note that normally CAST(scalar) should be fold in Spark JVM
side. However, for
// some cases e.g., scalar subquery, Spark will not fold it, so we
need to handle it
// here.
let array = scalar.to_array()?;
let scalar = ScalarValue::try_from_array(
- &parquet_convert_array(array, data_type, parquet_options)?,
+ &convert_array(array, data_type, mapping, parquet_options,
None)?,
0,
)?;
Ok(ColumnarValue::Scalar(scalar))
}
}
}
-fn parquet_convert_array(
- array: ArrayRef,
+/// Wrap a [`SparkError`] the way every native operator surfaces it to the JVM.
+pub(crate) fn spark_error(error: SparkError) -> DataFusionError {
+ DataFusionError::External(Box::new(error))
+}
+
+/// Outcome of matching one requested id or name against a struct's file
fields: the last
+/// file field that matched and whether more than one did. A plain `Copy`
value, so resolving
+/// a wide struct allocates nothing per id or per name; the matched names are
only gathered
+/// when an ambiguity is reported.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub(crate) struct FieldMatch {
+ pub(crate) index: usize,
+ pub(crate) ambiguous: bool,
+}
+
+impl FieldMatch {
+ pub(crate) fn new(index: usize, ambiguous: bool) -> Self {
+ Self { index, ambiguous }
+ }
+
+ /// The first file field carrying this id or name.
+ pub(crate) fn first(index: usize) -> Self {
+ Self::new(index, false)
+ }
+
+ /// A further file field carrying the same id or name: the later index
wins, as Spark's
+ /// `toMap` does for exact names, and the entry turns ambiguous.
+ pub(crate) fn also(self, index: usize) -> Self {
+ Self::new(index, true)
+ }
+}
+
+/// Record file field `index` under `key`, keeping the entry `Copy`-sized
however many fields
+/// share the key.
+pub(crate) fn record_field_match<K: Hash + Eq>(
+ matches: &mut HashMap<K, FieldMatch>,
+ key: K,
+ index: usize,
+) {
+ matches
+ .entry(key)
+ .and_modify(|m| *m = m.also(index))
+ .or_insert_with(|| FieldMatch::first(index));
+}
+
+/// Names of the fields carrying `id`, for the duplicate-id error message.
Bracketed and
+/// comma-joined the way Spark's `matchIdField` renders the list, so the
message reads
+/// `Found duplicate field(s) "1": [x, y] in id mapping mode` on both sides.
+pub(crate) fn field_names_with_id(fields: &Fields, id: i32) -> String {
+ let names = fields
+ .iter()
+ .filter(|f| field_id(f) == Some(id))
+ .map(|f| f.name().as_str())
+ .collect::<Vec<_>>()
+ .join(", ");
+ format!("[{names}]")
+}
+
+/// Which file field supplies each requested field, resolved once per file and
reused for
+/// every batch. Follows the requested type as Spark's `clipParquetSchema`
does: a struct
+/// lists one source per requested field, a list in any Arrow representation
or a map carries
+/// the mapping of its element or key and value types, and anything else is a
leaf converted
+/// by type.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) enum FieldMapping {
+ Struct(Vec<StructFieldSource>),
+ List(Box<FieldMapping>),
+ Map(Box<FieldMapping>, Box<FieldMapping>),
+ Leaf,
+}
+
+/// The file field behind one requested struct field.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub(crate) struct StructFieldSource {
+ /// Index of the file field supplying the requested field; `None`
null-fills it.
+ pub(crate) from_index: Option<usize>,
+ /// Mapping of the requested field's own type.
+ pub(crate) nested: FieldMapping,
+}
+
+impl FieldMapping {
+ /// Mapping of a list's element type. A `Leaf` list converts its elements
by type alone,
+ /// as the adapter hands one to every column whose type holds no struct.
+ pub(crate) fn list_element(&self) -> DataFusionResult<&FieldMapping> {
+ match self {
+ FieldMapping::List(inner) => Ok(inner),
+ FieldMapping::Leaf => Ok(&FieldMapping::Leaf),
+ other => Err(DataFusionError::Internal(format!(
+ "list column resolved to a non-list field mapping: {other:?}"
+ ))),
+ }
+ }
+
+ /// Mappings of a map's key and value types; see
[`FieldMapping::list_element`].
+ pub(crate) fn map_entries(&self) -> DataFusionResult<(&FieldMapping,
&FieldMapping)> {
+ match self {
+ FieldMapping::Map(key, value) => Ok((key, value)),
+ FieldMapping::Leaf => Ok((&FieldMapping::Leaf,
&FieldMapping::Leaf)),
+ other => Err(DataFusionError::Internal(format!(
+ "map column resolved to a non-map field mapping: {other:?}"
+ ))),
+ }
+ }
+
+ /// True when every requested field reads the file field at its own
position, so a
+ /// metadata-only relabel of the file array already yields the requested
layout.
+ pub(crate) fn is_positional(&self) -> bool {
+ match self {
+ FieldMapping::Struct(sources) => sources
+ .iter()
+ .enumerate()
+ .all(|(i, s)| s.from_index == Some(i) &&
s.nested.is_positional()),
+ FieldMapping::List(inner) => inner.is_positional(),
+ FieldMapping::Map(key, value) => key.is_positional() &&
value.is_positional(),
+ FieldMapping::Leaf => true,
+ }
+ }
+}
+
+/// The element field of a list in any Arrow representation. One place decides
which types
+/// are lists, so the mapping resolver, the struct-holding walk in the schema
adapter, and
+/// [`convert_array`] agree.
+pub(crate) fn list_element_field(data_type: &DataType) -> Option<&FieldRef> {
+ match data_type {
+ DataType::List(f)
+ | DataType::LargeList(f)
+ | DataType::FixedSizeList(f, _)
+ | DataType::ListView(f)
+ | DataType::LargeListView(f) => Some(f),
+ _ => None,
+ }
+}
+
+/// Resolve how `to_type` reads from `from_type`, recursing through struct,
list, and map
+/// types. Raises the ambiguity Spark reports from `clipParquetGroupFields`
when a requested
+/// id or case-insensitive name matches more than one file field at any level.
+pub(crate) fn resolve_field_mapping(
+ from_type: &DataType,
to_type: &DataType,
parquet_options: &SparkParquetOptions,
-) -> DataFusionResult<ArrayRef> {
- parquet_convert_array_impl(array, to_type, parquet_options, None)
+) -> Result<FieldMapping, SparkError> {
+ use DataType::*;
+ // Dictionary encoding is a physical detail: resolve against the value
type it wraps.
+ // Parquet dictionary encoding only wraps a leaf type, so this mirrors the
adapter's
+ // conversion check and never changes which mapping is built.
+ if let Dictionary(_, value_type) = from_type {
+ return resolve_field_mapping(value_type, to_type, parquet_options);
+ }
+ // The element mapping is the same whichever list representation either
side uses.
+ if let (Some(from_item), Some(to_item)) =
+ (list_element_field(from_type), list_element_field(to_type))
+ {
+ return Ok(FieldMapping::List(Box::new(resolve_field_mapping(
+ from_item.data_type(),
+ to_item.data_type(),
+ parquet_options,
+ )?)));
+ }
+ match (from_type, to_type) {
+ (Struct(from_fields), Struct(to_fields)) => {
+ resolve_struct_mapping(from_fields, to_fields, parquet_options)
+ }
+ (Map(from_entries, from_ordered), Map(to_entries, to_ordered))
+ if from_ordered == to_ordered =>
+ {
+ match (from_entries.data_type(), to_entries.data_type()) {
+ (Struct(from_kv), Struct(to_kv)) if from_kv.len() == 2 &&
to_kv.len() == 2 => {
+ let key = resolve_field_mapping(
+ from_kv[0].data_type(),
+ to_kv[0].data_type(),
+ parquet_options,
+ )?;
+ let value = resolve_field_mapping(
+ from_kv[1].data_type(),
+ to_kv[1].data_type(),
+ parquet_options,
+ )?;
+ Ok(FieldMapping::Map(Box::new(key), Box::new(value)))
+ }
+ _ => Ok(FieldMapping::Leaf),
+ }
+ }
+ _ => Ok(FieldMapping::Leaf),
+ }
}
-fn parquet_convert_array_impl(
+/// Match `to` (requested) struct fields to `from` (file) fields. Mirrors
Spark's
+/// `clipParquetGroupFields`: when the requested struct carries Parquet field
ids anywhere,
+/// id-bearing requested fields match only by id and the rest by name;
otherwise every field
+/// matches by name.
+fn resolve_struct_mapping(
+ from_fields: &Fields,
+ to_fields: &Fields,
+ parquet_options: &SparkParquetOptions,
+) -> Result<FieldMapping, SparkError> {
+ let should_match_by_id =
+ parquet_options.use_field_id && to_fields.iter().any(|f|
field_id(f).is_some());
+
+ let mut id_matches: HashMap<i32, FieldMatch> = HashMap::new();
+ if should_match_by_id {
+ for (i, field) in from_fields.iter().enumerate() {
+ if let Some(id) = field_id(field) {
+ record_field_match(&mut id_matches, id, i);
+ }
+ }
+ }
+
+ // Fold the file and requested names once via the same
`toLowerCase(Locale.ROOT)` the
+ // top-level schema adapter uses, so nested case-insensitive matching
agrees with it.
+ let mut all_names: Vec<&str> = Vec::with_capacity(from_fields.len() +
to_fields.len());
+ all_names.extend(from_fields.iter().map(|f| f.name().as_str()));
+ all_names.extend(to_fields.iter().map(|f| f.name().as_str()));
+ let all_folded = fold_names(&all_names, parquet_options.case_sensitive)
+ .map_err(|e| SparkError::Internal(e.to_string()))?;
+ let (from_folded, to_folded) = all_folded.split_at(from_fields.len());
+
+ let mut name_matches: HashMap<&str, FieldMatch> = HashMap::new();
+ for (i, folded) in from_folded.iter().enumerate() {
+ record_field_match(&mut name_matches, folded.as_str(), i);
+ }
+
+ let mut sources = Vec::with_capacity(to_fields.len());
+ for (to_pos, to_field) in to_fields.iter().enumerate() {
+ let from_index = match (should_match_by_id, field_id(to_field)) {
+ // A missing id match is a missing column, never a name match.
+ (true, Some(id)) => match id_matches.get(&id) {
+ Some(m) if m.ambiguous => {
+ return Err(SparkError::DuplicateFieldByFieldId {
+ required_id: id,
+ matched_fields: field_names_with_id(from_fields, id),
+ });
+ }
+ Some(m) => Some(m.index),
+ None => None,
+ },
+ _ => match name_matches.get(to_folded[to_pos].as_str()) {
+ // Spark's `matchCaseInsensitiveField` raises
`_LEGACY_ERROR_TEMP_2093` for a
+ // requested name that folds onto more than one file field,
whether the siblings
+ // differ by case or are byte-identical. In case-sensitive
mode the fold is
+ // identity, so a collision means byte-identical siblings.
Spark's
+ // `matchCaseSensitiveField` builds its map with `toMap` there
and the last field
+ // wins silently. Comet refuses instead of picking one, with
the error every
+ // other path raises for a duplicate the decoder cannot
represent.
+ Some(m) if m.ambiguous => {
+ let matched: Vec<&str> = from_folded
+ .iter()
+ .zip(from_fields.iter())
+ .filter(|(folded, _)| *folded == &to_folded[to_pos])
+ .map(|(_, f)| f.name().as_str())
+ .collect();
+ if parquet_options.case_sensitive {
+ return
Err(SparkError::Internal(duplicate_parquet_field_message(
+ to_field.name(),
+ )));
+ }
+ return Err(SparkError::duplicate_field_case_insensitive(
+ to_field.name(),
+ &matched,
+ ));
+ }
+ Some(m) => Some(m.index),
+ None => None,
+ },
+ };
+ let nested = match from_index {
+ Some(i) => resolve_field_mapping(
+ from_fields[i].data_type(),
+ to_field.data_type(),
+ parquet_options,
+ )?,
+ None => FieldMapping::Leaf,
+ };
+ sources.push(StructFieldSource { from_index, nested });
+ }
+ Ok(FieldMapping::Struct(sources))
+}
+
+/// Convert `array` to `to_type` through its resolved `mapping`.
`parent_nulls` masks the rows
+/// hidden beneath null ancestors, so only values Spark reads are checked for
overflow.
+fn convert_array(
Review Comment:
Dropped from this PR. `parquet_support.rs` is main plus the duplicate-id arm
in `match_struct_fields` and `field_names_with_id`. The once-per-file mapping,
the renames, `convert_list`, `list_element_field` and the `has_timestamp_unit`
change sit on a side branch for a follow-up PR, which will carry a case in
`native/core/benches/parquet_timestamp_conversion.rs` so the per-batch saving
is measured.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
