Skip to content

Commit

Permalink
SNOW-1361405 Emit Span for UDx (#122)
Browse files Browse the repository at this point in the history 
* create methodChain

* alias

* rollup

* groupByGroupingSets

* limit

* na

* chain

* add method chain

* update

* update test

* fix test

* fix test

* fix test

* fix test

* fix test

* reformat

* fix test

* rename

* temp udf

* temp udf

* temp udf

* perm udf

* session.udf

* udf

* udtf

* sproc

* sproc

* setup Java

* temp2

* fix test

* udf

* table function

* sproc

* sproc

* fix test

* fix test

* fix test
  • Loading branch information
@sfc-gh-bli
sfc-gh-bli authored Jul 23, 2024
1 parent 2a4b910 commit d000ced
Show file tree
Hide file tree
Showing 17 changed files with 2,868 additions and 1,720 deletions.
81 changes: 56 additions & 25 deletions src/main/java/com/snowflake/snowpark_java/Functions.java
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,12 @@
package com.snowflake.snowpark_java;

import static com.snowflake.snowpark.internal.OpenTelemetry.javaUDF;

import com.snowflake.snowpark.internal.JavaUtils;
import com.snowflake.snowpark_java.types.DataType;
import com.snowflake.snowpark_java.udf.*;
import java.util.List;
import java.util.function.Supplier;

/**
* Provides utility functions that generate Column expression that you can pass to DataFrame
Expand Down Expand Up @@ -3912,8 +3915,8 @@ public static Column callUDF(String udfName, Column... cols) {
// |* @return The result UserDefinedFunction reference
// |*/
// |public static UserDefinedFunction udf(JavaUDF$x<$types> func,$input DataType output) {
// | return getActiveSession().udf()
// | .registerTemporary(func,${if(x == 0) "" else " input,"} output);
// | return userDefinedFunction("udf", () -> getActiveSession().udf()
// | .registerTemporary(func,${if(x == 0) "" else " input,"} output));
// |}
// |""".stripMargin)
// }
Expand All @@ -3927,7 +3930,8 @@ public static Column callUDF(String udfName, Column... cols) {
* @return The result UserDefinedFunction reference
*/
public static UserDefinedFunction udf(JavaUDF0<?> func, DataType output) {
return getActiveSession().udf().registerTemporary(func, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, output));
}

/**
Expand All @@ -3940,7 +3944,8 @@ public static UserDefinedFunction udf(JavaUDF0<?> func, DataType output) {
* @return The result UserDefinedFunction reference
*/
public static UserDefinedFunction udf(JavaUDF1<?, ?> func, DataType input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -3953,7 +3958,8 @@ public static UserDefinedFunction udf(JavaUDF1<?, ?> func, DataType input, DataT
* @return The result UserDefinedFunction reference
*/
public static UserDefinedFunction udf(JavaUDF2<?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -3967,7 +3973,8 @@ public static UserDefinedFunction udf(JavaUDF2<?, ?, ?> func, DataType[] input,
*/
public static UserDefinedFunction udf(
JavaUDF3<?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -3981,7 +3988,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF4<?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -3995,7 +4003,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF5<?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4009,7 +4018,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF6<?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4023,7 +4033,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF7<?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4037,7 +4048,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF8<?, ?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4051,7 +4063,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF9<?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4065,7 +4078,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF10<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4079,7 +4093,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF11<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4093,7 +4108,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF12<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4107,7 +4123,8 @@ public static UserDefinedFunction udf(
*/
public static UserDefinedFunction udf(
JavaUDF13<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func, DataType[] input, DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4123,7 +4140,8 @@ public static UserDefinedFunction udf(
JavaUDF14<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4139,7 +4157,8 @@ public static UserDefinedFunction udf(
JavaUDF15<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4155,7 +4174,8 @@ public static UserDefinedFunction udf(
JavaUDF16<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4171,7 +4191,8 @@ public static UserDefinedFunction udf(
JavaUDF17<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4187,7 +4208,8 @@ public static UserDefinedFunction udf(
JavaUDF18<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4203,7 +4225,8 @@ public static UserDefinedFunction udf(
JavaUDF19<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4219,7 +4242,8 @@ public static UserDefinedFunction udf(
JavaUDF20<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4235,7 +4259,8 @@ public static UserDefinedFunction udf(
JavaUDF21<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

/**
Expand All @@ -4251,10 +4276,16 @@ public static UserDefinedFunction udf(
JavaUDF22<?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?> func,
DataType[] input,
DataType output) {
return getActiveSession().udf().registerTemporary(func, input, output);
return userDefinedFunction(
"udf", () -> getActiveSession().udf().registerTemporary(func, input, output));
}

private static Session getActiveSession() {
return new Session(JavaUtils.getActiveSession());
}

private static UserDefinedFunction userDefinedFunction(
String funcName, Supplier<UserDefinedFunction> func) {
return javaUDF("Functions", funcName, "", "", 0, func);
}
}
Loading

0 comments on commit d000ced

Please sign in to comment.