feat(test): enable udf in nexmark q14 (#14975)

Signed-off-by: Runji Wang <[email protected]>
risingwavelabs · Feb 8, 2024 · acb3fb6 · acb3fb6
1 parent 7cbf16a
commit acb3fb6
Show file tree

Hide file tree

Showing 20 changed files with 82 additions and 52 deletions.
diff --git a/ci/Dockerfile b/ci/Dockerfile
@@ -29,6 +29,7 @@ RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- --no-mo
 ENV PATH /root/.cargo/bin/:$PATH
 
 RUN rustup show
+RUN rustup default `rustup show active-toolchain | awk '{print $1}'`
 
 RUN curl -sSL "https://github.com/bufbuild/buf/releases/download/v1.4.0/buf-$(uname -s)-$(uname -m).tar.gz" | \
     tar -xvzf - -C /usr/local --strip-components 1

diff --git a/ci/build-ci-image.sh b/ci/build-ci-image.sh
@@ -10,7 +10,7 @@ cat ../rust-toolchain
 # shellcheck disable=SC2155
 
 # REMEMBER TO ALSO UPDATE ci/docker-compose.yml
-export BUILD_ENV_VERSION=v20240204
+export BUILD_ENV_VERSION=v20240208
 
 export BUILD_TAG="public.ecr.aws/x5u3w5h6/rw-build-env:${BUILD_ENV_VERSION}"
 

diff --git a/ci/docker-compose.yml b/ci/docker-compose.yml
@@ -71,7 +71,7 @@ services:
       retries: 5
 
   source-test-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240208
     depends_on:
       - mysql
       - db
@@ -81,7 +81,7 @@ services:
       - ..:/risingwave
 
   sink-test-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240208
     depends_on:
       - mysql
       - db
@@ -93,12 +93,12 @@ services:
       - ..:/risingwave
 
   rw-build-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240208
     volumes:
       - ..:/risingwave
 
   ci-flamegraph-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240208
     # NOTE(kwannoel): This is used in order to permit
     # syscalls for `nperf` (perf_event_open),
     # so it can do CPU profiling.
@@ -109,7 +109,7 @@ services:
       - ..:/risingwave
 
   regress-test-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240208
     depends_on:
       db:
         condition: service_healthy

diff --git a/ci/scripts/sql/nexmark/q14.drop.sql b/ci/scripts/sql/nexmark/q14.drop.sql
@@ -1,3 +1,4 @@
 -- noinspection SqlNoDataSourceInspectionForFile
 -- noinspection SqlResolveForFile
 DROP SINK nexmark_q14;
+DROP FUNCTION count_char;
diff --git a/ci/scripts/sql/nexmark/q14.sql b/ci/scripts/sql/nexmark/q14.sql
@@ -1,5 +1,16 @@
 -- noinspection SqlNoDataSourceInspectionForFile
 -- noinspection SqlResolveForFile
+
+CREATE FUNCTION count_char(s varchar, c varchar) RETURNS int LANGUAGE javascript AS $$
+    var count = 0;
+    for (var cc of s) {
+        if (cc === c) {
+            count++;
+        }
+    }
+    return count;
+$$;
+
 CREATE SINK nexmark_q14 AS
 SELECT auction,
        bidder,
@@ -15,11 +26,8 @@ SELECT auction,
                THEN 'nightTime'
            ELSE 'otherTime'
            END       AS bidTimeType,
-       date_time
-       -- extra
-       -- TODO: count_char is an UDF, add it back when we support similar functionality.
-       -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
-       -- count_char(extra, 'c') AS c_counts
+       date_time,
+       count_char(extra, 'c') AS c_counts
 FROM bid
 WHERE 0.908 * price > 1000000
   AND 0.908 * price < 50000000

diff --git a/e2e_test/source/nexmark_endless_sinks/nexmark_endless_part3.slt b/e2e_test/source/nexmark_endless_sinks/nexmark_endless_part3.slt
@@ -16,6 +16,9 @@ drop sink nexmark_q10;
 statement ok
 drop sink nexmark_q14;
 
+statement ok
+drop function count_char;
+
 statement ok
 drop sink nexmark_q15;
 

diff --git a/e2e_test/streaming/nexmark/drop_views.slt.part b/e2e_test/streaming/nexmark/drop_views.slt.part
@@ -34,6 +34,9 @@ drop materialized view nexmark_q10;
 statement ok
 drop materialized view nexmark_q14;
 
+statement ok
+drop function count_char;
+
 statement ok
 drop materialized view nexmark_q15;
 

diff --git a/e2e_test/streaming/nexmark/q14.slt.part b/e2e_test/streaming/nexmark/q14.slt.part
@@ -1,15 +1,15 @@
 query IIRTT rowsort
 SELECT * FROM nexmark_q14;
 ----
-1000	1001	12696452.316	nightTime	2015-07-15 00:00:11.003
-1000	1001	1753269.004	nightTime	2015-07-15 00:00:14.004
-1000	1001	2196534.628	nightTime	2015-07-15 00:00:04.001
-1000	1001	26038438.152	nightTime	2015-07-15 00:00:13.003
-1000	1001	39654430.240	nightTime	2015-07-15 00:00:04.001
-1000	1001	9503126.184	nightTime	2015-07-15 00:00:22.005
-1000	1006	1301754.200	nightTime	2015-07-15 00:00:12.003
-1006	1001	19666916.800	nightTime	2015-07-15 00:00:16.004
-1006	1001	3099951.044	nightTime	2015-07-15 00:00:16.004
-1008	1002	26289625.456	nightTime	2015-07-15 00:00:04.001
-1012	1001	3067656.208	nightTime	2015-07-15 00:00:20.005
-1012	1001	3299152.624	nightTime	2015-07-15 00:00:20.005
+1000	1001	12696452.316	nightTime	2015-07-15 00:00:11.003 1
+1000	1001	1753269.004		nightTime	2015-07-15 00:00:14.004 0
+1000	1001	2196534.628		nightTime	2015-07-15 00:00:04.001 3
+1000	1001	26038438.152	nightTime	2015-07-15 00:00:13.003 3
+1000	1001	39654430.240	nightTime	2015-07-15 00:00:04.001 2
+1000	1001	9503126.184		nightTime	2015-07-15 00:00:22.005 1
+1000	1006	1301754.200		nightTime	2015-07-15 00:00:12.003 2
+1006	1001	19666916.800	nightTime	2015-07-15 00:00:16.004 4
+1006	1001	3099951.044		nightTime	2015-07-15 00:00:16.004 3
+1008	1002	26289625.456	nightTime	2015-07-15 00:00:04.001 4
+1012	1001	3067656.208		nightTime	2015-07-15 00:00:20.005 3
+1012	1001	3299152.624		nightTime	2015-07-15 00:00:20.005 3
diff --git a/e2e_test/streaming/nexmark/sinks/q14.slt.part b/e2e_test/streaming/nexmark/sinks/q14.slt.part
@@ -1,3 +1,14 @@
+statement ok
+CREATE FUNCTION count_char(s varchar, c varchar) RETURNS int LANGUAGE javascript AS $$
+    var count = 0;
+    for (var cc of s) {
+        if (cc === c) {
+            count++;
+        }
+    }
+    return count;
+$$;
+
 statement ok
 CREATE SINK nexmark_q14 AS
 SELECT
@@ -15,11 +26,8 @@ SELECT
     THEN 'nightTime'
     ELSE 'otherTime'
   END AS bidTimeType,
-  date_time
-  -- extra
-  -- TODO: count_char is an UDF, add it back when we support similar functionality.
-  -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
-  -- count_char(extra, 'c') AS c_counts
+  date_time,
+  count_char(extra, 'c') AS c_counts
 FROM bid
 WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000
 WITH ( connector = 'blackhole', type = 'append-only', force_append_only = 'true');
diff --git a/e2e_test/streaming/nexmark/views/q14.slt.part b/e2e_test/streaming/nexmark/views/q14.slt.part
@@ -1,3 +1,14 @@
+statement ok
+CREATE FUNCTION count_char(s varchar, c varchar) RETURNS int LANGUAGE javascript AS $$
+    var count = 0;
+    for (var cc of s) {
+        if (cc === c) {
+            count++;
+        }
+    }
+    return count;
+$$;
+
 statement ok
 CREATE MATERIALIZED VIEW nexmark_q14 AS
 SELECT
@@ -15,10 +26,7 @@ SELECT
     THEN 'nightTime'
     ELSE 'otherTime'
   END AS bidTimeType,
-  date_time
-  -- extra
-  -- TODO: count_char is an UDF, add it back when we support similar functionality.
-  -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
-  -- count_char(extra, 'c') AS c_counts
+  date_time,
+  count_char(extra, 'c') AS c_counts
 FROM bid
 WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;
diff --git a/e2e_test/udf/wasm/src/lib.rs b/e2e_test/udf/wasm/src/lib.rs
@@ -1,6 +1,11 @@
 use arrow_udf::function;
 use rust_decimal::Decimal;
 
+#[function("count_char(varchar, varchar) -> int")]
+fn count_char(s: &str, c: &str) -> i32 {
+    s.matches(c).count() as i32
+}
+
 #[function("int_42() -> int")]
 fn int_42() -> i32 {
     42

diff --git a/src/expr/core/src/expr/expr_udf.rs b/src/expr/core/src/expr/expr_udf.rs
@@ -71,7 +71,6 @@ impl Expression for UserDefinedFunction {
         self.return_type.clone()
     }
 
-    #[cfg_or_panic(not(madsim))]
     async fn eval(&self, input: &DataChunk) -> Result<ArrayRef> {
         let vis = input.visibility();
         let mut columns = Vec::with_capacity(self.children.len());
@@ -82,7 +81,6 @@ impl Expression for UserDefinedFunction {
         self.eval_inner(columns, vis).await
     }
 
-    #[cfg_or_panic(not(madsim))]
     async fn eval_row(&self, input: &OwnedRow) -> Result<Datum> {
         let mut columns = Vec::with_capacity(self.children.len());
         for child in &self.children {
@@ -183,7 +181,6 @@ impl UserDefinedFunction {
     }
 }
 
-#[cfg_or_panic(not(madsim))]
 impl Build for UserDefinedFunction {
     fn build(
         prost: &ExprNode,
@@ -194,6 +191,7 @@ impl Build for UserDefinedFunction {
 
         let identifier = udf.get_identifier()?;
         let imp = match udf.language.as_str() {
+            #[cfg(not(madsim))]
             "wasm" | "rust" => {
                 let link = udf.get_link()?;
                 // Use `block_in_place` as an escape hatch to run async code here in sync context.
@@ -218,10 +216,13 @@ impl Build for UserDefinedFunction {
                 )?;
                 UdfImpl::JavaScript(rt)
             }
+            #[cfg(not(madsim))]
             _ => {
                 let link = udf.get_link()?;
                 UdfImpl::External(get_or_create_flight_client(link)?)
             }
+            #[cfg(madsim)]
+            l => panic!("UDF language {l:?} is not supported on madsim"),
         };
 
         let arg_schema = Arc::new(Schema::new(

diff --git a/src/frontend/planner_test/tests/testdata/input/nexmark.yaml b/src/frontend/planner_test/tests/testdata/input/nexmark.yaml
@@ -336,8 +336,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/input/nexmark_source.yaml b/src/frontend/planner_test/tests/testdata/input/nexmark_source.yaml
@@ -338,8 +338,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/input/nexmark_temporal_filter.yaml b/src/frontend/planner_test/tests/testdata/input/nexmark_temporal_filter.yaml
@@ -321,8 +321,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/input/nexmark_watermark.yaml b/src/frontend/planner_test/tests/testdata/input/nexmark_watermark.yaml
@@ -305,8 +305,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/output/nexmark.yaml b/src/frontend/planner_test/tests/testdata/output/nexmark.yaml
@@ -1185,8 +1185,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/output/nexmark_source.yaml b/src/frontend/planner_test/tests/testdata/output/nexmark_source.yaml
@@ -1003,8 +1003,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/output/nexmark_temporal_filter.yaml b/src/frontend/planner_test/tests/testdata/output/nexmark_temporal_filter.yaml
@@ -957,8 +957,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;

diff --git a/src/frontend/planner_test/tests/testdata/output/nexmark_watermark.yaml b/src/frontend/planner_test/tests/testdata/output/nexmark_watermark.yaml
@@ -1208,8 +1208,7 @@
       END AS bidTimeType,
       date_time,
       extra
-      -- TODO: count_char is an UDF, add it back when we support similar functionality.
-      -- https://github.com/nexmark/nexmark/blob/master/nexmark-flink/src/main/java/com/github/nexmark/flink/udf/CountChar.java
+      -- ignore UDF in planner test
       -- count_char(extra, 'c') AS c_counts
     FROM bid
     WHERE 0.908 * price > 1000000 AND 0.908 * price < 50000000;