andygrove commented on code in PR #5806: URL: https://github.com/apache/datafusion-comet/pull/5806#discussion_r3981709406
########## native/spark-expr/src/map_funcs/map_extract.rs: ########## @@ -0,0 +1,632 @@ +// 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. + +use arrow::array::{ + new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray, NullBufferBuilder, Scalar, + UInt32Array, +}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::cmp::eq; +use arrow::compute::take; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::utils::take_function_args; +use datafusion::common::{exec_err, Result as DataFusionResult}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`. +/// +/// Overrides DataFusion's `map_extract` under the same name, and differs from it in two ways: +/// +/// - it returns the matched **value** rather than a one-element list, so the planner does not +/// have to unwrap the list with a second `ListExtract` pass (see `planner.rs`); +/// - the lookup is vectorized. DataFusion's `general_map_extract_inner` re-slices the query key +/// and every candidate key into a fresh `ArrayRef` per comparison and compares them through +/// `dyn Array` equality, which made a constant-key lookup roughly 35x more expensive than any +/// other Comet map kernel and slower than Spark itself +/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)). Here a single Arrow +/// `eq` covers the whole batch of entries at once, the per-row work is a bit scan over the +/// resulting mask, and the values are gathered with one `take`. +/// +/// Spark's own lookup returns the first entry whose key compares equal, so the mask scan stops at +/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup key all produce +/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map overload. +/// +/// Key types whose Spark equality this cannot reproduce (floating point, non-default collations, +/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala` declines them so the +/// expression falls back to Spark. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkMapExtract { + signature: Signature, +} + +impl Default for SparkMapExtract { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapExtract { + pub fn new() -> Self { + Self { + // `user_defined` so `coerce_types` runs and casts the lookup key to the map's key + // type; Comet's planner applies that coercion to the argument expression. + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkMapExtract { + fn name(&self) -> &str { + "map_extract" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(map_entry_fields(map_type)?.1.data_type().clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult<Vec<DataType>> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(vec![ + map_type.clone(), + map_entry_fields(map_type)?.0.data_type().clone(), + ]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> { + let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?; + spark_map_extract(map_arg, key_arg, args.number_rows) + } +} + +/// The `(key, value)` fields of a `Map`'s entry struct. +fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef, &FieldRef)> { + match map_type { + DataType::Map(entries, _) => match entries.data_type() { + DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0], &fields[1])), + other => exec_err!("map_extract: map entries must be a two-field struct, got {other}"), + }, + other => exec_err!("map_extract: the first argument must be a map, got {other}"), + } +} + +/// Look up `key_arg` in each row of `map_arg`, returning the matched value or `NULL`. +pub fn spark_map_extract( + map_arg: &ColumnarValue, + key_arg: &ColumnarValue, + number_rows: usize, +) -> DataFusionResult<ColumnarValue> { + let map_ref: ArrayRef = match map_arg { + ColumnarValue::Array(array) => Arc::clone(array), + ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?, + }; + let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else { + return exec_err!( + "map_extract: the first argument must be a map, got {}", + map_ref.data_type() + ); + }; + + let num_rows = map_array.len(); + let value_type = map_array.value_type(); + + // Arrow keeps a sliced `MapArray`'s entries child intact and slices only the offsets, so the + // offsets index the *unsliced* keys/values and the visible entries are the half-open range + // [entries_start, entries_end). Comparing only that window keeps a native OFFSET from paying + // for the entries it skipped. + let offsets = map_array.offsets(); + let entries_start = offsets[0] as usize; + let entries_end = offsets[num_rows] as usize; + if entries_start == entries_end { Review Comment: Fixed. Both checks moved into a `validate_lookup_key` helper that runs before the empty-window return, so the lookup key type and, for an array key, its length are now checked on every batch regardless of what the maps hold. `argument_checks_do_not_depend_on_the_data` pins it with an all-empty/all-NULL map: it now errors on a mismatched key type and on a short key array where it previously answered NULLs. ########## native/spark-expr/src/map_funcs/map_extract.rs: ########## @@ -0,0 +1,632 @@ +// 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. + +use arrow::array::{ + new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray, NullBufferBuilder, Scalar, + UInt32Array, +}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::cmp::eq; +use arrow::compute::take; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::utils::take_function_args; +use datafusion::common::{exec_err, Result as DataFusionResult}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`. +/// +/// Overrides DataFusion's `map_extract` under the same name, and differs from it in two ways: +/// +/// - it returns the matched **value** rather than a one-element list, so the planner does not +/// have to unwrap the list with a second `ListExtract` pass (see `planner.rs`); +/// - the lookup is vectorized. DataFusion's `general_map_extract_inner` re-slices the query key +/// and every candidate key into a fresh `ArrayRef` per comparison and compares them through +/// `dyn Array` equality, which made a constant-key lookup roughly 35x more expensive than any +/// other Comet map kernel and slower than Spark itself +/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)). Here a single Arrow +/// `eq` covers the whole batch of entries at once, the per-row work is a bit scan over the +/// resulting mask, and the values are gathered with one `take`. +/// +/// Spark's own lookup returns the first entry whose key compares equal, so the mask scan stops at +/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup key all produce +/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map overload. +/// +/// Key types whose Spark equality this cannot reproduce (floating point, non-default collations, +/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala` declines them so the +/// expression falls back to Spark. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkMapExtract { + signature: Signature, +} + +impl Default for SparkMapExtract { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapExtract { + pub fn new() -> Self { + Self { + // `user_defined` so `coerce_types` runs and casts the lookup key to the map's key + // type; Comet's planner applies that coercion to the argument expression. + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkMapExtract { + fn name(&self) -> &str { + "map_extract" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(map_entry_fields(map_type)?.1.data_type().clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult<Vec<DataType>> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(vec![ + map_type.clone(), + map_entry_fields(map_type)?.0.data_type().clone(), + ]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> { + let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?; + spark_map_extract(map_arg, key_arg, args.number_rows) + } +} + +/// The `(key, value)` fields of a `Map`'s entry struct. +fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef, &FieldRef)> { + match map_type { + DataType::Map(entries, _) => match entries.data_type() { + DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0], &fields[1])), + other => exec_err!("map_extract: map entries must be a two-field struct, got {other}"), + }, + other => exec_err!("map_extract: the first argument must be a map, got {other}"), + } +} + +/// Look up `key_arg` in each row of `map_arg`, returning the matched value or `NULL`. +pub fn spark_map_extract( + map_arg: &ColumnarValue, + key_arg: &ColumnarValue, + number_rows: usize, +) -> DataFusionResult<ColumnarValue> { + let map_ref: ArrayRef = match map_arg { + ColumnarValue::Array(array) => Arc::clone(array), + ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?, + }; + let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else { + return exec_err!( + "map_extract: the first argument must be a map, got {}", + map_ref.data_type() + ); + }; + + let num_rows = map_array.len(); + let value_type = map_array.value_type(); + + // Arrow keeps a sliced `MapArray`'s entries child intact and slices only the offsets, so the + // offsets index the *unsliced* keys/values and the visible entries are the half-open range + // [entries_start, entries_end). Comparing only that window keeps a native OFFSET from paying + // for the entries it skipped. + let offsets = map_array.offsets(); + let entries_start = offsets[0] as usize; + let entries_end = offsets[num_rows] as usize; + if entries_start == entries_end { + // Every row is empty or NULL, so nothing can match. + return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows))); + } + let window_len = entries_end - entries_start; + let keys = map_array.keys().slice(entries_start, window_len); + + let matched = match key_arg { + ColumnarValue::Scalar(scalar) => { + if scalar.is_null() { + // Spark map keys are never NULL, so a NULL lookup key matches nothing. + return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows))); + } + let key = scalar.to_array_of_size(1)?; + key_match_mask(&keys, &key, true)? + } + ColumnarValue::Array(key_array) => { + if key_array.len() != num_rows { + return exec_err!( + "map_extract: expected {num_rows} lookup keys, got {}", + key_array.len() + ); + } + // One vectorized compare needs a lookup key per *entry*, not per row, so gather each + // row's key across that row's entries. Entries in a gap between two rows (offsets are + // only required to be monotonic) keep index 0; the per-row scan below never reads + // those positions. + let mut gather = vec![0u32; window_len]; + for row in 0..num_rows { + let start = offsets[row] as usize - entries_start; + let end = offsets[row + 1] as usize - entries_start; + gather[start..end].fill(row as u32); + } + let per_entry_key = take(key_array, &UInt32Array::from(gather), None)?; + key_match_mask(&keys, &per_entry_key, false)? + } + }; + + // Gather the first matching entry of each row. Map offsets are `i32`, so an entry index always + // fits in `u32`. + // + // A NULL map row reads NULL whatever its entries hold. Arrow does not require a null row's + // offset range to be empty, and neither Comet's struct-field helper (which adds a parent null + // mask while preserving the child buffers) nor the UDF execution layer clears those entries, so + // a null row can carry a live `a -> 7` that would otherwise match. Spark returns NULL for a + // NULL map under both ANSI modes, for `element_at` and for `GetMapValue` alike, and only + // `element_at` has a nullable-input guard upstream of this kernel. + let map_nulls = map_array.nulls(); + let mut indices = vec![0u32; num_rows]; + let mut nulls = NullBufferBuilder::new(num_rows); + for row in 0..num_rows { + if map_nulls.is_some_and(|n| n.is_null(row)) { + nulls.append(false); + continue; + } + let start = offsets[row] as usize - entries_start; + let end = offsets[row + 1] as usize - entries_start; + let found = (start..end).find(|&i| matched.value(i)); + if let Some(i) = found { + indices[row] = (i + entries_start) as u32; + } + nulls.append(found.is_some()); + } + let indices = UInt32Array::new(indices.into(), nulls.finish()); + + Ok(ColumnarValue::Array(take( + map_array.values(), + &indices, + None, + )?)) +} + +/// A bit per map entry: set where the stored key equals the lookup key. `lookup` is either a +/// length-1 array broadcast over every entry (constant key) or one key per entry. +fn key_match_mask( + keys: &ArrayRef, + lookup: &ArrayRef, + lookup_is_scalar: bool, +) -> DataFusionResult<BooleanBuffer> { + // The planner casts the lookup key to the map's declared key type, so a mismatch here means + // the runtime encoding is not the declared one (a dictionary-encoded key column, say). Reject + // it rather than comparing incomparable encodings and reporting every row as a miss. + if keys.data_type() != lookup.data_type() { + return exec_err!( + "map_extract: lookup key type {} does not match the map key type {}", + lookup.data_type(), + keys.data_type() + ); + } + let compared = if lookup_is_scalar { + eq(keys, &Scalar::new(Arc::clone(lookup))) + } else { + eq(keys, lookup) + }; + match compared { + Ok(mask) => { + // A NULL on either side compares as NULL, which is not a match. + let (values, nulls) = mask.into_parts(); + Ok(match nulls { + Some(nulls) if nulls.null_count() > 0 => &values & nulls.inner(), + _ => values, + }) + } + // `eq` rejects nested key types. `MapKeySupport` declines those before they reach the + // native lookup, but keep DataFusion's element-wise comparison as a backstop so this + // kernel is never less capable than the one it replaces. + Err(_) => Ok(elementwise_match_mask(keys, lookup, lookup_is_scalar)), Review Comment: Taken. It dispatches on `keys.data_type().is_nested()` now, which is exactly the predicate `compare_op` applies after unwrapping one dictionary level, and with the length and type checked up front nesting was the only remaining thing `eq` rejects. Both `eq` calls use `?`, so a genuine failure surfaces instead of silently taking the per-row path. ########## native/spark-expr/src/map_funcs/mod.rs: ########## @@ -15,5 +15,7 @@ // specific language governing permissions and limitations // under the License. +mod map_extract; mod map_sort; +pub use map_extract::{spark_map_extract, SparkMapExtract}; Review Comment: Done. `spark_map_extract` is private again; only `SparkMapExtract` is re-exported from `map_funcs` and `lib.rs`. ########## native/spark-expr/src/comet_scalar_funcs.rs: ########## @@ -321,6 +321,9 @@ fn all_scalar_functions() -> Vec<Arc<ScalarUDF>> { )), Arc::new(ScalarUDF::new_from_impl(SparkMakeDate::default())), Arc::new(ScalarUDF::new_from_impl(SparkMakeTime::default())), + // Overrides datafusion-functions-nested' `map_extract` with a vectorized lookup that + // returns the value itself rather than a one-element list (#5795). + Arc::new(ScalarUDF::new_from_impl(SparkMapExtract::default())), Review Comment: Added the alias rather than documenting the gap, since it is three lines and leaves nothing to rediscover. `SparkMapExtract::aliases()` returns `["element_at"]`, so the override now replaces both registry entries. I re-grepped and nothing emits that name today, so this is consistency rather than a fix, and there is a one-line test so it does not quietly come undone. ########## native/spark-expr/src/map_funcs/map_extract.rs: ########## @@ -0,0 +1,632 @@ +// 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. + +use arrow::array::{ + new_null_array, Array, ArrayRef, BooleanBufferBuilder, MapArray, NullBufferBuilder, Scalar, + UInt32Array, +}; +use arrow::buffer::BooleanBuffer; +use arrow::compute::kernels::cmp::eq; +use arrow::compute::take; +use arrow::datatypes::{DataType, FieldRef}; +use datafusion::common::utils::take_function_args; +use datafusion::common::{exec_err, Result as DataFusionResult}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use std::sync::Arc; + +/// Spark's map lookup: `GetMapValue` (`m[k]`) and `element_at(<map>, k)`. +/// +/// Overrides DataFusion's `map_extract` under the same name, and differs from it in two ways: +/// +/// - it returns the matched **value** rather than a one-element list, so the planner does not +/// have to unwrap the list with a second `ListExtract` pass (see `planner.rs`); +/// - the lookup is vectorized. DataFusion's `general_map_extract_inner` re-slices the query key +/// and every candidate key into a fresh `ArrayRef` per comparison and compares them through +/// `dyn Array` equality, which made a constant-key lookup roughly 35x more expensive than any +/// other Comet map kernel and slower than Spark itself +/// ([#5795](https://github.com/apache/datafusion-comet/issues/5795)). Here a single Arrow +/// `eq` covers the whole batch of entries at once, the per-row work is a bit scan over the +/// resulting mask, and the values are gathered with one `take`. +/// +/// Spark's own lookup returns the first entry whose key compares equal, so the mask scan stops at +/// the first match too. A missing key, a `NULL` map row, and a `NULL` lookup key all produce +/// `NULL`, matching `GetMapValueUtil.getValueEval` and `ElementAt`'s map overload. +/// +/// Key types whose Spark equality this cannot reproduce (floating point, non-default collations, +/// complex keys) never reach here: `MapKeySupport` in `serde/maps.scala` declines them so the +/// expression falls back to Spark. +#[derive(Debug, Hash, Eq, PartialEq)] +pub struct SparkMapExtract { + signature: Signature, +} + +impl Default for SparkMapExtract { + fn default() -> Self { + Self::new() + } +} + +impl SparkMapExtract { + pub fn new() -> Self { + Self { + // `user_defined` so `coerce_types` runs and casts the lookup key to the map's key + // type; Comet's planner applies that coercion to the argument expression. + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkMapExtract { + fn name(&self) -> &str { + "map_extract" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> DataFusionResult<DataType> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(map_entry_fields(map_type)?.1.data_type().clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> DataFusionResult<Vec<DataType>> { + let [map_type, _] = take_function_args(self.name(), arg_types)?; + Ok(vec![ + map_type.clone(), + map_entry_fields(map_type)?.0.data_type().clone(), + ]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> DataFusionResult<ColumnarValue> { + let [map_arg, key_arg] = take_function_args(self.name(), &args.args)?; + spark_map_extract(map_arg, key_arg, args.number_rows) + } +} + +/// The `(key, value)` fields of a `Map`'s entry struct. +fn map_entry_fields(map_type: &DataType) -> DataFusionResult<(&FieldRef, &FieldRef)> { + match map_type { + DataType::Map(entries, _) => match entries.data_type() { + DataType::Struct(fields) if fields.len() == 2 => Ok((&fields[0], &fields[1])), + other => exec_err!("map_extract: map entries must be a two-field struct, got {other}"), + }, + other => exec_err!("map_extract: the first argument must be a map, got {other}"), + } +} + +/// Look up `key_arg` in each row of `map_arg`, returning the matched value or `NULL`. +pub fn spark_map_extract( + map_arg: &ColumnarValue, + key_arg: &ColumnarValue, + number_rows: usize, +) -> DataFusionResult<ColumnarValue> { + let map_ref: ArrayRef = match map_arg { + ColumnarValue::Array(array) => Arc::clone(array), + ColumnarValue::Scalar(scalar) => scalar.to_array_of_size(number_rows)?, + }; + let Some(map_array) = map_ref.as_any().downcast_ref::<MapArray>() else { + return exec_err!( + "map_extract: the first argument must be a map, got {}", + map_ref.data_type() + ); + }; + + let num_rows = map_array.len(); + let value_type = map_array.value_type(); + + // Arrow keeps a sliced `MapArray`'s entries child intact and slices only the offsets, so the + // offsets index the *unsliced* keys/values and the visible entries are the half-open range + // [entries_start, entries_end). Comparing only that window keeps a native OFFSET from paying + // for the entries it skipped. + let offsets = map_array.offsets(); + let entries_start = offsets[0] as usize; + let entries_end = offsets[num_rows] as usize; + if entries_start == entries_end { + // Every row is empty or NULL, so nothing can match. + return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows))); + } + let window_len = entries_end - entries_start; + let keys = map_array.keys().slice(entries_start, window_len); + + let matched = match key_arg { + ColumnarValue::Scalar(scalar) => { + if scalar.is_null() { + // Spark map keys are never NULL, so a NULL lookup key matches nothing. + return Ok(ColumnarValue::Array(new_null_array(value_type, num_rows))); + } + let key = scalar.to_array_of_size(1)?; + key_match_mask(&keys, &key, true)? + } + ColumnarValue::Array(key_array) => { + if key_array.len() != num_rows { + return exec_err!( + "map_extract: expected {num_rows} lookup keys, got {}", + key_array.len() + ); + } + // One vectorized compare needs a lookup key per *entry*, not per row, so gather each + // row's key across that row's entries. Entries in a gap between two rows (offsets are + // only required to be monotonic) keep index 0; the per-row scan below never reads + // those positions. + let mut gather = vec![0u32; window_len]; + for row in 0..num_rows { + let start = offsets[row] as usize - entries_start; + let end = offsets[row + 1] as usize - entries_start; + gather[start..end].fill(row as u32); + } + let per_entry_key = take(key_array, &UInt32Array::from(gather), None)?; + key_match_mask(&keys, &per_entry_key, false)? + } + }; + + // Gather the first matching entry of each row. Map offsets are `i32`, so an entry index always + // fits in `u32`. + // + // A NULL map row reads NULL whatever its entries hold. Arrow does not require a null row's + // offset range to be empty, and neither Comet's struct-field helper (which adds a parent null + // mask while preserving the child buffers) nor the UDF execution layer clears those entries, so + // a null row can carry a live `a -> 7` that would otherwise match. Spark returns NULL for a + // NULL map under both ANSI modes, for `element_at` and for `GetMapValue` alike, and only + // `element_at` has a nullable-input guard upstream of this kernel. + let map_nulls = map_array.nulls(); Review Comment: I could not find a producing path either, so I softened the wording rather than claim one. The comment now says every producer traced -- the Parquet readers, Spark`s `ArrowWriter`, and arrow`s own `filter` / `take` / `concat` -- gives a null row an empty range, and that this is a representation the format permits rather than one known to arrive here, so the guard exists to pin the contract at the kernel boundary. The test fixture`s doc comment had the same phrasing and got the same treatment. -- 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]
