neilconway commented on code in PR #25201:
URL: https://github.com/apache/datafusion/pull/25201#discussion_r4030468227
##########
datafusion/functions/src/utils.rs:
##########
@@ -347,6 +353,180 @@ pub fn decimal64_to_i64(value: i64, scale: i8) ->
Result<i64, ArrowError> {
}
}
+/// Finds, for each row of `map`, the first entry whose key equals that row's
+/// lookup key.
+///
+/// `keys` holds either a single key, which every row is looked up with, or
+/// one key per map row. The result has one element per map row: the index of
+/// the matching entry into `map.values()`, or null when the row is null, the
+/// lookup key is null, or no entry matches. It can be passed directly to
+/// [`arrow::compute::take`] on `map.values()`.
+///
+/// Non-nested keys must have the map's key type, up to dictionary encoding.
+/// Nested keys must have the same structure, and may differ in field names
+/// and nullability. Keys are compared the way `ORDER BY` compares values:
+/// floating point keys use total ordering, so `-0.0` and `0.0` are different
+/// keys and NaN matches NaN.
+pub fn map_lookup(map: &MapArray, keys: &dyn Array) -> Result<UInt32Array> {
+ let map_keys = map.keys();
+ let single_key = match keys.len() {
+ 1 => true,
+ len if len == map.len() => false,
+ len => {
+ return internal_err!(
+ "map_lookup expects one lookup key or one per map row ({}),
got {len}",
+ map.len()
+ );
+ }
+ };
+ let key_type = map_keys.data_type();
+ // A nested lookup key only has to be nested here; `make_comparator`
+ // checks its structure. A non-nested lookup key must have the map's
+ // key type, ignoring dictionary encoding.
+ let compatible = if key_type.is_nested() {
+ keys.data_type().is_nested()
+ } else {
+
strip_dictionary(key_type).equals_datatype(strip_dictionary(keys.data_type()))
+ };
+ if !compatible {
+ return exec_err!(
+ "The key type {} does not match the map key type {}",
+ keys.data_type(),
+ key_type
+ );
+ }
+ // The comparison kernels need both sides to use the same encoding.
+ let cast_keys;
+ let keys: &dyn Array = if key_type.is_nested() || keys.data_type() ==
key_type {
+ keys
+ } else {
+ cast_keys = cast(keys, key_type)?;
+ cast_keys.as_ref()
+ };
+
+ let offsets = map.value_offsets();
+ let (first, last) = (offsets[0] as usize, offsets[map.len()] as usize);
+ // No row has any entries, so nothing can match. Map keys are never
+ // null, so a null lookup key matches nothing either.
+ if first == last || (single_key && keys.logical_null_count() > 0) {
+ return Ok(UInt32Array::new_null(map.len()));
+ }
+ let key_nulls = if single_key {
+ None
+ } else {
+ keys.logical_nulls()
+ };
+ let mut scanner = RowScanner::new(map, key_nulls.as_ref());
+
+ // Scan with a comparator, which stops at the first match in each row.
+ // Count the comparisons over a sample of rows to see whether stopping
+ // early pays off.
+ let cmp = make_comparator(map_keys.as_ref(), keys,
SortOptions::default())?;
+ let compare =
+ |entry: usize, row: usize| cmp(entry, if single_key { 0 } else { row
}).is_eq();
+ let sample = map.len().min(SAMPLE_ROWS);
+ let mut comparisons = 0;
+ let sampled_entries = scanner.scan(0..sample, |entry, row| {
+ comparisons += 1;
+ compare(entry, row)
+ });
+
+ // If the sampled rows compared more than half of their entries, stopping
+ // early is not paying off, so the remaining rows are cheaper to compare
all
+ // at once with the vectorized `eq`. We can only use `eq` when we have a
+ // single, non-nested key. The exact break-even point depends on the key
+ // type and the hardware; half keeps the cost of a wrong guess to about a
+ // third in either direction.
+ let rest = sample..map.len();
+ if single_key
+ && !key_type.is_nested()
+ && !rest.is_empty()
+ && comparisons * 2 > sampled_entries
+ {
+ let range_start = offsets[sample] as usize;
+ let in_range = map_keys.slice(range_start, last - range_start);
+ let matches = eq(&Scalar::new(keys.slice(0, 1)), &in_range)?;
+ // Neither side has nulls, so the value bits alone are meaningful.
+ let bits = matches.values();
+ scanner.scan(rest, |entry, _| bits.value(entry - range_start));
Review Comment:
Thanks, fixed.
--
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]
