neilconway commented on code in PR #25201:
URL: https://github.com/apache/datafusion/pull/25201#discussion_r4030485732
##########
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:
Fair enough, added the fix to this PR.
--
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]
