-
Notifications
You must be signed in to change notification settings - Fork 28.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[SPARK-23940][SQL] Add transform_values SQL function #22045
Changes from 5 commits
68392e3
fd05645
cdecd32
7df3ac5
b73106d
daf7935
56d08ef
9e89565
3382e1a
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -497,6 +497,60 @@ case class ArrayAggregate( | |
override def prettyName: String = "aggregate" | ||
} | ||
|
||
/** | ||
* Returns a map that applies the function to each value of the map. | ||
*/ | ||
@ExpressionDescription( | ||
usage = "_FUNC_(expr, func) - Transforms values in the map using the function.", | ||
examples = """ | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ditto. |
||
Examples: | ||
> SELECT _FUNC_(map(array(1, 2, 3), array(1, 2, 3), (k, v) -> v + 1); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: we need one more right parenthesis after the second There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. @ueshin Thanks for the review! and yes I agree, I made the same mistakes in both the PR's. I was waiting for the transform_key to converge so that I can make the same changes here. |
||
map(array(1, 2, 3), array(2, 3, 4)) | ||
> SELECT _FUNC_(map(array(1, 2, 3), array(1, 2, 3), (k, v) -> k + v); | ||
map(array(1, 2, 3), array(2, 4, 6)) | ||
""", | ||
since = "2.4.0") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ditto. |
||
case class TransformValues( | ||
argument: Expression, | ||
function: Expression) | ||
extends MapBasedSimpleHigherOrderFunction with CodegenFallback { | ||
|
||
override def nullable: Boolean = argument.nullable | ||
|
||
override def dataType: DataType = { | ||
val map = argument.dataType.asInstanceOf[MapType] | ||
MapType(map.keyType, function.dataType, function.nullable) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We can use |
||
} | ||
|
||
@transient val MapType(keyType, valueType, valueContainsNull) = argument.dataType | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
|
||
override def bind(f: (Expression, Seq[(DataType, Boolean)]) => LambdaFunction) | ||
: TransformValues = { | ||
copy(function = f(function, (keyType, false) :: (valueType, valueContainsNull) :: Nil)) | ||
} | ||
|
||
@transient lazy val (keyVar, valueVar) = { | ||
val LambdaFunction( | ||
_, (keyVar: NamedLambdaVariable) :: (valueVar: NamedLambdaVariable) :: Nil, _) = function | ||
(keyVar, valueVar) | ||
} | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: how about: @transient lazy val LambdaFunction(_,
(keyVar: NamedLambdaVariable) :: (valueVar: NamedLambdaVariable) :: Nil, _) = function |
||
|
||
override def nullSafeEval(inputRow: InternalRow, argumentValue: Any): Any = { | ||
val map = argumentValue.asInstanceOf[MapData] | ||
val f = functionForEval | ||
val resultValues = new GenericArrayData(new Array[Any](map.numElements)) | ||
var i = 0 | ||
while (i < map.numElements) { | ||
keyVar.value.set(map.keyArray().get(i, keyVar.dataType)) | ||
valueVar.value.set(map.valueArray().get(i, valueVar.dataType)) | ||
resultValues.update(i, f.eval(inputRow)) | ||
i += 1 | ||
} | ||
new ArrayBasedMapData(map.keyArray(), resultValues) | ||
} | ||
override def prettyName: String = "transform_values" | ||
} | ||
|
||
/** | ||
* Merges two given maps into a single map by applying function to the pair of values with | ||
* the same key. | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -95,6 +95,12 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper | |
aggregate(expr, zero, merge, identity) | ||
} | ||
|
||
def transformValues(expr: Expression, f: (Expression, Expression) => Expression): Expression = { | ||
val valueType = expr.dataType.asInstanceOf[MapType].valueType | ||
val keyType = expr.dataType.asInstanceOf[MapType].keyType | ||
TransformValues(expr, createLambda(keyType, false, valueType, true, f)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We should use |
||
} | ||
|
||
test("ArrayTransform") { | ||
val ai0 = Literal.create(Seq(1, 2, 3), ArrayType(IntegerType, containsNull = false)) | ||
val ai1 = Literal.create(Seq[Integer](1, null, 3), ArrayType(IntegerType, containsNull = true)) | ||
|
@@ -283,6 +289,61 @@ class HigherOrderFunctionsSuite extends SparkFunSuite with ExpressionEvalHelper | |
15) | ||
} | ||
|
||
test("TransformValues") { | ||
val ai0 = Literal.create( | ||
Map(1 -> 1, 2 -> 2, 3 -> 3), | ||
MapType(IntegerType, IntegerType)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you add |
||
val ai1 = Literal.create( | ||
Map(1 -> 1, 2 -> null, 3 -> 3), | ||
MapType(IntegerType, IntegerType)) | ||
val ain = Literal.create( | ||
Map.empty[Int, Int], | ||
MapType(IntegerType, IntegerType)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you add tests for |
||
|
||
val plusOne: (Expression, Expression) => Expression = (k, v) => v + 1 | ||
val valueUpdate: (Expression, Expression) => Expression = (k, v) => k * k | ||
|
||
checkEvaluation(transformValues(ai0, plusOne), Map(1 -> 2, 2 -> 3, 3 -> 4)) | ||
checkEvaluation(transformValues(ai0, valueUpdate), Map(1 -> 1, 2 -> 4, 3 -> 9)) | ||
checkEvaluation( | ||
transformValues(transformValues(ai0, plusOne), valueUpdate), Map(1 -> 1, 2 -> 4, 3 -> 9)) | ||
checkEvaluation(transformValues(ai1, plusOne), Map(1 -> 2, 2 -> null, 3 -> 4)) | ||
checkEvaluation(transformValues(ai1, valueUpdate), Map(1 -> 1, 2 -> 4, 3 -> 9)) | ||
checkEvaluation( | ||
transformValues(transformValues(ai1, plusOne), valueUpdate), Map(1 -> 1, 2 -> 4, 3 -> 9)) | ||
checkEvaluation(transformValues(ain, plusOne), Map.empty[Int, Int]) | ||
|
||
val as0 = Literal.create( | ||
Map("a" -> "xy", "bb" -> "yz", "ccc" -> "zx"), MapType(StringType, StringType)) | ||
val as1 = Literal.create( | ||
Map("a" -> "xy", "bb" -> null, "ccc" -> "zx"), MapType(StringType, StringType)) | ||
val asn = Literal.create(Map.empty[StringType, StringType], MapType(StringType, StringType)) | ||
|
||
val concatValue: (Expression, Expression) => Expression = (k, v) => Concat(Seq(k, v)) | ||
val valueTypeUpdate: (Expression, Expression) => Expression = | ||
(k, v) => Length(v) + 1 | ||
|
||
checkEvaluation( | ||
transformValues(as0, concatValue), Map("a" -> "axy", "bb" -> "bbyz", "ccc" -> "ccczx")) | ||
checkEvaluation(transformValues(as0, valueTypeUpdate), | ||
Map("a" -> 3, "bb" -> 3, "ccc" -> 3)) | ||
checkEvaluation( | ||
transformValues(transformValues(as0, concatValue), concatValue), | ||
Map("a" -> "aaxy", "bb" -> "bbbbyz", "ccc" -> "cccccczx")) | ||
checkEvaluation(transformValues(as1, concatValue), | ||
Map("a" -> "axy", "bb" -> null, "ccc" -> "ccczx")) | ||
checkEvaluation(transformValues(as1, valueTypeUpdate), | ||
Map("a" -> 3, "bb" -> null, "ccc" -> 3)) | ||
checkEvaluation( | ||
transformValues(transformValues(as1, concatValue), concatValue), | ||
Map("a" -> "aaxy", "bb" -> null, "ccc" -> "cccccczx")) | ||
checkEvaluation(transformValues(asn, concatValue), Map.empty[String, String]) | ||
checkEvaluation(transformValues(asn, valueTypeUpdate), Map.empty[String, Int]) | ||
checkEvaluation( | ||
transformValues(transformValues(asn, concatValue), valueTypeUpdate), | ||
Map.empty[String, Int]) | ||
} | ||
|
||
test("MapZipWith") { | ||
def map_zip_with( | ||
left: Expression, | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -51,3 +51,17 @@ select exists(ys, y -> y > 30) as v from nested; | |
|
||
-- Check for element existence in a null array | ||
select exists(cast(null as array<int>), y -> y > 30) as v; | ||
|
||
create or replace temporary view nested as values | ||
(1, map(1,1,2,2,3,3)), | ||
(2, map(4,4,5,5,6,6)) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit:
|
||
as t(x, ys); | ||
|
||
-- Identity Transform Keys in a map | ||
select transform_values(ys, (k, v) -> v) as v from nested; | ||
|
||
-- Transform Keys in a map by adding constant | ||
select transform_values(ys, (k, v) -> v + 1) as v from nested; | ||
|
||
-- Transform Keys in a map using values | ||
select transform_values(ys, (k, v) -> k + v) as v from nested; |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: indent