comphead commented on code in PR #25201:
URL: https://github.com/apache/datafusion/pull/25201#discussion_r4007688037
##########
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:
**Test gap:** this branch is never reached on a sliced map.
`sliced_map_and_keys` has 4 rows, so `rest` is empty and the branch is skipped;
`vectorized_scan_after_missing_sample` has `offsets[0] == 0`, so `range_start`
is only ever exercised against a zero base.
I verified the indexing is correct today (200-row rotated-key map,
`slice(37, 150)`, matches a naive per-row scan), so this is coverage rather
than a live bug. But sliced map batches with more than 32 rows are common after
limit, coalesce and batch splitting, and an off-by-`offsets[0]` here silently
returns values from the wrong entries instead of erroring.
##########
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:
**P2:** this can return a later duplicate rather than the first match,
contradicting the rustdoc above ("the first entry whose key equals that row's
lookup key"). `self.hint` carries over from the previous row, so for a row that
holds the lookup key twice the hint can land on the second one.
Repro, added to the `map_lookup_tests` module on this branch:
```rust
// keys [1, 7 | 7, 7], two rows of length 2, lookup key 7
let got = map_lookup(&map, &Int32Array::from(vec![7])).unwrap();
// first match per row would be [1, 2]; actual is [1, 3]
```
Reordering so row 0 matches at position 0 (`keys [7, 1 | 7, 7]`) gives `[0,
2]`, the first match. So the answer depends on the preceding row, and therefore
on `batch_size`, since the hint is per call and each call is one batch. Both
pre-PR implementations used `(start..end).find(...)` and always returned the
first match.
Arrow does not enforce key uniqueness (`MapArray::try_new` validates
offsets, nulls and child count only), so duplicates arrive through Parquet,
IPC, FFI and custom `TableProvider`s. DataFusion's own `map()` / `make_map()`
reject them, so this is not reachable from a SQL literal.
Either scan `start..hinted` when the hint matches, or state in the rustdoc
that the entry returned for a row with duplicate keys is unspecified. The doc
as written is the part that is wrong today.
##########
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!(
Review Comment:
`DataType::Null` keys fail here, before reaching the null short-circuit
below, so `SELECT column1[NULL]` still errors (the `.slt` expectation updated
in this PR) while the new `MAP {1:'a', 2:'b'}[arrow_cast(NULL, 'Int64')]` case
returns `NULL`. Two spellings of the same SQL NULL behave differently.
Not a regression, but this PR rewrites the path and changes the message, so
it is the natural place to fix it: accept `DataType::Null` and return
`UInt32Array::new_null(map.len())`.
##########
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
Review Comment:
Nothing asserts the two strategies agree. The rustdoc makes a specific
semantic claim (total ordering, so `-0.0` and `0.0` are different keys and NaN
matches NaN), but `map_extract_float_keys` uses a single-row map and can
therefore only reach the comparator.
They agree today only because `ArrowNativeTypeOp::is_eq` for floats is
bitwise and `eq` routes through it, which is an arrow implementation detail
rather than a stable contract. If it moves to IEEE equality, `map[0.0]` starts
matching a `-0.0` key only for batches larger than 32 rows where this heuristic
flips, and nothing here catches it. `SAMPLE_ROWS` is private and the threshold
is inline, so there is no seam to force a strategy in a test; a test-only
override would make the equivalence test easy to write.
Separately on the threshold itself: the description says the measured
break-even is around 60% but this uses 50%, and
`map_extract/utf8_view/shuffled/1024x32 (+19.8%)` is exactly the marginal case.
At width 32 with a rotated key the sample compares about 16 entries per row
plus the failed hint probe, so `comparisons * 2` lands just over
`sampled_entries` and flips to `eq`. Moving toward the measured value may
recover that case without giving up the `missing` wins.
##########
datafusion/functions/benches/misc/get_field.rs:
##########
@@ -15,78 +15,156 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{ArrayRef, Int32Builder, MapBuilder, StringBuilder};
-use arrow::datatypes::{DataType, Field};
-use criterion::{Criterion, criterion_group};
+use arrow::array::{Array, ArrayRef, Int32Array, MapArray, StringViewArray,
StructArray};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::datatypes::{DataType, Field, FieldRef};
+use criterion::{Bencher, BenchmarkId, Criterion, criterion_group};
use datafusion_common::ScalarValue;
use datafusion_common::config::ConfigOptions;
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
use datafusion_functions::core::get_field;
use std::hint::black_box;
use std::sync::Arc;
-/// A map array with `size` rows, each holding `entries` key/value pairs.
-/// Every tenth row is null.
-fn map_array(size: usize, entries: usize) -> ArrayRef {
- let mut builder = MapBuilder::new(None, StringBuilder::new(),
Int32Builder::new());
+const ROWS: usize = 1024;
Review Comment:
The old bench had an 8192-row case and this pins `ROWS = 1024`. 8192 is
DataFusion's default `batch_size`, and the design amortizes a fixed 32-row
comparator sample over the batch, so 1024 rows both understates the benefit and
overstates the sample cost. Worth keeping one 8192-row shape, plus something
just above the threshold (say 40 rows, as produced by a selective filter) where
the sample is a large fraction of the batch rather than 0.4% of it.
Also, swapping plain `Utf8` for `Utf8View` drops a common shape:
`map<string, ...>` read without `schema_force_view_types` takes a different
comparator (`compare_bytes` vs `compare_byte_view`).
--
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]
