-
Notifications
You must be signed in to change notification settings - Fork 1k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: add functions as UDF input types (#6955)
- Loading branch information
1 parent
c0cccaa
commit 1b39ab5
Showing
17 changed files
with
683 additions
and
118 deletions.
There are no files selected for viewing
73 changes: 73 additions & 0 deletions
73
ksqldb-common/src/main/java/io/confluent/ksql/function/types/LambdaType.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,73 @@ | ||
/* | ||
* Copyright 2021 Confluent Inc. | ||
* | ||
* Licensed under the Confluent Community License (the "License"; you may not use | ||
* this file except in compliance with the License. You may obtain a copy of the | ||
* License at | ||
* | ||
* http://www.confluent.io/confluent-community-license | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT | ||
* WARRANTIES OF ANY KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations under the License. | ||
*/ | ||
|
||
package io.confluent.ksql.function.types; | ||
|
||
import com.google.common.collect.ImmutableList; | ||
import java.util.List; | ||
import java.util.Objects; | ||
|
||
public final class LambdaType extends ObjectType { | ||
|
||
private final ImmutableList<ParamType> inputTypes; | ||
private final ParamType returnType; | ||
|
||
private LambdaType( | ||
final List<ParamType> inputTypes, | ||
final ParamType returnType | ||
) { | ||
this.inputTypes = ImmutableList.copyOf( | ||
Objects.requireNonNull(inputTypes, "inputTypes")); | ||
this.returnType = Objects.requireNonNull(returnType, "returnType"); | ||
} | ||
|
||
public static LambdaType of( | ||
final List<ParamType> inputTypes, | ||
final ParamType returnType | ||
) { | ||
return new LambdaType(inputTypes, returnType); | ||
} | ||
|
||
public List<ParamType> inputTypes() { | ||
return inputTypes; | ||
} | ||
|
||
public ParamType returnType() { | ||
return returnType; | ||
} | ||
|
||
@Override | ||
public boolean equals(final Object o) { | ||
if (this == o) { | ||
return true; | ||
} | ||
if (o == null || getClass() != o.getClass()) { | ||
return false; | ||
} | ||
final LambdaType lambdaType = (LambdaType) o; | ||
return Objects.equals(inputTypes, lambdaType.inputTypes) | ||
&& Objects.equals(returnType, lambdaType.returnType); | ||
} | ||
|
||
@Override | ||
public int hashCode() { | ||
return Objects.hash(inputTypes, returnType); | ||
} | ||
|
||
@Override | ||
public String toString() { | ||
return "LAMBDA<" + inputTypes + ", " + returnType + ">"; | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
45 changes: 0 additions & 45 deletions
45
ksqldb-engine/src/main/java/io/confluent/ksql/engine/rewrite/LambdaContext.java
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
55 changes: 55 additions & 0 deletions
55
ksqldb-engine/src/main/java/io/confluent/ksql/function/udf/array/ArrayTransform.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
/* | ||
* Copyright 2021 Confluent Inc. | ||
* | ||
* Licensed under the Confluent Community License (the "License"; you may not use | ||
* this file except in compliance with the License. You may obtain a copy of the | ||
* License at | ||
* | ||
* http://www.confluent.io/confluent-community-license | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT | ||
* WARRANTIES OF ANY KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations under the License. | ||
*/ | ||
|
||
package io.confluent.ksql.function.udf.array; | ||
|
||
import io.confluent.ksql.function.FunctionCategory; | ||
import io.confluent.ksql.function.udf.Udf; | ||
import io.confluent.ksql.function.udf.UdfDescription; | ||
import io.confluent.ksql.function.udf.UdfParameter; | ||
import io.confluent.ksql.util.KsqlConstants; | ||
import java.util.List; | ||
import java.util.function.Function; | ||
import java.util.stream.Collectors; | ||
|
||
/** | ||
* Transform an array with a function | ||
*/ | ||
@SuppressWarnings("MethodMayBeStatic") // UDF methods can not be static. | ||
@UdfDescription( | ||
name = "array_transform", | ||
category = FunctionCategory.ARRAY, | ||
description = "Apply a function to each element in an array. " | ||
+ "The transformed array is returned.", | ||
author = KsqlConstants.CONFLUENT_AUTHOR | ||
) | ||
public class ArrayTransform { | ||
|
||
@Udf | ||
public <T, R> List<R> arrayTransform( | ||
@UdfParameter(description = "The array") final List<T> array, | ||
@UdfParameter(description = "The lambda function") final Function<T, R> function | ||
) { | ||
if (array == null) { | ||
return null; | ||
} | ||
return array.stream().map(item -> { | ||
if (item == null) { | ||
return null; | ||
} | ||
return function.apply(item); | ||
}).collect(Collectors.toList()); | ||
} | ||
} |
59 changes: 59 additions & 0 deletions
59
ksqldb-engine/src/main/java/io/confluent/ksql/function/udf/map/MapTransform.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
/* | ||
* Copyright 2021 Confluent Inc. | ||
* | ||
* Licensed under the Confluent Community License (the "License"; you may not use | ||
* this file except in compliance with the License. You may obtain a copy of the | ||
* License at | ||
* | ||
* http://www.confluent.io/confluent-community-license | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT | ||
* WARRANTIES OF ANY KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations under the License. | ||
*/ | ||
|
||
package io.confluent.ksql.function.udf.map; | ||
|
||
import io.confluent.ksql.function.FunctionCategory; | ||
import io.confluent.ksql.function.udf.Udf; | ||
import io.confluent.ksql.function.udf.UdfDescription; | ||
import io.confluent.ksql.function.udf.UdfParameter; | ||
import io.confluent.ksql.util.KsqlConstants; | ||
import java.util.Map; | ||
import java.util.function.BiFunction; | ||
import java.util.stream.Collectors; | ||
|
||
/** | ||
* Transform a map's key and values using two lambda functions | ||
*/ | ||
@UdfDescription( | ||
name = "map_transform", | ||
category = FunctionCategory.MAP, | ||
description = "Apply one function to each key and " | ||
+ "one function to each value of a map. " | ||
+ "The two arguments for each function are in order: key, value. " | ||
+ "The first function provided will be applied to each key and the " | ||
+ "second one applied to each value. " | ||
+ "The transformed map is returned.", | ||
author = KsqlConstants.CONFLUENT_AUTHOR | ||
) | ||
public class MapTransform { | ||
|
||
@Udf | ||
public <K,V,R,T> Map<R,T> mapTransform( | ||
@UdfParameter(description = "The map") final Map<K, V> map, | ||
@UdfParameter(description = "The key lambda function") final BiFunction<K, V, R> biFunction1, | ||
@UdfParameter(description = "The value lambda function") final BiFunction<K, V, T> biFunction2 | ||
) { | ||
if (map == null) { | ||
return null; | ||
} | ||
|
||
return map.entrySet() | ||
.stream() | ||
.collect(Collectors.toMap( | ||
entry -> biFunction1.apply(entry.getKey(), entry.getValue()), | ||
entry -> biFunction2.apply(entry.getKey(), entry.getValue()))); | ||
} | ||
} |
60 changes: 60 additions & 0 deletions
60
ksqldb-engine/src/main/java/io/confluent/ksql/function/udf/map/ReduceMap.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,60 @@ | ||
/* | ||
* Copyright 2021 Confluent Inc. | ||
* | ||
* Licensed under the Confluent Community License (the "License"; you may not use | ||
* this file except in compliance with the License. You may obtain a copy of the | ||
* License at | ||
* | ||
* http://www.confluent.io/confluent-community-license | ||
* | ||
* Unless required by applicable law or agreed to in writing, software | ||
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT | ||
* WARRANTIES OF ANY KIND, either express or implied. See the License for the | ||
* specific language governing permissions and limitations under the License. | ||
*/ | ||
|
||
package io.confluent.ksql.function.udf.map; | ||
|
||
import io.confluent.ksql.execution.codegen.helpers.TriFunction; | ||
import io.confluent.ksql.function.FunctionCategory; | ||
import io.confluent.ksql.function.udf.Udf; | ||
import io.confluent.ksql.function.udf.UdfDescription; | ||
import io.confluent.ksql.function.udf.UdfParameter; | ||
import io.confluent.ksql.util.KsqlConstants; | ||
import java.util.Map; | ||
import java.util.Map.Entry; | ||
|
||
/** | ||
* Reduce a map using an initial state and function | ||
*/ | ||
@UdfDescription( | ||
name = "map_reduce", | ||
category = FunctionCategory.MAP, | ||
description = "Reduce the input map down to a single value " | ||
+ "using an initial state and a function. " | ||
+ "The initial state (s) is passed into the scope of the function. " | ||
+ "Each invocation returns a new value for s, " | ||
+ "which the next invocation will receive. " | ||
+ "The final value for s is returned." | ||
+ "The three arguments for the function are in order: key, value, state.", | ||
author = KsqlConstants.CONFLUENT_AUTHOR | ||
) | ||
public class ReduceMap { | ||
|
||
@Udf | ||
public <K,V,S> S reduceMap( | ||
@UdfParameter(description = "The map") final Map<K, V> map, | ||
@UdfParameter(description = "The initial state") final S initialState, | ||
@UdfParameter(description = "The reduce function") final TriFunction<K, V, S, S> triFunction | ||
) { | ||
if (map == null) { | ||
return null; | ||
} | ||
|
||
S state = initialState; | ||
for (Entry<K, V> entry : map.entrySet()) { | ||
state = triFunction.apply(entry.getKey(), entry.getValue(), state); | ||
} | ||
return state; | ||
} | ||
} |
Oops, something went wrong.