neilconway commented on code in PR #25201:
URL: https://github.com/apache/datafusion/pull/25201#discussion_r4030453744
##########
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));
+ } else {
+ scanner.scan(rest, compare);
+ }
+ Ok(scanner.finish())
+}
+
+/// Number of rows [`map_lookup`] scans with the comparator before deciding
+/// whether the rest of the batch is better served by the vectorized `eq`.
+const SAMPLE_ROWS: usize = 32;
+
+/// The value type of a dictionary-encoded type, or the type itself.
+fn strip_dictionary(data_type: &DataType) -> &DataType {
+ match data_type {
+ DataType::Dictionary(_, value_type) => value_type,
+ other => other,
+ }
+}
+
+/// Scans map rows for the first entry that satisfies a predicate.
+struct RowScanner<'a> {
+ offsets: &'a [i32],
+ /// Rows to skip: null map rows and rows whose lookup key is null.
+ skip: Option<NullBuffer>,
+ found: UInt32Builder,
+ /// Position within its row of the most recent match.
+ hint: usize,
+}
+
+impl<'a> RowScanner<'a> {
+ fn new(map: &'a MapArray, key_nulls: Option<&NullBuffer>) -> Self {
+ Self {
+ offsets: map.value_offsets(),
+ skip: NullBuffer::union(map.nulls(), key_nulls),
+ found: UInt32Builder::with_capacity(map.len()),
+ hint: 0,
+ }
+ }
+
+ /// Scans `rows`, recording the first entry for which `is_match(entry,
row)`
+ /// holds, or null for a skipped row or a row without a match. Returns the
+ /// number of entries in the rows that were scanned.
+ fn scan(
+ &mut self,
+ rows: Range<usize>,
+ mut is_match: impl FnMut(usize, usize) -> bool,
+ ) -> usize {
+ let mut entries = 0;
+ for row in rows {
+ if self.skip.as_ref().is_some_and(|skip| skip.is_null(row)) {
+ self.found.append_null();
+ continue;
+ }
+ let start = self.offsets[row] as usize;
+ let end = self.offsets[row + 1] as usize;
+ entries += end - start;
+
+ // Rows in a batch usually share the same key order, so try the
+ // position where the previous row matched first. When that guess
+ // is right, the lookup costs one comparison wherever the key sits.
+ let hinted = start + self.hint;
+ let found = if hinted < end && is_match(hinted, row) {
+ Some(hinted)
+ } else {
+ (start..end).find(|&entry| entry != hinted && is_match(entry,
row))
+ };
Review Comment:
Thanks, good point. I think we should just update the docs here; a map with
duplicate keys is a corner case and I think returning any of the associated
values is defensible.
--
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]
