andygrove commented on code in PR #98:
URL: https://github.com/apache/datafusion-spark/pull/98#discussion_r4565486093
##########
crates/spark/src/function/map/utils.rs:
##########
@@ -203,50 +187,52 @@ fn map_deduplicate_keys(
"map_deduplicate_keys: keys and values lists in the same
row must have equal lengths"
);
}
- key_to_output_idx.clear();
- for cur_entry_idx in 0..num_keys_entries {
- let key = ScalarValue::try_from_array(
- &flat_keys,
- cur_keys_offset + cur_entry_idx,
- )?
- .compacted();
- let abs_value_idx = (cur_values_offset + cur_entry_idx) as i32;
-
- if let Some(&output_idx) = key_to_output_idx.get(&key) {
- if last_value_wins {
- value_indices[output_idx] = abs_value_idx;
- keys_mask_builder.append_value(false);
- continue;
+ if num_keys_entries != 0 {
+ let mut seen_keys = HashSet::new();
+ let mut keys_mask_one = vec![false; num_keys_entries];
+ let mut values_mask_one = vec![false; num_values_entries];
+ for cur_entry_idx in (0..num_keys_entries).rev() {
Review Comment:
This changes the dedup logic as well as the config, and it changes the
output order. Scanning in reverse and keeping the last-seen position puts a
duplicated key where it *last* appears. Spark's `ArrayBasedMapBuilder` keeps it
where it *first* appears and overwrites the value.
With `LAST_WIN`, `map_from_arrays(array(1,2,1), array('a','b','c'))` now
returns `{2: b, 1: c}`, but Spark returns `{1: c, 2: b}`. The existing unit
test catches this: `cargo test -p datafusion-spark --lib map::` fails
`sliced_input_last_win` on this branch (left `[(2,"b"),(1,"c")]`, right
`[(1,"c"),(2,"b")]`).
Could we keep the dedup implementation from `main` and limit this PR to the
config move?
##########
crates/spark/src/function/map/utils.rs:
##########
@@ -203,50 +187,52 @@ fn map_deduplicate_keys(
"map_deduplicate_keys: keys and values lists in the same
row must have equal lengths"
);
}
- key_to_output_idx.clear();
- for cur_entry_idx in 0..num_keys_entries {
- let key = ScalarValue::try_from_array(
- &flat_keys,
- cur_keys_offset + cur_entry_idx,
- )?
- .compacted();
- let abs_value_idx = (cur_values_offset + cur_entry_idx) as i32;
-
- if let Some(&output_idx) = key_to_output_idx.get(&key) {
- if last_value_wins {
- value_indices[output_idx] = abs_value_idx;
- keys_mask_builder.append_value(false);
- continue;
+ if num_keys_entries != 0 {
+ let mut seen_keys = HashSet::new();
Review Comment:
Smaller points on this block, which mostly go away if the `main`
implementation is kept:
- `values_mask_one` is always the same as `keys_mask_one`, because the
function already errors when key and value counts differ. So the second mask
and second `filter` repeat the same work.
- Each row allocates a new `HashSet` and two `Vec<bool>`s, and null rows
allocate vectors. `main` reused one map with `clear()` and used
`append_n(false)`.
##########
crates/spark/src/session_state.rs:
##########
@@ -42,6 +43,12 @@ pub trait SessionStateBuilderSpark {
impl SessionStateBuilderSpark for SessionStateBuilder {
fn with_spark_features(mut self) -> Self {
+ self.config()
Review Comment:
Two issues with registering `SparkOptions` here:
1. This always inserts `SparkOptions::default()`, so it replaces a
`SparkOptions` the caller already put in the config. For example,
`with_config(SessionConfig::new().with_option_extension(SparkOptions {
map_key_dedup_policy: LastWin, .. }))` followed by `with_spark_features()`
resets the policy to `EXCEPTION`. Could it insert only when no `SparkOptions`
is present?
2. It changes whatever config the builder holds at this point, so a later
`.with_config(cfg)` drops the extension. `SET spark.map_key_dedup_policy` then
fails, and the map functions quietly use `EXCEPTION`, because `is_some_and`
treats a missing extension as the default. It may be worth documenting that
`with_config` must come first, or registering the extension in a way that
survives it.
##########
crates/sqllogictest/test_files/map/map_key_dedup_policy.slt:
##########
@@ -0,0 +1,16 @@
+# Spark map construction uses EXCEPTION by default.
+statement ok
+set spark.map_key_dedup_policy = 'LAST_WIN';
+
+query ?
+SELECT map_from_arrays(array(1, 2, 1), array('a', 'b', 'c'));
+----
+{2: b, 1: c}
Review Comment:
This expects `{2: b, 1: c}`, which is the order from the reverse scan in
`utils.rs`. Spark returns `{1: c, 2: b}`, and `map_from_arrays.slt:146` expects
`{1: c, 2: b}` for the same query, so the two files disagree.
##########
crates/spark/src/function/map/utils.rs:
##########
@@ -203,50 +187,52 @@ fn map_deduplicate_keys(
"map_deduplicate_keys: keys and values lists in the same
row must have equal lengths"
);
}
- key_to_output_idx.clear();
- for cur_entry_idx in 0..num_keys_entries {
- let key = ScalarValue::try_from_array(
- &flat_keys,
- cur_keys_offset + cur_entry_idx,
- )?
- .compacted();
- let abs_value_idx = (cur_values_offset + cur_entry_idx) as i32;
-
- if let Some(&output_idx) = key_to_output_idx.get(&key) {
- if last_value_wins {
- value_indices[output_idx] = abs_value_idx;
- keys_mask_builder.append_value(false);
- continue;
+ if num_keys_entries != 0 {
+ let mut seen_keys = HashSet::new();
+ let mut keys_mask_one = vec![false; num_keys_entries];
+ let mut values_mask_one = vec![false; num_values_entries];
+ for cur_entry_idx in (0..num_keys_entries).rev() {
+ let key =
+ ScalarValue::try_from_array(&flat_keys,
cur_keys_offset + cur_entry_idx)?
+ .compacted();
+ if seen_keys.contains(&key) {
+ if !last_value_wins {
+ return exec_err!(
+ "[DUPLICATED_MAP_KEY] Duplicate map key {key}
was found, \
+ please check the input data. To allow
duplicate keys with \
+ last-value-wins semantics, set \
+ `spark.map_key_dedup_policy` to `LAST_WIN`."
+ );
+ }
+ } else {
+ keys_mask_one[cur_entry_idx] = true;
+ values_mask_one[cur_entry_idx] = true;
+ seen_keys.insert(key);
+ new_last_offset += 1;
}
- return exec_err!(
- "[DUPLICATED_MAP_KEY] Duplicate map key {key} was
found, \
- please check the input data. To allow duplicate keys
with \
- last-value-wins semantics, set \
- `datafusion.spark.map_key_dedup_policy` to
`LAST_WIN`."
- );
}
- keys_mask_builder.append_value(true);
- key_to_output_idx.insert(key, value_indices.len());
- value_indices.push(abs_value_idx);
- new_last_offset += 1;
+ keys_mask_builder.append_array(&keys_mask_one.into());
+ values_mask_builder.append_array(&values_mask_one.into());
}
} else {
- // The result entry is NULL — no keys/values emitted. Still pad the
- // mask so it stays aligned with the key entries of later rows.
- keys_mask_builder.append_n(num_keys_entries, false);
+ let keys_mask_one = vec![false; num_keys_entries];
+ let values_mask_one = vec![false; num_values_entries];
+ keys_mask_builder.append_array(&keys_mask_one.into());
+ values_mask_builder.append_array(&values_mask_one.into());
}
new_offsets.push(new_last_offset);
cur_keys_offset += num_keys_entries;
cur_values_offset += num_values_entries;
}
let keys_mask = keys_mask_builder.finish();
- // Every key entry of the input rows must have exactly one mask bit;
- // otherwise the mask selects the wrong keys for all later rows.
- debug_assert_eq!(keys_mask.len(), cur_keys_offset - first_keys_offset);
let input_keys = flat_keys.slice(first_keys_offset, keys_mask.len());
let needed_keys = filter(&input_keys, &keys_mask)?;
- let value_indices_array = Int32Array::from(value_indices);
- let needed_values = take(&flat_values, &value_indices_array, None)?;
+ let values_mask = values_mask_builder.finish();
Review Comment:
The `debug_assert` that the keys mask covers every key entry was removed,
along with the comments explaining the first-occurrence design. Could we keep
them? If the mask length ever drifts from the key count, `filter` would
silently pick the wrong keys for later rows.
##########
crates/spark/src/function/map/map_from_arrays.rs:
##########
@@ -69,8 +68,11 @@ impl ScalarUDFImpl for MapFromArrays {
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
- let last_value_wins =
- args.config_options.spark.map_key_dedup_policy ==
MapKeyDedupPolicy::LastWin;
+ let last_value_wins = args
Review Comment:
All three map functions now read only the `SparkOptions` extension. But
`map_from_arrays.slt:141`, `map_from_entries.slt:169` and
`str_to_map.slt:119/146/150-152` still set
`datafusion.spark.map_key_dedup_policy`, which these functions no longer read.
Their `LAST_WIN` queries would then hit `DUPLICATED_MAP_KEY`, so I'd expect
`cargo test -p datafusion-spark-sqllogictest` to fail. Could those files move
to `spark.map_key_dedup_policy`?
Users who set the old key would also lose `LAST_WIN` without an error, which
may be worth a note in the PR description.
##########
crates/spark/src/function/map/map_from_entries.rs:
##########
@@ -89,8 +85,11 @@ impl ScalarUDFImpl for MapFromEntries {
}
fn invoke_with_args(&self, args: ScalarFunctionArgs) ->
Result<ColumnarValue> {
- let last_value_wins =
- args.config_options.spark.map_key_dedup_policy ==
MapKeyDedupPolicy::LastWin;
+ let last_value_wins = args
Review Comment:
The same `SparkOptions` lookup and `LastWin` comparison is repeated in
`map_from_arrays`, `map_from_entries` and `str_to_map`. A helper in
`config.rs`, such as `fn map_last_value_wins(options: &ConfigOptions) -> bool`,
would keep it in one place.
--
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]
