From f545601811939d20412bdc062d5fe3bfa3c6e62d Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 19:30:25 -0700 Subject: [PATCH 01/16] add aether udf skills --- LICENSE | 405 +++++++++++++- skills/.gitignore | 165 ++++++ skills/README.md | 188 +++++++ skills/docs/dev/TESTING.md | 33 ++ skills/docs/dev/VERSIONS.md | 53 ++ skills/examples/CalculateRiskUDF.scala | 24 + skills/examples/FormatPhoneUDF.java | 26 + skills/examples/IntegerMultiplyBy2UDF.java | 78 +++ skills/pyproject.toml | 40 ++ skills/tests/test_export/__init__.py | 6 + skills/tests/test_export/cuda_fixtures.py | 208 +++++++ skills/tests/test_export/java_fixtures.py | 201 +++++++ skills/tests/test_export/scala_fixtures.py | 193 +++++++ skills/tests/test_export/test_jvm.py | 523 ++++++++++++++++++ skills/tests/test_export/utils.py | 91 +++ skills/tests/test_skill_frontmatter.py | 40 ++ skills/udf-benchmark/CUDF_MICROBENCHMARKS.md | 30 + skills/udf-benchmark/SKILL.md | 80 +++ skills/udf-convert-to-cuda/SKILL.md | 169 ++++++ .../examples/CosineSimilarityJni.cpp | 67 +++ .../CosineSimilarityNativeRapidsUDF.java | 56 ++ .../examples/cosine_similarity.cu | 119 ++++ .../examples/cosine_similarity.hpp | 22 + .../references/JNI_CUDA_GUIDE.md | 162 ++++++ .../references/NATIVE_BUILD_ENV.md | 92 +++ .../templates/cuda/Dockerfile | 65 +++ .../cuda/native/scripts/extract-cudf-libs.sh | 81 +++ .../cuda/native/src/main/cpp/CMakeLists.txt | 119 ++++ .../main/cpp/src/PlaceholderUDFNameJni.cpp | 58 ++ .../src/main/cpp/src/placeholder_udf_name.cu | 18 + .../src/main/cpp/src/placeholder_udf_name.hpp | 13 + .../main/java/com/udf/NativeUDFLoader.java | 29 + .../PlaceholderUDFNameNativeRapidsUDF.java | 45 ++ skills/udf-convert-to-cudf/SKILL.md | 128 +++++ .../examples/URLDecode.java | 57 ++ .../examples/URLDecodeExtendsFunction.scala | 44 ++ .../examples/URLDecodeHive.java | 57 ++ .../examples/URLDecodeWithField.scala | 48 ++ .../references/RAPIDS_UDF.md | 111 ++++ skills/udf-convert-to-sql/SKILL.md | 87 +++ .../examples/FormatPhone.java | 27 + .../examples/FormatPhone.scala | 22 + .../examples/FormatPhoneHive.java | 26 + .../examples/NormalizeTags.java | 37 ++ .../examples/NormalizeTags.scala | 25 + .../examples/NormalizeTagsHive.java | 32 ++ .../examples/format_phone.sql | 17 + .../examples/normalize_tags.sql | 15 + skills/udf-gen-test/SKILL.md | 148 +++++ .../templates/java/.mvn/jvm.config | 16 + skills/udf-gen-test/templates/java/pom.xml | 342 ++++++++++++ .../templates/java/run_gen_data.sh | 73 +++ .../templates/java/run_micro_benchmark.sh | 60 ++ .../templates/java/run_spark_benchmark.sh | 70 +++ .../src/main/java/com/udf/SparkUtils.java | 126 +++++ .../main/java/com/udf/bench/BenchUtils.java | 112 ++++ .../src/main/java/com/udf/bench/GenData.java | 110 ++++ .../java/com/udf/bench/MicroBenchRunner.java | 313 +++++++++++ .../java/com/udf/bench/SparkBenchRunner.java | 163 ++++++ .../test/java/com/udf/CudfComparisonTest.java | 70 +++ .../test/java/com/udf/SqlComparisonTest.java | 76 +++ .../java/src/test/java/com/udf/TestUtils.java | 78 +++ .../java/src/test/java/com/udf/UnitTest.java | 122 ++++ .../templates/scala/.mvn/jvm.config | 15 + skills/udf-gen-test/templates/scala/pom.xml | 382 +++++++++++++ .../templates/scala/run_gen_data.sh | 73 +++ .../templates/scala/run_micro_benchmark.sh | 60 ++ .../templates/scala/run_spark_benchmark.sh | 70 +++ .../scala/src/main/scala/com/udf/Arm.scala | 63 +++ .../src/main/scala/com/udf/SparkUtils.scala | 84 +++ .../main/scala/com/udf/bench/BenchUtils.scala | 109 ++++ .../main/scala/com/udf/bench/GenData.scala | 101 ++++ .../com/udf/bench/MicroBenchRunner.scala | 281 ++++++++++ .../com/udf/bench/SparkBenchRunner.scala | 176 ++++++ .../scala/com/udf/CudfComparisonTest.scala | 56 ++ .../scala/com/udf/SqlComparisonTest.scala | 59 ++ .../src/test/scala/com/udf/TestUtils.scala | 47 ++ .../src/test/scala/com/udf/UnitTest.scala | 96 ++++ skills/udf-judge-conversion/SKILL.md | 92 +++ skills/udf-optimize-cudf/SKILL.md | 149 +++++ .../references/OPTIMIZATION_PATTERNS.md | 22 + 81 files changed, 8045 insertions(+), 1 deletion(-) create mode 100644 skills/.gitignore create mode 100644 skills/README.md create mode 100644 skills/docs/dev/TESTING.md create mode 100644 skills/docs/dev/VERSIONS.md create mode 100644 skills/examples/CalculateRiskUDF.scala create mode 100644 skills/examples/FormatPhoneUDF.java create mode 100644 skills/examples/IntegerMultiplyBy2UDF.java create mode 100644 skills/pyproject.toml create mode 100644 skills/tests/test_export/__init__.py create mode 100644 skills/tests/test_export/cuda_fixtures.py create mode 100644 skills/tests/test_export/java_fixtures.py create mode 100644 skills/tests/test_export/scala_fixtures.py create mode 100644 skills/tests/test_export/test_jvm.py create mode 100644 skills/tests/test_export/utils.py create mode 100644 skills/tests/test_skill_frontmatter.py create mode 100644 skills/udf-benchmark/CUDF_MICROBENCHMARKS.md create mode 100644 skills/udf-benchmark/SKILL.md create mode 100644 skills/udf-convert-to-cuda/SKILL.md create mode 100644 skills/udf-convert-to-cuda/examples/CosineSimilarityJni.cpp create mode 100644 skills/udf-convert-to-cuda/examples/CosineSimilarityNativeRapidsUDF.java create mode 100644 skills/udf-convert-to-cuda/examples/cosine_similarity.cu create mode 100644 skills/udf-convert-to-cuda/examples/cosine_similarity.hpp create mode 100644 skills/udf-convert-to-cuda/references/JNI_CUDA_GUIDE.md create mode 100644 skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md create mode 100644 skills/udf-convert-to-cuda/templates/cuda/Dockerfile create mode 100644 skills/udf-convert-to-cuda/templates/cuda/native/scripts/extract-cudf-libs.sh create mode 100644 skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/CMakeLists.txt create mode 100644 skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/PlaceholderUDFNameJni.cpp create mode 100644 skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.cu create mode 100644 skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.hpp create mode 100644 skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/NativeUDFLoader.java create mode 100644 skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java create mode 100644 skills/udf-convert-to-cudf/SKILL.md create mode 100644 skills/udf-convert-to-cudf/examples/URLDecode.java create mode 100644 skills/udf-convert-to-cudf/examples/URLDecodeExtendsFunction.scala create mode 100644 skills/udf-convert-to-cudf/examples/URLDecodeHive.java create mode 100644 skills/udf-convert-to-cudf/examples/URLDecodeWithField.scala create mode 100644 skills/udf-convert-to-cudf/references/RAPIDS_UDF.md create mode 100644 skills/udf-convert-to-sql/SKILL.md create mode 100644 skills/udf-convert-to-sql/examples/FormatPhone.java create mode 100644 skills/udf-convert-to-sql/examples/FormatPhone.scala create mode 100644 skills/udf-convert-to-sql/examples/FormatPhoneHive.java create mode 100644 skills/udf-convert-to-sql/examples/NormalizeTags.java create mode 100644 skills/udf-convert-to-sql/examples/NormalizeTags.scala create mode 100644 skills/udf-convert-to-sql/examples/NormalizeTagsHive.java create mode 100644 skills/udf-convert-to-sql/examples/format_phone.sql create mode 100644 skills/udf-convert-to-sql/examples/normalize_tags.sql create mode 100644 skills/udf-gen-test/SKILL.md create mode 100644 skills/udf-gen-test/templates/java/.mvn/jvm.config create mode 100644 skills/udf-gen-test/templates/java/pom.xml create mode 100644 skills/udf-gen-test/templates/java/run_gen_data.sh create mode 100644 skills/udf-gen-test/templates/java/run_micro_benchmark.sh create mode 100644 skills/udf-gen-test/templates/java/run_spark_benchmark.sh create mode 100644 skills/udf-gen-test/templates/java/src/main/java/com/udf/SparkUtils.java create mode 100644 skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/BenchUtils.java create mode 100644 skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/GenData.java create mode 100644 skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java create mode 100644 skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/SparkBenchRunner.java create mode 100644 skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java create mode 100644 skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java create mode 100644 skills/udf-gen-test/templates/java/src/test/java/com/udf/TestUtils.java create mode 100644 skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java create mode 100644 skills/udf-gen-test/templates/scala/.mvn/jvm.config create mode 100644 skills/udf-gen-test/templates/scala/pom.xml create mode 100644 skills/udf-gen-test/templates/scala/run_gen_data.sh create mode 100644 skills/udf-gen-test/templates/scala/run_micro_benchmark.sh create mode 100644 skills/udf-gen-test/templates/scala/run_spark_benchmark.sh create mode 100644 skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala create mode 100644 skills/udf-gen-test/templates/scala/src/main/scala/com/udf/SparkUtils.scala create mode 100644 skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/BenchUtils.scala create mode 100644 skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/GenData.scala create mode 100644 skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala create mode 100644 skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/SparkBenchRunner.scala create mode 100644 skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala create mode 100644 skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala create mode 100644 skills/udf-gen-test/templates/scala/src/test/scala/com/udf/TestUtils.scala create mode 100644 skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala create mode 100644 skills/udf-judge-conversion/SKILL.md create mode 100644 skills/udf-optimize-cudf/SKILL.md create mode 100644 skills/udf-optimize-cudf/references/OPTIMIZATION_PATTERNS.md diff --git a/LICENSE b/LICENSE index 261eeb9e9f8..884b334376f 100644 --- a/LICENSE +++ b/LICENSE @@ -1,3 +1,406 @@ +Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +This code is dual-licensed with documentation/skills under the CC-BY-4.0 AND source code under Apache-2.0 license terms. + + +Attribution 4.0 International + +======================================================================= + +Creative Commons Corporation ("Creative Commons") is not a law firm and +does not provide legal services or legal advice. Distribution of +Creative Commons public licenses does not create a lawyer-client or +other relationship. Creative Commons makes its licenses and related +information available on an "as-is" basis. Creative Commons gives no +warranties regarding its licenses, any material licensed under their +terms and conditions, or any related information. Creative Commons +disclaims all liability for damages resulting from their use to the +fullest extent possible. + +Using Creative Commons Public Licenses + +Creative Commons public licenses provide a standard set of terms and +conditions that creators and other rights holders may use to share +original works of authorship and other material subject to copyright +and certain other rights specified in the public license below. The +following considerations are for informational purposes only, are not +exhaustive, and do not form part of our licenses. + + Considerations for licensors: Our public licenses are + intended for use by those authorized to give the public + permission to use material in ways otherwise restricted by + copyright and certain other rights. Our licenses are + irrevocable. Licensors should read and understand the terms + and conditions of the license they choose before applying it. + Licensors should also secure all rights necessary before + applying our licenses so that the public can reuse the + material as expected. Licensors should clearly mark any + material not subject to the license. This includes other CC- + licensed material, or material used under an exception or + limitation to copyright. More considerations for licensors: + wiki.creativecommons.org/Considerations_for_licensors + + Considerations for the public: By using one of our public + licenses, a licensor grants the public permission to use the + licensed material under specified terms and conditions. If + the licensor's permission is not necessary for any reason--for + example, because of any applicable exception or limitation to + copyright--then that use is not regulated by the license. Our + licenses grant only permissions under copyright and certain + other rights that a licensor has authority to grant. Use of + the licensed material may still be restricted for other + reasons, including because others have copyright or other + rights in the material. A licensor may make special requests, + such as asking that all changes be marked or described. + Although not required by our licenses, you are encouraged to + respect those requests where reasonable. More considerations + for the public: + wiki.creativecommons.org/Considerations_for_licensees + +======================================================================= + +Creative Commons Attribution 4.0 International Public License + +By exercising the Licensed Rights (defined below), You accept and agree +to be bound by the terms and conditions of this Creative Commons +Attribution 4.0 International Public License ("Public License"). To the +extent this Public License may be interpreted as a contract, You are +granted the Licensed Rights in consideration of Your acceptance of +these terms and conditions, and the Licensor grants You such rights in +consideration of benefits the Licensor receives from making the +Licensed Material available under these terms and conditions. + + +Section 1 -- Definitions. + + a. Adapted Material means material subject to Copyright and Similar + Rights that is derived from or based upon the Licensed Material + and in which the Licensed Material is translated, altered, + arranged, transformed, or otherwise modified in a manner requiring + permission under the Copyright and Similar Rights held by the + Licensor. For purposes of this Public License, where the Licensed + Material is a musical work, performance, or sound recording, + Adapted Material is always produced where the Licensed Material is + synched in timed relation with a moving image. + + b. Adapter's License means the license You apply to Your Copyright + and Similar Rights in Your contributions to Adapted Material in + accordance with the terms and conditions of this Public License. + + c. Copyright and Similar Rights means copyright and/or similar rights + closely related to copyright including, without limitation, + performance, broadcast, sound recording, and Sui Generis Database + Rights, without regard to how the rights are labeled or + categorized. For purposes of this Public License, the rights + specified in Section 2(b)(1)-(2) are not Copyright and Similar + Rights. + + d. Effective Technological Measures means those measures that, in the + absence of proper authority, may not be circumvented under laws + fulfilling obligations under Article 11 of the WIPO Copyright + Treaty adopted on December 20, 1996, and/or similar international + agreements. + + e. Exceptions and Limitations means fair use, fair dealing, and/or + any other exception or limitation to Copyright and Similar Rights + that applies to Your use of the Licensed Material. + + f. Licensed Material means the artistic or literary work, database, + or other material to which the Licensor applied this Public + License. + + g. Licensed Rights means the rights granted to You subject to the + terms and conditions of this Public License, which are limited to + all Copyright and Similar Rights that apply to Your use of the + Licensed Material and that the Licensor has authority to license. + + h. Licensor means the individual(s) or entity(ies) granting rights + under this Public License. + + i. Share means to provide material to the public by any means or + process that requires permission under the Licensed Rights, such + as reproduction, public display, public performance, distribution, + dissemination, communication, or importation, and to make material + available to the public including in ways that members of the + public may access the material from a place and at a time + individually chosen by them. + + j. Sui Generis Database Rights means rights other than copyright + resulting from Directive 96/9/EC of the European Parliament and of + the Council of 11 March 1996 on the legal protection of databases, + as amended and/or succeeded, as well as other essentially + equivalent rights anywhere in the world. + + k. You means the individual or entity exercising the Licensed Rights + under this Public License. Your has a corresponding meaning. + + +Section 2 -- Scope. + + a. License grant. + + 1. Subject to the terms and conditions of this Public License, + the Licensor hereby grants You a worldwide, royalty-free, + non-sublicensable, non-exclusive, irrevocable license to + exercise the Licensed Rights in the Licensed Material to: + + a. reproduce and Share the Licensed Material, in whole or + in part; and + + b. produce, reproduce, and Share Adapted Material. + + 2. Exceptions and Limitations. For the avoidance of doubt, where + Exceptions and Limitations apply to Your use, this Public + License does not apply, and You do not need to comply with + its terms and conditions. + + 3. Term. The term of this Public License is specified in Section + 6(a). + + 4. Media and formats; technical modifications allowed. The + Licensor authorizes You to exercise the Licensed Rights in + all media and formats whether now known or hereafter created, + and to make technical modifications necessary to do so. The + Licensor waives and/or agrees not to assert any right or + authority to forbid You from making technical modifications + necessary to exercise the Licensed Rights, including + technical modifications necessary to circumvent Effective + Technological Measures. For purposes of this Public License, + simply making modifications authorized by this Section 2(a) + (4) never produces Adapted Material. + + 5. Downstream recipients. + + a. Offer from the Licensor -- Licensed Material. Every + recipient of the Licensed Material automatically + receives an offer from the Licensor to exercise the + Licensed Rights under the terms and conditions of this + Public License. + + b. No downstream restrictions. You may not offer or impose + any additional or different terms or conditions on, or + apply any Effective Technological Measures to, the + Licensed Material if doing so restricts exercise of the + Licensed Rights by any recipient of the Licensed + Material. + + 6. No endorsement. Nothing in this Public License constitutes or + may be construed as permission to assert or imply that You + are, or that Your use of the Licensed Material is, connected + with, or sponsored, endorsed, or granted official status by, + the Licensor or others designated to receive attribution as + provided in Section 3(a)(1)(A)(i). + + b. Other rights. + + 1. Moral rights, such as the right of integrity, are not + licensed under this Public License, nor are publicity, + privacy, and/or other similar personality rights; however, to + the extent possible, the Licensor waives and/or agrees not to + assert any such rights held by the Licensor to the limited + extent necessary to allow You to exercise the Licensed + Rights, but not otherwise. + + 2. Patent and trademark rights are not licensed under this + Public License. + + 3. To the extent possible, the Licensor waives any right to + collect royalties from You for the exercise of the Licensed + Rights, whether directly or through a collecting society + under any voluntary or waivable statutory or compulsory + licensing scheme. In all other cases the Licensor expressly + reserves any right to collect such royalties. + + +Section 3 -- License Conditions. + +Your exercise of the Licensed Rights is expressly made subject to the +following conditions. + + a. Attribution. + + 1. If You Share the Licensed Material (including in modified + form), You must: + + a. retain the following if it is supplied by the Licensor + with the Licensed Material: + + i. identification of the creator(s) of the Licensed + Material and any others designated to receive + attribution, in any reasonable manner requested by + the Licensor (including by pseudonym if + designated); + + ii. a copyright notice; + + iii. a notice that refers to this Public License; + + iv. a notice that refers to the disclaimer of + warranties; + + v. a URI or hyperlink to the Licensed Material to the + extent reasonably practicable; + + b. indicate if You modified the Licensed Material and + retain an indication of any previous modifications; and + + c. indicate the Licensed Material is licensed under this + Public License, and include the text of, or the URI or + hyperlink to, this Public License. + + 2. You may satisfy the conditions in Section 3(a)(1) in any + reasonable manner based on the medium, means, and context in + which You Share the Licensed Material. For example, it may be + reasonable to satisfy the conditions by providing a URI or + hyperlink to a resource that includes the required + information. + + 3. If requested by the Licensor, You must remove any of the + information required by Section 3(a)(1)(A) to the extent + reasonably practicable. + + 4. If You Share Adapted Material You produce, the Adapter's + License You apply must not prevent recipients of the Adapted + Material from complying with this Public License. + + +Section 4 -- Sui Generis Database Rights. + +Where the Licensed Rights include Sui Generis Database Rights that +apply to Your use of the Licensed Material: + + a. for the avoidance of doubt, Section 2(a)(1) grants You the right + to extract, reuse, reproduce, and Share all or a substantial + portion of the contents of the database; + + b. if You include all or a substantial portion of the database + contents in a database in which You have Sui Generis Database + Rights, then the database in which You have Sui Generis Database + Rights (but not its individual contents) is Adapted Material; and + + c. You must comply with the conditions in Section 3(a) if You Share + all or a substantial portion of the contents of the database. + +For the avoidance of doubt, this Section 4 supplements and does not +replace Your obligations under this Public License where the Licensed +Rights include other Copyright and Similar Rights. + + +Section 5 -- Disclaimer of Warranties and Limitation of Liability. + + a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE + EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS + AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF + ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, + IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, + WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR + PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, + ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT + KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT + ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. + + b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE + TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, + NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, + INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, + COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR + USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN + ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR + DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR + IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. + + c. The disclaimer of warranties and limitation of liability provided + above shall be interpreted in a manner that, to the extent + possible, most closely approximates an absolute disclaimer and + waiver of all liability. + + +Section 6 -- Term and Termination. + + a. This Public License applies for the term of the Copyright and + Similar Rights licensed here. However, if You fail to comply with + this Public License, then Your rights under this Public License + terminate automatically. + + b. Where Your right to use the Licensed Material has terminated under + Section 6(a), it reinstates: + + 1. automatically as of the date the violation is cured, provided + it is cured within 30 days of Your discovery of the + violation; or + + 2. upon express reinstatement by the Licensor. + + For the avoidance of doubt, this Section 6(b) does not affect any + right the Licensor may have to seek remedies for Your violations + of this Public License. + + c. For the avoidance of doubt, the Licensor may also offer the + Licensed Material under separate terms or conditions or stop + distributing the Licensed Material at any time; however, doing so + will not terminate this Public License. + + d. Sections 1, 5, 6, 7, and 8 survive termination of this Public + License. + + +Section 7 -- Other Terms and Conditions. + + a. The Licensor shall not be bound by any additional or different + terms or conditions communicated by You unless expressly agreed. + + b. Any arrangements, understandings, or agreements regarding the + Licensed Material not stated herein are separate from and + independent of the terms and conditions of this Public License. + + +Section 8 -- Interpretation. + + a. For the avoidance of doubt, this Public License does not, and + shall not be interpreted to, reduce, limit, restrict, or impose + conditions on any use of the Licensed Material that could lawfully + be made without permission under this Public License. + + b. To the extent possible, if any provision of this Public License is + deemed unenforceable, it shall be automatically reformed to the + minimum extent necessary to make it enforceable. If the provision + cannot be reformed, it shall be severed from this Public License + without affecting the enforceability of the remaining terms and + conditions. + + c. No term or condition of this Public License will be waived and no + failure to comply consented to unless expressly agreed to by the + Licensor. + + d. Nothing in this Public License constitutes or may be interpreted + as a limitation upon, or waiver of, any privileges and immunities + that apply to the Licensor or You, including from the legal + processes of any jurisdiction or authority. + + +======================================================================= + +Creative Commons is not a party to its public +licenses. Notwithstanding, Creative Commons may elect to apply one of +its public licenses to material it publishes and in those instances +will be considered the “Licensor.” The text of the Creative Commons +public licenses is dedicated to the public domain under the CC0 Public +Domain Dedication. Except for the limited purpose of indicating that +material is shared under a Creative Commons public license or as +otherwise permitted by the Creative Commons policies published at +creativecommons.org/policies, Creative Commons does not authorize the +use of the trademark "Creative Commons" or any other trademark or logo +of Creative Commons without its prior written consent including, +without limitation, in connection with any unauthorized modifications +to any of its public licenses or any other arrangements, +understandings, or agreements concerning use of licensed material. For +the avoidance of doubt, this paragraph does not form part of the +public licenses. + +Creative Commons may be contacted at creativecommons.org. + + + + Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ @@ -186,7 +589,7 @@ same "printed page" as the copyright notice for easier identification within third-party archives. - Copyright [yyyy] [name of copyright owner] + Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. diff --git a/skills/.gitignore b/skills/.gitignore new file mode 100644 index 00000000000..16a0e72b69f --- /dev/null +++ b/skills/.gitignore @@ -0,0 +1,165 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/_build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pdm +.pdm-python +.pdm-build/ + +# pixi +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# Abstra +.abstra/ + +# Visual Studio Code +.vscode/ +.cursor/ +.claude/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml + +# Scala +.scala-build/ +.metals/ +.bsp/ + +# Maven config under skills is source, not generated output. +!**/.mvn/ +!**/.mvn/** diff --git a/skills/README.md b/skills/README.md new file mode 100644 index 00000000000..8e5694a5b3d --- /dev/null +++ b/skills/README.md @@ -0,0 +1,188 @@ +# Project Aether Agent Skills + +Aether Agent is a set of skills to convert Apache Spark User-Defined Functions (UDFs) for GPU acceleration with the [RAPIDS Accelerator for Apache Spark](https://github.com/NVIDIA/spark-rapids). It provides: + +1. **Test generation** -- Create unit tests and test data for existing UDFs. +2. **Conversion** -- Convert a UDF to a GPU-compatible implementation (SQL, cuDF RapidsUDF, or native CUDA RapidsUDF). +3. **Benchmarking** -- Generate synthetic data and benchmark the original UDF against the GPU implementation. +4. **Optimization** -- Iteratively profile and optimize a cuDF RapidsUDF for GPU performance. + +
+Table of Contents + +- [Supported Formats](#supported-formats) +- [Prerequisites](#prerequisites) +- [Selecting an LLM](#selecting-an-llm) +- [Quick Start](#quick-start) + - [Installing Skills](#installing-skills) + - [Using Skills](#using-skills) + - [Quick Start](#quick-start-1) + +
+ +## Supported Formats + +| UDF Type | cuDF RapidsUDF | CUDA RapidsUDF | Spark SQL | +|-----------|----------------|------------------------|-----------| +| Java UDF | Yes | Yes | Yes | +| Hive UDF | Yes | Yes | Yes | +| Scala UDF | Yes | Yes | Yes | +| Java UDAF | -- | -- | Yes | +| Hive UDAF | -- | -- | Yes | +| Scala UDAF | -- | -- | Yes | + +## Prerequisites + +- **[Maven](https://maven.apache.org/install.html)** is required to build/compile UDFs. +- **[JDK](https://docs.oracle.com/en/java/javase/index.html)** must be installed on the system. +- **Local GPU** with [CUDA toolkit](https://developer.nvidia.com/cuda/toolkit) is required (see [Spark RAPIDS compatibility](https://nvidia.github.io/spark-rapids/docs/download.html) for version requirements). + +If a local GPU is not available, another option is to run Aether Agent from a cloud instance, such as AWS EC2. + +## Selecting an LLM + +For best results, we recommend the latest reasoning models from OpenAI, Anthropic, or Google. As a good proxy, models near the top of the [Terminal-Bench 2.0 leaderboard](https://www.tbench.ai/leaderboard/terminal-bench/2.0) tend to perform well. + +## Quick Start + +Skills require any IDE or LLM that supports the [agent skills spec](https://skill.md/) (e.g., Cursor, Codex, Claude Code). + +### Installing Skills + +Copy the skills from this repo into your project: + +```bash +# Claude Code +mkdir -p /path/to/your/project/.claude/skills/ +cp -r skills/* /path/to/your/project/.claude/skills/ + +# Codex +mkdir -p /path/to/your/project/.agents/skills/ +cp -r skills/* /path/to/your/project/.agents/skills/ + +# Cursor +mkdir -p /path/to/your/project/.cursor/skills/ +cp -r skills/* /path/to/your/project/.cursor/skills/ + +# Kiro +mkdir -p /path/to/your/project/.kiro/skills/ +cp -r skills/* /path/to/your/project/.kiro/skills/ +``` + +### Using Skills + +Skills follow a multi-step workflow: + +1. **[udf-gen-test](udf-gen-test/SKILL.md)** -- Generate a unit test for the UDF +2. **[udf-convert-to-cudf](udf-convert-to-cudf/SKILL.md)**, **[udf-convert-to-cuda](udf-convert-to-cuda/SKILL.md)**, or **[udf-convert-to-sql](udf-convert-to-sql/SKILL.md)** -- Convert the UDF to a GPU-compatible implementation +3. **[udf-judge-conversion](udf-judge-conversion/SKILL.md)** -- Review generated tests and implementations for coverage gaps, bugs, and edge cases +4. **[udf-benchmark](udf-benchmark/SKILL.md)** -- Benchmark CPU vs GPU performance +5. **[udf-optimize-cudf](udf-optimize-cudf/SKILL.md)** -- Iteratively profile and optimize the cuDF RapidsUDF + +To invoke a skill, use your IDE's skill command, or simply describe the task and let the agent load the skill automatically. + +```bash +# Manual invocation +❯ Use the /udf-gen-test skill to generate a unit test for @FormatPhoneUDF.java + +# Automatic invocation +❯ Generate a unit test for @FormatPhoneUDF.java +``` + +Each skill builds on the output of the previous one -- udf-gen-test produces a project with a passing unit test, which the conversion skills use as input, and the udf-benchmark skill uses the conversion output. + +You can invoke multiple steps in a single prompt: + +```bash +❯ Generate a unit test for @FormatPhoneUDF.java, then convert it to cuDF, native CUDA, or SQL and benchmark +``` + +### Quick Start + +Once you've installed the skills, try the workflow with one of the provided example UDFs: +- Java: [FormatPhoneUDF.java](examples/FormatPhoneUDF.java) +- Hive: [IntegerMultiplyBy2UDF.java](examples/IntegerMultiplyBy2UDF.java) +- Scala: [CalculateRiskUDF.scala](examples/CalculateRiskUDF.scala) + +#### Workflow + +(Example dialogue is from Claude Code). + +1. Open your project in your IDE of choice with the skills installed. +2. Ask the agent to generate a unit test: + + ```bash + ❯ Please generate a unit test for @FormatPhoneUDF.java + ⎿ Read FormatPhoneUDF.java (23 lines) + + ● Skill(udf-gen-test) + ⎿  Successfully loaded skill + ● I'll generate a unit test for FormatPhoneUDF.java. Let me start by setting up the project. + ● ... + ● Unit test generation is complete. Here's a summary: + + Results + + - Project directory: FormatPhoneUDF/ + - Unit test: FormatPhoneUDF/src/test/java/com/udf/UnitTest.java + - Status: All tests passing + - Test coverage: 12 cases + ``` + +3. Once the test passes, convert the UDF: + + ```bash + ❯ Please convert the UDF to cuDF + + ● I'll use the udf-convert-to-cudf skill to convert this UDF to a GPU-accelerated RapidsUDF. + ● Skill(udf-convert-to-cudf) + ⎿  Successfully loaded skill + ● ... + ● Results + + The RapidsUDF conversion is complete. + + - RapidsUDF: FormatPhoneUDF/src/main/java/com/udf/FormatPhoneRapidsUDF.java + - Comparison test: passes with no memory leaks + ``` + + Alternatively, ask for a judge pass as part of the conversion request: + + ```bash + ❯ Please convert the UDF to cuDF and run the judge review agent on the result + ``` + + **Note:** The conversion agent will only invoke the judge when you explicitly ask for it. Otherwise, the agent performs a local self-review. You can also invoke the udf-judge-conversion skill on its own, at any time. + +4. Finally, benchmark the result: + + ```bash + ❯ Please benchmark the implementations on 100M rows. + + ● Skill(udf-benchmark) + ⎿  Successfully loaded skill + ● ... + ● Benchmark Results — 100M rows + + ┌─────────┬─────────────┐ + │ Mode │ Runtime (s) │ + ├─────────┼─────────────┤ + │ CPU │ 16.27 │ + ├─────────┼─────────────┤ + │ GPU │ 6.52 │ + ├─────────┼─────────────┤ + │ Speedup │ 2.50x │ + └─────────┴─────────────┘ + + The GPU RapidsUDF implementation is 2.5x faster than the CPU UDF on 100 million rows. + ``` + +5. Optionally for cuDF RapidsUDF conversions, optimize the implementation: + + ```bash + ❯ Please optimize the implementation + + ● Skill(udf-optimize-cudf) + ⎿  Successfully loaded skill + ● ... + ``` diff --git a/skills/docs/dev/TESTING.md b/skills/docs/dev/TESTING.md new file mode 100644 index 00000000000..505a96a0258 --- /dev/null +++ b/skills/docs/dev/TESTING.md @@ -0,0 +1,33 @@ +# Testing + +## Setup + +Set up a local dev environment: + +```bash +python -m venv .venv +source .venv/bin/activate +pip install -e ".[dev]" +``` + +## Fast Tests + +Run the fast tests: + +```bash +pytest -m "not slow" +``` + +These are generally lightweight skill validation tests, such as verifying skill frontmatter. + +## Integration Tests + +Run the integration tests: + +```bash +pytest -m slow -s tests/test_export +``` + +These tests deterministically fill in the Java/Scala template projects from `skills/udf-gen-test/templates/` with fixture implementations, then actually compile and run Spark tests and benchmark scripts locally. + +Thus they require JDK, Maven and Maven repository access, and a GPU environment for GPU paths. diff --git a/skills/docs/dev/VERSIONS.md b/skills/docs/dev/VERSIONS.md new file mode 100644 index 00000000000..e56fffdb236 --- /dev/null +++ b/skills/docs/dev/VERSIONS.md @@ -0,0 +1,53 @@ +# Version Update Guide + +## Files To Update + +### Java udf-gen-test Maven template + +File: `skills/udf-gen-test/templates/java/pom.xml` + +Update these properties together: + +- `` +- `` +- `` +- `` if the RAPIDS artifact classifier changes +- `` +- `` + +### Scala udf-gen-test Maven template + +File: `skills/udf-gen-test/templates/scala/pom.xml` + +Update these properties together: + +- `` +- `` +- `` +- `` +- `` if the RAPIDS artifact classifier changes +- `` +- `` + +### Native CUDA dependency extraction + +File: `skills/udf-convert-to-cuda/templates/cuda/native/scripts/extract-cudf-libs.sh` + +Update these default values: + +- `SCALA_VERSION` +- `RAPIDS4SPARK_VERSION` +- `CUDA_VERSION` if the RAPIDS artifact classifier changes +- `CUDF_BRANCH` + +### Native CUDA CMake template + +File: `skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/CMakeLists.txt` + +Update these values: + +- `RAPIDS_CMAKE_BRANCH` +- `project(RAPIDSUDFJNI VERSION ...)` +- `rapids_cpm_find(cudf ...)` + +`RAPIDS_CMAKE_BRANCH` should generally match the RAPIDS/cuDF branch or tag used by the Maven templates and `extract-cudf-libs.sh`. The `rapids_cpm_find(cudf...)` version should use the RAPIDS major/minor CPM version, for example `26.04.00` for `26.04.0`. diff --git a/skills/examples/CalculateRiskUDF.scala b/skills/examples/CalculateRiskUDF.scala new file mode 100644 index 00000000000..8e9dc6070c8 --- /dev/null +++ b/skills/examples/CalculateRiskUDF.scala @@ -0,0 +1,24 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package examples + +/** + * Calculate risk score based on credit score. + * + * @param creditScore Credit score + * @return Risk score + */ +class CalculateRiskUDF extends Function1[Integer, String] with Serializable { + override def apply(creditScore: Integer): String = { + Option(creditScore) match { + case Some(score) if score >= 750 => "LOW" + case Some(score) if score >= 650 => "MEDIUM" + case Some(score) if score >= 500 => "HIGH" + case Some(score) if score < 500 => "VERY_HIGH" + case None => "UNKNOWN" + } + } +} diff --git a/skills/examples/FormatPhoneUDF.java b/skills/examples/FormatPhoneUDF.java new file mode 100644 index 00000000000..f747887911d --- /dev/null +++ b/skills/examples/FormatPhoneUDF.java @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package examples; + +import org.apache.spark.sql.api.java.UDF1; + +/** Strip non-digit characters and format as (XXX) XXX-XXXX. */ +public class FormatPhoneUDF implements UDF1 { + @Override + public String call(String phone) throws Exception { + if (phone == null) { + return null; + } + String digits = phone.replaceAll("[^0-9]", ""); + if (digits.length() != 10) { + return null; + } + return String.format("(%s) %s-%s", + digits.substring(0, 3), + digits.substring(3, 6), + digits.substring(6)); + } +} diff --git a/skills/examples/IntegerMultiplyBy2UDF.java b/skills/examples/IntegerMultiplyBy2UDF.java new file mode 100644 index 00000000000..c6b0cb0055f --- /dev/null +++ b/skills/examples/IntegerMultiplyBy2UDF.java @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package examples; + +import org.apache.hadoop.hive.ql.exec.Description; +import org.apache.hadoop.hive.ql.exec.UDFArgumentException; +import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException; +import org.apache.hadoop.hive.ql.metadata.HiveException; +import org.apache.hadoop.hive.ql.udf.generic.GenericUDF; +import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector; +import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector; +import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory; +import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo; +import org.apache.hadoop.hive.serde2.typeinfo.TypeInfoUtils; +import org.apache.hadoop.io.LongWritable; +import org.apache.log4j.Logger; + +@Description(name = "integer_multiply_by_2", value = "_FUNC_(x) - Returns x * 2 for integer values") +public class IntegerMultiplyBy2UDF extends GenericUDF { + private static final Logger LOG = Logger.getLogger(IntegerMultiplyBy2UDF.class); + private PrimitiveObjectInspector inputOI; + + @Override + public ObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException { + if (arguments.length != 1) { + throw new UDFArgumentException("Exactly one argument is expected."); + } + + ObjectInspector oi = arguments[0]; + if (oi.getCategory() != ObjectInspector.Category.PRIMITIVE) { + throw new UDFArgumentTypeException(0, "Argument must be PRIMITIVE, but " + oi.getCategory().name() + " was passed."); + } + + inputOI = (PrimitiveObjectInspector) oi; + + // Check if input is numeric + if (inputOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.INT && + inputOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.LONG && + inputOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.SHORT && + inputOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.BYTE) { + throw new UDFArgumentTypeException(0, "Argument must be numeric (INT/LONG/SHORT/BYTE), but " + inputOI.getPrimitiveCategory().name() + " was passed."); + } + + // Return LongWritable type for the result + return PrimitiveObjectInspectorFactory.writableLongObjectInspector; + } + + @Override + public Object evaluate(DeferredObject[] arguments) throws HiveException { + if (arguments == null || arguments.length != 1) { + return null; + } + + Object input = arguments[0].get(); + if (input == null) { + return null; + } + + long value = getLongValue(input); + return new LongWritable(value * 2); + } + + @Override + public String getDisplayString(String[] children) { + return "integer_multiply_by_2(" + (children != null ? String.join(",", children) : "") + ")"; + } + + private long getLongValue(Object obj) { + if (obj instanceof Number) { + return ((Number) obj).longValue(); + } else { + throw new IllegalArgumentException("Cannot convert " + obj.getClass().getName() + " to long"); + } + } +} diff --git a/skills/pyproject.toml b/skills/pyproject.toml new file mode 100644 index 00000000000..e977287231d --- /dev/null +++ b/skills/pyproject.toml @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[project] +name = "aether-agent" +version = "0.1.0" +description = "Convert Spark UDFs into GPU implementations" +authors = [ + {name = "Rishi Chandra", email = "rishic@nvidia.com"} +] +readme = "README.md" +requires-python = ">=3.10" +classifiers = [ + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] + +[project.optional-dependencies] +dev = [ + "pytest==8.4.1", + "PyYAML==6.0.3", + "isort==6.0.1", + "black==25.1.0", + "ruff==0.12.8", +] + +[tool.setuptools] +packages = [] + +[tool.pyright] +typeCheckingMode = "standard" + +[tool.pytest.ini_options] +markers = [ + "slow: integration tests", +] diff --git a/skills/tests/test_export/__init__.py b/skills/tests/test_export/__init__.py new file mode 100644 index 00000000000..2ae3327da34 --- /dev/null +++ b/skills/tests/test_export/__init__.py @@ -0,0 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Tests for the export package. +""" diff --git a/skills/tests/test_export/cuda_fixtures.py b/skills/tests/test_export/cuda_fixtures.py new file mode 100644 index 00000000000..51175478c85 --- /dev/null +++ b/skills/tests/test_export/cuda_fixtures.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 + +""" +CUDA source fixtures for JVM export integration tests. +""" + +NATIVE_RAPIDS_UDF_SOURCE = """\ +package com.udf; + +import ai.rapids.cudf.ColumnVector; +import com.nvidia.spark.RapidsUDF; +import org.apache.hadoop.hive.ql.exec.UDF; +import org.apache.spark.sql.api.java.UDF1; + +public class IntegerMultiplyBy2NativeRapidsUDF extends UDF + implements UDF1, RapidsUDF { + public Integer evaluate(Integer value) { + if (value == null) return null; + return value * 2; + } + + @Override + public Integer call(Integer value) { + return evaluate(value); + } + + @Override + public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { + if (args.length != 1) { + throw new IllegalArgumentException("Unexpected argument count: " + args.length); + } + if (numRows != args[0].getRowCount()) { + throw new IllegalArgumentException( + "Expected " + numRows + " rows, received " + args[0].getRowCount()); + } + + NativeUDFLoader.ensureLoaded(); + return new ColumnVector(integerMultiplyBy2(args[0].getNativeView())); + } + + private static native long integerMultiplyBy2(long inputView); +} +""" + +JNI_SOURCE = """\ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +// SPDX-License-Identifier: Apache-2.0 + +#include "integer_multiply_by_2.hpp" + +#include +#include +#include + +#include + +#include +#include + +namespace { + +constexpr char const* RUNTIME_ERROR_CLASS = "java/lang/RuntimeException"; +constexpr char const* ILLEGAL_ARG_CLASS = "java/lang/IllegalArgumentException"; + +void throw_java_exception(JNIEnv* env, char const* class_name, char const* message) +{ + jclass ex_class = env->FindClass(class_name); + if (ex_class != nullptr) { + env->ThrowNew(ex_class, message); + } +} + +} // namespace + +extern "C" { + +JNIEXPORT jlong JNICALL +Java_com_udf_IntegerMultiplyBy2NativeRapidsUDF_integerMultiplyBy2(JNIEnv* env, + jclass, + jlong input_view) +{ + try { + auto input = reinterpret_cast(input_view); + if (input == nullptr) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, "input column view is null"); + return 0; + } + if (input->type().id() != cudf::type_id::INT32) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, "input must be INT32"); + return 0; + } + + std::unique_ptr result = integer_multiply_by_2(*input); + return reinterpret_cast(result.release()); + } catch (std::bad_alloc const& e) { + auto message = std::string("Unable to allocate native memory: ") + e.what(); + throw_java_exception(env, RUNTIME_ERROR_CLASS, message.c_str()); + } catch (std::invalid_argument const& e) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, e.what()); + } catch (std::exception const& e) { + throw_java_exception(env, RUNTIME_ERROR_CLASS, e.what()); + } + return 0; +} + +} +""" + +CUDA_SOURCE = """\ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +// SPDX-License-Identifier: Apache-2.0 + +#include "integer_multiply_by_2.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace { + +__global__ void multiply_by_2_kernel(int32_t const* input, int32_t* output, cudf::size_type size) +{ + auto const idx = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + if (idx < size) { + output[idx] = input[idx] * 2; + } +} + +} // namespace + +std::unique_ptr integer_multiply_by_2( + cudf::column_view const& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + if (input.type().id() != cudf::type_id::INT32) { + throw std::invalid_argument("input must be INT32"); + } + + auto const row_count = input.size(); + auto null_mask = cudf::copy_bitmask(input, stream, mr); + auto result = cudf::make_numeric_column( + input.type(), row_count, std::move(null_mask), input.null_count(), stream, mr); + + if (row_count > 0) { + constexpr int threads_per_block = 256; + int const blocks = (row_count + threads_per_block - 1) / threads_per_block; + multiply_by_2_kernel<<>>( + input.data(), result->mutable_view().data(), row_count); + CUDF_CHECK_CUDA(stream.value()); + } + + return result; +} +""" + +HEADER_SOURCE = """\ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +// SPDX-License-Identifier: Apache-2.0 + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include + +std::unique_ptr integer_multiply_by_2( + cudf::column_view const& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +""" + +CMAKE_SOURCE_FILES = """\ +set(SOURCE_FILES + "src/IntegerMultiplyBy2Jni.cpp" + "src/integer_multiply_by_2.cu" +) +""" + +PLACEHOLDER_FILES = ( + "src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java", + "native/src/main/cpp/src/PlaceholderUDFNameJni.cpp", + "native/src/main/cpp/src/placeholder_udf_name.cu", + "native/src/main/cpp/src/placeholder_udf_name.hpp", +) + +NATIVE_SOURCE_FILES = { + "src/main/java/com/udf/IntegerMultiplyBy2NativeRapidsUDF.java": NATIVE_RAPIDS_UDF_SOURCE, + "native/src/main/cpp/src/IntegerMultiplyBy2Jni.cpp": JNI_SOURCE, + "native/src/main/cpp/src/integer_multiply_by_2.cu": CUDA_SOURCE, + "native/src/main/cpp/src/integer_multiply_by_2.hpp": HEADER_SOURCE, +} diff --git a/skills/tests/test_export/java_fixtures.py b/skills/tests/test_export/java_fixtures.py new file mode 100644 index 00000000000..2845308800d --- /dev/null +++ b/skills/tests/test_export/java_fixtures.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Java source code fixtures for integration tests. +""" + +from .utils import replace_java_todo_method + +NAME = "java" +REPLACE_TODO_FN = replace_java_todo_method +TEST_SELECTOR_FLAG = "-Dtest" + +# --------------------------------------------------------------------------- +# UDF source code +# --------------------------------------------------------------------------- + +UDF_SOURCE = """\ +package com.udf; + +import org.apache.hadoop.hive.ql.exec.UDF; + +public class IntegerMultiplyBy2UDF extends UDF { + public Integer evaluate(Integer value) { + if (value == null) return null; + return value * 2; + } +} +""" + +RAPIDS_UDF_SOURCE = """\ +package com.udf; + +import ai.rapids.cudf.ColumnVector; +import ai.rapids.cudf.Scalar; +import com.nvidia.spark.RapidsUDF; +import org.apache.hadoop.hive.ql.exec.UDF; + +public class IntegerMultiplyBy2RapidsUDF extends UDF implements RapidsUDF { + public Integer evaluate(Integer value) { + if (value == null) return null; + return value * 2; + } + + @Override + public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { + try (Scalar two = Scalar.fromInt(2)) { + return args[0].mul(two); + } + } +} +""" + +SQL_SOURCE = """\ +SELECT *, + value * 2 AS result +FROM test_table +""" + +# --------------------------------------------------------------------------- +# Unit test methods +# --------------------------------------------------------------------------- + +UNIT_TEST_METHODS = { + "createTestData": """\ + public static Dataset createTestData(SparkSession spark) { + StructType schema = new StructType(new StructField[]{ + DataTypes.createStructField("id", DataTypes.IntegerType, false), + DataTypes.createStructField("value", DataTypes.IntegerType, true) + }); + List data = Arrays.asList( + RowFactory.create(1, 123), + RowFactory.create(2, 0), + RowFactory.create(3, -5), + RowFactory.create(4, null) + ); + return spark.createDataFrame(data, schema); + }""", + "registerUDF": """\ + public static void registerUDF(SparkSession spark, String udfName) { + spark.sql("CREATE TEMPORARY FUNCTION " + udfName + + " AS 'com.udf.IntegerMultiplyBy2UDF'"); + }""", + "executeUDF": """\ + public static Dataset executeUDF(SparkSession spark, String udfName, Dataset testDF) { + testDF.createOrReplaceTempView("test_table"); + return spark.sql("SELECT *, " + udfName + + "(value) AS result FROM test_table"); + }""", + "verifyUDFResults": """\ + public static void verifyUDFResults(Dataset resultDF, Dataset testDF) { + Row[] results = (Row[]) resultDF.sort("id").collect(); + Assert.assertEquals(246, (int) results[0].getAs("result")); + Assert.assertEquals(0, (int) results[1].getAs("result")); + Assert.assertEquals(-10, (int) results[2].getAs("result")); + Assert.assertTrue(results[3].isNullAt(results[3].fieldIndex("result"))); + }""", +} + +RAPIDS_UDF_REGISTER = """\ + public static void registerRapidsUDF(SparkSession spark, String udfName) { + spark.sql("CREATE TEMPORARY FUNCTION " + udfName + + " AS 'com.udf.IntegerMultiplyBy2RapidsUDF'"); + }""" + +NATIVE_RAPIDS_UDF_REGISTER = """\ + public static void registerRapidsUDF(SparkSession spark, String udfName) { + spark.sql("CREATE TEMPORARY FUNCTION " + udfName + + " AS 'com.udf.IntegerMultiplyBy2NativeRapidsUDF'"); + }""" + +# --------------------------------------------------------------------------- +# BenchUtils methods +# --------------------------------------------------------------------------- + +BENCH_GENERATE = """\ + public static Dataset generateSyntheticData( + SparkSession spark, long numRows, int numPartitions) { + Dataset baseDF = spark.range(0, numRows, 1, numPartitions).toDF("id"); + return baseDF.select( + col("id"), + expr("CAST(rand() * 1000 AS INT)").alias("value") + ); + }""" + +BENCH_CPU = """\ + public static Dataset executeCpu(SparkSession spark, Dataset df) { + df.createOrReplaceTempView("bench_table"); + spark.sql("CREATE TEMPORARY FUNCTION integer_multiply_by_2" + + " AS 'com.udf.IntegerMultiplyBy2UDF'"); + return spark.sql("SELECT *, integer_multiply_by_2(value)" + + " AS result FROM bench_table"); + }""" + +BENCH_GPU_CUDF = """\ + public static Dataset executeGpu(SparkSession spark, Dataset df) { + df.createOrReplaceTempView("bench_table"); + spark.sql("CREATE TEMPORARY FUNCTION integer_multiply_by_2_rapids" + + " AS 'com.udf.IntegerMultiplyBy2RapidsUDF'"); + return spark.sql("SELECT *, integer_multiply_by_2_rapids(value)" + + " AS result FROM bench_table"); + }""" + +BENCH_GPU_CUDA = """\ + public static Dataset executeGpu(SparkSession spark, Dataset df) { + df.createOrReplaceTempView("bench_table"); + spark.sql("CREATE TEMPORARY FUNCTION integer_multiply_by_2_native" + + " AS 'com.udf.IntegerMultiplyBy2NativeRapidsUDF'"); + return spark.sql("SELECT *, integer_multiply_by_2_native(value)" + + " AS result FROM bench_table"); + }""" + +BENCH_GPU_SQL = """\ + public static Dataset executeGpu(SparkSession spark, Dataset df) { + df.createOrReplaceTempView("bench_table"); + try { + String sqlContent = new String( + java.nio.file.Files.readAllBytes( + java.nio.file.Paths.get("src/main/resources/integer_multiply_by_2.sql"))); + String benchSql = sqlContent.replace("test_table", "bench_table"); + return spark.sql(benchSql); + } catch (java.io.IOException e) { + throw new RuntimeException(e); + } + }""" + +# --------------------------------------------------------------------------- +# MicroBenchRunner methods +# --------------------------------------------------------------------------- + +MICRO_PREPARE_CPU = """\ + public static Object[] prepareCpuData(HostColumnVector[] hostColumns, int numRows) { + Integer[] values = new Integer[numRows]; + for (int i = 0; i < numRows; i++) { + values[i] = hostColumns[1].isNull(i) + ? null : hostColumns[1].getInt(i); + } + return new Object[] { values }; + }""" + +MICRO_EXECUTE_CPU = """\ + public static void executeCpu(Object[] data, int numRows) { + Integer[] values = (Integer[]) data[0]; + com.udf.IntegerMultiplyBy2UDF udf = new com.udf.IntegerMultiplyBy2UDF(); + for (int i = 0; i < numRows; i++) { + udf.evaluate(values[i]); + } + }""" + +MICRO_EXECUTE_GPU = """\ + public static ColumnVector executeGpu(Table table, int numRows) { + com.udf.IntegerMultiplyBy2RapidsUDF udf = new com.udf.IntegerMultiplyBy2RapidsUDF(); + return udf.evaluateColumnar(numRows, table.getColumn(1)); + }""" + +MICRO_EXECUTE_GPU_CUDA = """\ + public static ColumnVector executeGpu(Table table, int numRows) { + com.udf.IntegerMultiplyBy2NativeRapidsUDF udf = + new com.udf.IntegerMultiplyBy2NativeRapidsUDF(); + return udf.evaluateColumnar(numRows, table.getColumn(1)); + }""" diff --git a/skills/tests/test_export/scala_fixtures.py b/skills/tests/test_export/scala_fixtures.py new file mode 100644 index 00000000000..8bf96649166 --- /dev/null +++ b/skills/tests/test_export/scala_fixtures.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Scala source code fixtures for integration tests. +""" + +from .utils import replace_scala_todo_method + +NAME = "scala" +REPLACE_TODO_FN = replace_scala_todo_method +TEST_SELECTOR_FLAG = "-Dsuites" + +# --------------------------------------------------------------------------- +# UDF source code +# --------------------------------------------------------------------------- + +UDF_SOURCE = """\ +package com.udf + +class IntegerMultiplyBy2UDF extends Function1[Integer, Integer] with Serializable { + override def apply(value: Integer): Integer = { + if (value == null) null else value * 2 + } +} +""" + +RAPIDS_UDF_SOURCE = """\ +package com.udf + +import ai.rapids.cudf._ +import com.nvidia.spark.RapidsUDF +import Arm.withResource + +class IntegerMultiplyBy2RapidsUDF extends Function1[Integer, Integer] with Serializable with RapidsUDF { + override def apply(value: Integer): Integer = { + if (value == null) null else value * 2 + } + + override def evaluateColumnar(numRows: Int, args: ColumnVector*): ColumnVector = { + withResource(Scalar.fromInt(2)) { two => + args.head.mul(two) + } + } +} +""" + +SQL_SOURCE = """\ +SELECT *, + value * 2 AS result +FROM test_table +""" + +# --------------------------------------------------------------------------- +# Unit test methods +# --------------------------------------------------------------------------- + +UNIT_TEST_METHODS = { + "createTestData": """\ + def createTestData(spark: SparkSession): DataFrame = { + val schema = StructType(Seq( + StructField("id", IntegerType, nullable = false), + StructField("value", IntegerType, nullable = true) + )) + val testData = Seq( + Row(1, 123), + Row(2, 0), + Row(3, -5), + Row(4, null) + ) + spark.createDataFrame(spark.sparkContext.parallelize(testData), schema) + }""", + "registerUDF": """\ + def registerUDF(spark: SparkSession, udfName: String): Unit = { + spark.udf.register(udfName, new IntegerMultiplyBy2UDF()) + }""", + "executeUDF": """\ + def executeUDF(spark: SparkSession, udfName: String, testDF: DataFrame): DataFrame = { + testDF.createOrReplaceTempView("test_table") + spark.sql(s"SELECT *, $udfName(value) AS result FROM test_table") + }""", + "verifyUDFResults": """\ + def verifyUDFResults(resultDF: DataFrame, testDF: DataFrame): Unit = { + val results = resultDF.collect().sortBy(_.getAs[Int]("id")) + assert(results(0).getAs[Int]("result") === 246) + assert(results(1).getAs[Int]("result") === 0) + assert(results(2).getAs[Int]("result") === -10) + assert(results(3).isNullAt(results(3).fieldIndex("result"))) + }""", +} + +RAPIDS_UDF_REGISTER = """\ + def registerRapidsUDF(spark: SparkSession, udfName: String): Unit = { + spark.udf.register(udfName, new IntegerMultiplyBy2RapidsUDF()) + }""" + +NATIVE_RAPIDS_UDF_REGISTER = """\ + def registerRapidsUDF(spark: SparkSession, udfName: String): Unit = { + spark.udf.register( + udfName, + new IntegerMultiplyBy2NativeRapidsUDF(), + org.apache.spark.sql.types.IntegerType) + }""" + +# --------------------------------------------------------------------------- +# BenchUtils methods +# --------------------------------------------------------------------------- + +BENCH_GENERATE = """\ + def generateSyntheticData( + spark: SparkSession, + numRows: Long, + numPartitions: Int + ): DataFrame = { + val baseDF = spark.range(0, numRows, 1, numPartitions) + baseDF.select( + col("id"), + (rand() * 1000).cast(IntegerType).alias("value") + ) + }""" + +BENCH_CPU = """\ + def executeCpu(spark: SparkSession, df: DataFrame): DataFrame = { + import com.udf.IntegerMultiplyBy2UDF + df.createOrReplaceTempView("bench_table") + spark.udf.register("integer_multiply_by_2", new IntegerMultiplyBy2UDF()) + spark.sql("SELECT *, integer_multiply_by_2(value) AS result FROM bench_table") + }""" + +BENCH_GPU_CUDF = """\ + def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { + import com.udf.IntegerMultiplyBy2RapidsUDF + df.createOrReplaceTempView("bench_table") + spark.udf.register("integer_multiply_by_2_rapids", new IntegerMultiplyBy2RapidsUDF()) + spark.sql("SELECT *, integer_multiply_by_2_rapids(value) AS result FROM bench_table") + }""" + +BENCH_GPU_CUDA = """\ + def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { + df.createOrReplaceTempView("bench_table") + spark.udf.register( + "integer_multiply_by_2_native", + new com.udf.IntegerMultiplyBy2NativeRapidsUDF(), + org.apache.spark.sql.types.IntegerType) + spark.sql("SELECT *, integer_multiply_by_2_native(value) AS result FROM bench_table") + }""" + +BENCH_GPU_SQL = """\ + def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { + df.createOrReplaceTempView("bench_table") + val sqlContent = scala.io.Source.fromFile("src/main/resources/integer_multiply_by_2.sql").mkString + val benchSql = sqlContent.replace("test_table", "bench_table") + spark.sql(benchSql) + }""" + +# --------------------------------------------------------------------------- +# MicroBenchRunner methods +# --------------------------------------------------------------------------- + +MICRO_PREPARE_CPU = """\ + def prepareCpuData( + hostColumns: Array[HostColumnVector], + numRows: Int + ): Array[AnyRef] = { + val values = Array.tabulate(numRows) { i => + if (hostColumns(1).isNull(i)) null + else Int.box(hostColumns(1).getInt(i)) + } + Array[AnyRef](values) + }""" + +MICRO_EXECUTE_CPU = """\ + def executeCpu(data: Array[AnyRef], numRows: Int): Unit = { + val values = data(0).asInstanceOf[Array[Integer]] + val udf = new com.udf.IntegerMultiplyBy2UDF() + var i = 0 + while (i < numRows) { + udf.apply(values(i)) + i += 1 + } + }""" + +MICRO_EXECUTE_GPU = """\ + def executeGpu(table: Table, numRows: Int): ColumnVector = { + val udf = new com.udf.IntegerMultiplyBy2RapidsUDF() + udf.evaluateColumnar(numRows, table.getColumn(1)) + }""" + +MICRO_EXECUTE_GPU_CUDA = """\ + def executeGpu(table: Table, numRows: Int): ColumnVector = { + val udf = new com.udf.IntegerMultiplyBy2NativeRapidsUDF() + udf.evaluateColumnar(numRows, table.getColumn(1)) + }""" diff --git a/skills/tests/test_export/test_jvm.py b/skills/tests/test_export/test_jvm.py new file mode 100644 index 00000000000..44cade95b59 --- /dev/null +++ b/skills/tests/test_export/test_jvm.py @@ -0,0 +1,523 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Integration tests for the JVM export directories. +""" + +import os +import shutil +import stat +import tempfile +from pathlib import Path + +import pytest + +from . import cuda_fixtures, java_fixtures, scala_fixtures +from .utils import run_mvn, run_script + +pytestmark = pytest.mark.slow + +SKILLS_DIR = Path(__file__).resolve().parents[2] +TEMPLATES_DIR = SKILLS_DIR / "udf-gen-test" / "templates" +CUDA_TEMPLATES_DIR = ( + SKILLS_DIR + / "udf-convert-to-cuda" + / "templates" + / "cuda" +) + +LANG_CONFIGS = [java_fixtures, scala_fixtures] +TARGETS = ["cudf", "sql", "cuda"] +LANG_TARGET_PARAMS = [(cfg, target) for cfg in LANG_CONFIGS for target in TARGETS] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _comparison_test_methods(cfg, target): + """Return the comparison test method stubs for the given target.""" + if target == "cudf": + return {"registerRapidsUDF": cfg.RAPIDS_UDF_REGISTER} + if target == "cuda": + return {"registerRapidsUDF": cfg.NATIVE_RAPIDS_UDF_REGISTER} + # target == "sql" - no additional methods + return {} + + +def _bench_utils_methods(cfg, target): + """Return the BenchUtils method stubs for the given target.""" + methods = { + "generateSyntheticData": cfg.BENCH_GENERATE, + "executeCpu": cfg.BENCH_CPU, + } + if target == "cudf": + methods["executeGpu"] = cfg.BENCH_GPU_CUDF + elif target == "cuda": + methods["executeGpu"] = cfg.BENCH_GPU_CUDA + else: + methods["executeGpu"] = cfg.BENCH_GPU_SQL + return methods + + +def _micro_bench_methods(cfg, target): + """Return the MicroBenchRunner method stubs.""" + methods = { + "prepareCpuData": cfg.MICRO_PREPARE_CPU, + "executeCpu": cfg.MICRO_EXECUTE_CPU, + } + if target == "cuda": + methods["executeGpu"] = cfg.MICRO_EXECUTE_GPU_CUDA + else: + methods["executeGpu"] = cfg.MICRO_EXECUTE_GPU + return methods + + +def _build_project_dir(cfg): + """Copy export directory to a temp directory and resolve pom.xml.""" + export_dir = TEMPLATES_DIR / cfg.NAME + tmp_dir = tempfile.mkdtemp(prefix=f"test_{cfg.NAME}_") + project_dir = os.path.join(tmp_dir, cfg.NAME) + shutil.copytree(str(export_dir), project_dir) + + return tmp_dir, project_dir + + +def _copy_cuda_templates(project_dir): + """Copy CUDA add-on templates and replace placeholders with fixture sources.""" + shutil.copytree(str(CUDA_TEMPLATES_DIR), project_dir, dirs_exist_ok=True) + + extract_script = Path(project_dir) / "native" / "scripts" / "extract-cudf-libs.sh" + extract_script.chmod( + extract_script.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH + ) + + project_path = Path(project_dir) + for rel_path in cuda_fixtures.PLACEHOLDER_FILES: + (project_path / rel_path).unlink(missing_ok=True) + + for rel_path, source in cuda_fixtures.NATIVE_SOURCE_FILES.items(): + path = project_path / rel_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source) + + cmake_path = project_path / "native" / "src" / "main" / "cpp" / "CMakeLists.txt" + cmake = cmake_path.read_text() + cmake = cmake.replace( + """\ +set(SOURCE_FILES + "src/PlaceholderUDFNameJni.cpp" + "src/placeholder_udf_name.cu" +) +""", + cuda_fixtures.CMAKE_SOURCE_FILES, + ) + cmake_path.write_text(cmake) + + +def _fill_stubs(cfg, project_dir, target): + """Write UDF sources and fill in all TODO stubs in the project.""" + ext = f".{cfg.NAME}" + + def _replace_stubs(path, methods): + with open(path, "r") as f: + source = f.read() + for method_name, impl in methods.items(): + source = cfg.REPLACE_TODO_FN(source, method_name, impl) + with open(path, "w") as f: + f.write(source) + + src_dir = os.path.join(project_dir, "src", "main", cfg.NAME, "com", "udf") + test_dir = os.path.join(project_dir, "src", "test", cfg.NAME, "com", "udf") + + # Write CPU UDF source. + with open(os.path.join(src_dir, f"IntegerMultiplyBy2UDF{ext}"), "w") as f: + f.write(cfg.UDF_SOURCE) + + # Fill in unit test stubs. + _replace_stubs(os.path.join(test_dir, f"UnitTest{ext}"), cfg.UNIT_TEST_METHODS) + + if target == "cudf": + # Write RapidsUDF source. + with open(os.path.join(src_dir, f"IntegerMultiplyBy2RapidsUDF{ext}"), "w") as f: + f.write(cfg.RAPIDS_UDF_SOURCE) + + # Fill comparison test stubs. + _replace_stubs( + os.path.join(test_dir, f"CudfComparisonTest{ext}"), + _comparison_test_methods(cfg, "cudf"), + ) + + # Fill MicroBenchRunner stubs. + _replace_stubs( + os.path.join(src_dir, "bench", f"MicroBenchRunner{ext}"), + _micro_bench_methods(cfg, target), + ) + elif target == "cuda": + _copy_cuda_templates(project_dir) + + # Fill comparison test stubs. + _replace_stubs( + os.path.join(test_dir, f"CudfComparisonTest{ext}"), + _comparison_test_methods(cfg, "cuda"), + ) + + # Fill MicroBenchRunner stubs. + _replace_stubs( + os.path.join(src_dir, "bench", f"MicroBenchRunner{ext}"), + _micro_bench_methods(cfg, target), + ) + else: + # Write SQL file. + resources_dir = os.path.join(project_dir, "src", "main", "resources") + os.makedirs(resources_dir, exist_ok=True) + with open(os.path.join(resources_dir, "integer_multiply_by_2.sql"), "w") as f: + f.write(cfg.SQL_SOURCE) + + # Replace SQL file path placeholder in the comparison test. + sql_test_path = os.path.join(test_dir, f"SqlComparisonTest{ext}") + with open(sql_test_path, "r") as f: + content = f.read() + content = content.replace("placeholder_udf_name", "integer_multiply_by_2") + with open(sql_test_path, "w") as f: + f.write(content) + + # Fill comparison test stubs. + _replace_stubs(sql_test_path, _comparison_test_methods(cfg, "sql")) + + # Fill bench utils stubs. + _replace_stubs( + os.path.join(src_dir, "bench", f"BenchUtils{ext}"), + _bench_utils_methods(cfg, target), + ) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module", params=LANG_CONFIGS, ids=lambda c: c.NAME) +def project_dir(request): + """Clean copy of the export template with resolved pom.xml (stubs not filled).""" + cfg = request.param + tmp_dir, proj = _build_project_dir(cfg) + yield (proj, cfg) + shutil.rmtree(tmp_dir, ignore_errors=True) + + +@pytest.fixture( + scope="module", + params=LANG_TARGET_PARAMS, + ids=lambda p: f"{p[0].NAME}-{p[1]}", # "language-target" +) +def project_with_fixtures(request): + """Clean copy with resolved pom.xml and all stubs filled in.""" + cfg, target = request.param + tmp_dir, proj = _build_project_dir(cfg) + _fill_stubs(cfg, proj, target) + yield (proj, cfg, target) + shutil.rmtree(tmp_dir, ignore_errors=True) + + +@pytest.fixture( + scope="class", + params=LANG_TARGET_PARAMS, + ids=lambda p: f"{p[0].NAME}-{p[1]}", # "language-target" +) +def project_with_broken_gpu(request): + """Project with deliberately broken GPU implementation.""" + cfg, target = request.param + tmp_dir, proj = _build_project_dir(cfg) + _fill_stubs(cfg, proj, target) + + def _break_gpu_source(source: str) -> str: + # Change multiplier from 2 to 3 + source = source.replace("fromInt(2)", "fromInt(3)") + source = source.replace("* 2", "* 3") + return source + + def _insert_memory_leak(source: str) -> str: + # Inject an unclosed Scalar. + idx = source.index("evaluateColumnar") + brace = source.index("{", idx) + return ( + source[: brace + 1] + '\nScalar.fromString("LEAKED");' + source[brace + 1 :] + ) + + if target == "cudf": + path = os.path.join( + proj, + "src", + "main", + cfg.NAME, + "com", + "udf", + f"IntegerMultiplyBy2RapidsUDF.{cfg.NAME}", + ) + elif target == "cuda": + path = os.path.join( + proj, + "native", + "src", + "main", + "cpp", + "src", + "integer_multiply_by_2.cu", + ) + else: + path = os.path.join( + proj, "src", "main", "resources", "integer_multiply_by_2.sql" + ) + + # Read and overwrite with the broken source. + with open(path, "r") as f: + content = f.read() + + broken = _break_gpu_source(content) + if target == "cudf": + broken = _insert_memory_leak(broken) + + with open(path, "w") as f: + f.write(broken) + + yield (proj, cfg, target) + shutil.rmtree(tmp_dir, ignore_errors=True) + + +@pytest.fixture(scope="class", params=LANG_CONFIGS, ids=lambda c: c.NAME) +def project_with_broken_schema(request): + """Project with wrong column name in generateSyntheticData.""" + cfg = request.param + tmp_dir, proj = _build_project_dir(cfg) + _fill_stubs(cfg, proj, "cudf") + + def _break_bench_source(source: str) -> str: + # Cause a schema error due to unresolved column. + return source.replace('.alias("value")', '.alias("wrong_column")') + + ext = f".{cfg.NAME}" + bench_path = os.path.join( + proj, + "src", + "main", + cfg.NAME, + "com", + "udf", + "bench", + f"BenchUtils{ext}", + ) + + # Read and overwrite with the broken source. + with open(bench_path, "r") as f: + source = f.read() + with open(bench_path, "w") as f: + f.write(_break_bench_source(source)) + + yield (proj, cfg) + shutil.rmtree(tmp_dir, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestCompilation: + """Verify the export directory compiles.""" + + def test_compile_smoke(self, project_dir): + """ + Smoke test: compile the project as-is with TODO stubs, without + completing any of the methods, so we can catch any simple compile errors. + """ + proj, cfg = project_dir + result = run_mvn(proj, "clean", "compile") + assert result.returncode == 0, f"{cfg.NAME} smoke compile failed" + + def test_compile_with_fixtures(self, project_with_fixtures): + """ + Compile after writing UDF sources and filling in TODO stubs, to make + sure our fixtures are valid source code. + """ + proj, cfg, _target = project_with_fixtures + # test-compile compiles both main and test sources. + result = run_mvn(proj, "clean", "test-compile") + assert result.returncode == 0, f"{cfg.NAME} compile with fixtures failed" + + +class TestComparisonTest: + """Run the comparison test suite.""" + + def test_run_comparison_test(self, project_with_fixtures): + """Execute the comparison test (CudfComparisonTest or SqlComparisonTest).""" + proj, cfg, target = project_with_fixtures + if target == "sql": + suite = "com.udf.SqlComparisonTest" + else: + suite = "com.udf.CudfComparisonTest" + + result = run_mvn( + proj, + "test", + extra_args=[ + *(["-Pcuda-native-udf"] if target == "cuda" else []), + f"{cfg.TEST_SELECTOR_FLAG}={suite}", + ], + ) + assert result.returncode == 0, f"{suite} failed" + + +class TestBench: + """Test the benchmark pipeline (GenData + BenchRunner).""" + + def test_validate(self, project_with_fixtures): + """GenData validation: generate a small dataset and run validation.""" + proj, _, target = project_with_fixtures + result = run_mvn( + proj, + "compile", + "exec:java", + extra_args=[ + *(["-Pcuda-native-udf"] if target == "cuda" else []), + "-Dexec.mainClass=com.udf.bench.GenData", + "-Dexec.classpathScope=compile", + "-Dexec.args=--rows 1000 --validate --spark-conf spark.master=local[*]", + ], + ) + assert result.returncode == 0, "GenData validate failed" + + def test_spark_e2e(self, project_with_fixtures): + """End-to-end: GenData generates data, SparkBenchRunner benchmarks CPU/GPU.""" + proj, _, target = project_with_fixtures + + data_dir = os.path.join(proj, "data", "bench_input") + result_path = os.path.join(proj, "results", "bench_result.json") + mvn_args = ["--mvn-arg", "-Pcuda-native-udf"] if target == "cuda" else [] + try: + # GenData: generate parquet + gen_result = run_script( + os.path.join(proj, "run_gen_data.sh"), + args=["--rows", "1000", "--output-path", data_dir, *mvn_args], + ) + assert gen_result.returncode == 0, "run_gen_data.sh failed" + + # BenchRunner: run both benchmarks + for mode in ["cpu", "gpu"]: + bench_result = run_script( + os.path.join(proj, "run_spark_benchmark.sh"), + args=[ + "--mode", + mode, + "--data-path", + data_dir, + "--result-path", + result_path, + *mvn_args, + ], + ) + assert ( + bench_result.returncode == 0 + ), f"run_spark_benchmark.sh --mode {mode} failed" + assert os.path.isfile( + result_path + ), f"Result file not created: {result_path}" + finally: + shutil.rmtree(data_dir, ignore_errors=True) + if os.path.isfile(result_path): + os.remove(result_path) + + def test_micro_e2e(self, project_with_fixtures): + """End-to-end: GenData generates data, MicroBenchRunner benchmarks CPU/GPU.""" + proj, _, target = project_with_fixtures + if target not in {"cudf", "cuda"}: + pytest.skip("MicroBenchRunner only applies to RapidsUDF targets") + + data_dir = os.path.join(proj, "data", "micro_input") + mvn_args = ["--mvn-arg", "-Pcuda-native-udf"] if target == "cuda" else [] + try: + # GenData: generate parquet + gen_result = run_script( + os.path.join(proj, "run_gen_data.sh"), + args=["--rows", "1000", "--output-path", data_dir, *mvn_args], + ) + assert gen_result.returncode == 0, "run_gen_data.sh failed" + + # MicroBenchRunner: run both benchmarks + bench_result = run_script( + os.path.join(proj, "run_micro_benchmark.sh"), + args=["--mode", "all", "--data-path", data_dir, *mvn_args], + ) + assert bench_result.returncode == 0, ( + "run_micro_benchmark.sh failed:\n" + + bench_result.stdout + + bench_result.stderr + ) + finally: + shutil.rmtree(data_dir, ignore_errors=True) + + +class TestErrors: + """Verify that errors are caught by the test harness.""" + + def test_comparison_catches_gpu_error(self, project_with_broken_gpu): + """Comparison test should fail when GPU implementation produces wrong results.""" + proj, cfg, target = project_with_broken_gpu + if target == "sql": + suite = "com.udf.SqlComparisonTest" + else: + suite = "com.udf.CudfComparisonTest" + + result = run_mvn( + proj, + "test", + extra_args=[ + *(["-Pcuda-native-udf"] if target == "cuda" else []), + f"{cfg.TEST_SELECTOR_FLAG}={suite}", + ], + ) + assert ( + result.returncode != 0 + ), f"{suite} should have failed with broken GPU implementation" + combined = result.stdout + result.stderr + assert ( # 123 * 2 vs. 123 * 3, since we swapped multiplier + "246" in combined and "369" in combined + ), "Expected to see mismatch in test output" + + if target == "cudf": + # Re-run with debug.memory.leaks=true to verify memory leak is detected + leak_result = run_mvn( + proj, + "test", + extra_args=[ + f"{cfg.TEST_SELECTOR_FLAG}={suite}", + "-Ddebug.memory.leaks=true", + ], + ) + assert ( + leak_result.returncode != 0 + ), f"{suite} should have failed with broken GPU implementation" + leak_output = leak_result.stdout + leak_result.stderr + assert "A SCALAR WAS LEAKED" in leak_output, "Expected memory leak" + + def test_bench_validate_catches_error(self, project_with_broken_schema): + """GenData --validate should fail when synthetic data has wrong schema.""" + proj, cfg = project_with_broken_schema + result = run_mvn( + proj, + "compile", + "exec:java", + extra_args=[ + "-Dexec.mainClass=com.udf.bench.GenData", + "-Dexec.classpathScope=compile", + "-Dexec.args=--rows 1000 --validate --spark-conf spark.master=local[*]", + ], + ) + assert ( + result.returncode != 0 + ), "GenData --validate should have failed with wrong schema" + assert ( + "org.apache.spark.sql.AnalysisException" in result.stderr + ), "Expected to see schema error message" diff --git a/skills/tests/test_export/utils.py b/skills/tests/test_export/utils.py new file mode 100644 index 00000000000..d80f0e38b54 --- /dev/null +++ b/skills/tests/test_export/utils.py @@ -0,0 +1,91 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Shared test utilities for JVM skill integration tests. +""" + +import re +import subprocess +import sys +from typing import Optional + + +def run_mvn( + work_dir: str, + *goals: str, + extra_args: Optional[list[str]] = None, + timeout: int = 300, +) -> subprocess.CompletedProcess: + """Run Maven in work_dir with the given goals.""" + cmd = ["mvn", *goals, "-q"] + if extra_args: + cmd.extend(extra_args) + result = subprocess.run( + cmd, + cwd=work_dir, + capture_output=True, + text=True, + timeout=timeout, + ) + + # Write output to stdout/stderr so it is visible via pytest -s (and on errors) + if result.stdout: + sys.stdout.write(result.stdout) + if result.stderr: + sys.stderr.write(result.stderr) + return result + + +def run_script( + script_path: str, + args: Optional[list[str]] = None, + timeout: int = 300, +) -> subprocess.CompletedProcess: + """Run a bash script with the given arguments.""" + cmd = ["bash", script_path] + if args: + cmd.extend(args) + result = subprocess.run( + cmd, + capture_output=True, + text=True, + timeout=timeout, + ) + + # Write output to stdout/stderr so it is visible via pytest -s (and on errors) + if result.stdout: + sys.stdout.write(result.stdout) + if result.stderr: + sys.stderr.write(result.stderr) + return result + + +def replace_java_todo_method(source: str, method_name: str, new_body: str) -> str: + """ + Replace a TODO method in a Java source file with a real implementation. + """ + pattern = re.compile(r" public static \S+ " + re.escape(method_name) + r"\b") + match = pattern.search(source) + if not match: + raise ValueError(f"Could not find TODO method '{method_name}' in source") + + start = match.start() + brace_pos = source.index("{", start) + end_pos = source.index("}", brace_pos + 1) + 1 + return source[:start] + new_body + source[end_pos:] + + +def replace_scala_todo_method(source: str, method_name: str, new_body: str) -> str: + """ + Replace a TODO method stub in a Scala source file with a real implementation. + Assumes stubs look like "def foo(...) = ???" + """ + pattern = re.compile( + r" def " + re.escape(method_name) + r"\b.*?\?\?\?", + re.DOTALL, + ) + result = pattern.sub(new_body, source, count=1) + if result == source: + raise ValueError(f"Could not find TODO method '{method_name}' in source") + return result diff --git a/skills/tests/test_skill_frontmatter.py b/skills/tests/test_skill_frontmatter.py new file mode 100644 index 00000000000..cd2fe7c4a62 --- /dev/null +++ b/skills/tests/test_skill_frontmatter.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +""" +Tests for checked-in skill metadata. +""" + +from pathlib import Path + +import pytest +import yaml + +SKILLS_DIR = Path(__file__).resolve().parents[1] +SKILL_FILES = sorted(SKILLS_DIR.glob("*/SKILL.md")) + + +def _read_frontmatter(path: Path) -> str: + """Return the YAML frontmatter body from a SKILL.md file.""" + lines = path.read_text(encoding="utf-8").splitlines() + if not lines or lines[0] != "---": + raise ValueError(f"{path} must start with YAML frontmatter") + + try: + end = lines.index("---", 1) + except ValueError as e: + raise ValueError(f"{path} must close YAML frontmatter with ---") from e + + return "\n".join(lines[1:end]) + + +@pytest.mark.parametrize("skill_file", SKILL_FILES, ids=lambda p: p.parent.name) +def test_skill_frontmatter_loads(skill_file: Path) -> None: + frontmatter = _read_frontmatter(skill_file) + parsed = yaml.safe_load(frontmatter) + + assert isinstance(parsed, dict), f"{skill_file} frontmatter must parse to a map" + assert isinstance(parsed.get("name"), str), f"{skill_file} must define name" + assert isinstance( + parsed.get("description"), str + ), f"{skill_file} must define description" diff --git a/skills/udf-benchmark/CUDF_MICROBENCHMARKS.md b/skills/udf-benchmark/CUDF_MICROBENCHMARKS.md new file mode 100644 index 00000000000..703384979de --- /dev/null +++ b/skills/udf-benchmark/CUDF_MICROBENCHMARKS.md @@ -0,0 +1,30 @@ + + +# cuDF Microbenchmarks + +Measures fine-grained CPU vs. GPU performance without Spark overhead on in-memory data. + +## Contents +- [ ] Implement MicroBenchRunner +- [ ] Run microbenchmarks + +## Implement MicroBenchRunner + +Fill in the three TODO methods following the docstrings. + +## Run Microbenchmarks + +Generate data first (reuse from GenData output), then run: +```bash +./run_micro_benchmark.sh --mode all --data-path data/bench_data__rows.parquet --rows +``` + +Note that the specified number of rows will be coalesced into a single cuDF table. +A large table size (>1GB) will demonstrate better GPU performance. + +## Next Steps + +To profile and iteratively optimize GPU performance, use the **udf-optimize-cudf** skill. diff --git a/skills/udf-benchmark/SKILL.md b/skills/udf-benchmark/SKILL.md new file mode 100644 index 00000000000..66e5005a748 --- /dev/null +++ b/skills/udf-benchmark/SKILL.md @@ -0,0 +1,80 @@ +--- +name: udf-benchmark +description: Assists with benchmarking and profiling the performance of an Apache Spark UDF on the GPU. This is step 3 of 3 in the UDF conversion workflow (udf-gen-test -> udf-convert-to-* -> udf-benchmark). Use this skill when you have a CPU UDF and a RapidsUDF or SQL implementation, and need to benchmark the performance of the CPU UDF against the GPU implementation. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# UDF Benchmark + +## Workflow + +- [ ] Step 1: Implement BenchUtils (fill in TODO methods) +- [ ] Step 2: Validate with a small dataset +- [ ] Step 3: Generate full benchmark data and run benchmarks +- [ ] Step 4: cuDF microbenchmarks (skip for SQL targets) + +**Before making any edits, create a visible TODO checklist for every workflow step in this skill and keep it updated.** Do not produce a final answer until every required checklist item is marked complete. + +## Prerequisites + +- Project directory from Steps 1-2 (udf-gen-test, udf-convert-to-*) with passing tests + +Derive `` and `` from the UDF class name. + +> **Note:** Commands require access to `/tmp` (Spark temp storage) and `/dev` (GPU device). If commands fail due to sandbox restrictions, re-run them unsandboxed. + +## Step 1: Implement BenchUtils + +Read `src/main//com/udf/bench/BenchUtils.`. Replace placeholders with the actual camel/snake UDF name. + +Fill in the TODO methods following the docstrings. For variable-length inputs, generate sizable rows representative of enterprise-scale data. Refer to the unit test for schema and example data. + +## Step 2: Validate + +Make scripts executable: +```bash +chmod +x *.sh +``` + +Run validation mode to test with a small dataset: +```bash +./run_gen_data.sh --rows 1000 --validate +``` + +This runs both the CPU and GPU implementations on the dataset. +If validation fails, analyze the error and fix the BenchUtils implementation. + +## Step 3: Generate Data and Run Benchmarks + +### Generate benchmark data (10M rows): +```bash +./run_gen_data.sh --rows 10000000 +``` + +### Run benchmarks: +```bash +# CPU benchmark +./run_spark_benchmark.sh --mode cpu --data-path data/bench_data_10000000_rows.parquet + +# GPU benchmark +./run_spark_benchmark.sh --mode gpu --data-path data/bench_data_10000000_rows.parquet +``` + +Results are saved to the `results/` directory as JSON files. + +## Step 4: cuDF Microbenchmarks + +> Skip this step for SQL targets. This only applies to cuDF RapidsUDF conversions. + +Follow [CUDF_MICROBENCHMARKS.md](CUDF_MICROBENCHMARKS.md) to implement and run in-memory microbenchmarks. + +## Output + +Upon successful completion: +- Benchmark utilities: `src/main//com/udf/bench/BenchUtils.` +- Microbenchmarks (cuDF): `src/main//com/udf/bench/MicroBenchRunner.` +- Generated data: `data/` +- Benchmark results: `results/` diff --git a/skills/udf-convert-to-cuda/SKILL.md b/skills/udf-convert-to-cuda/SKILL.md new file mode 100644 index 00000000000..f43f7a98fb3 --- /dev/null +++ b/skills/udf-convert-to-cuda/SKILL.md @@ -0,0 +1,169 @@ +--- +name: udf-convert-to-cuda +description: Assists with converting a non-aggregating Apache Spark UDF to a native CUDA RapidsUDF using JNI and libcudf. This is step 2 of 3 in the UDF conversion workflow (udf-gen-test -> udf-convert-to-cuda -> udf-benchmark). Use this skill when you have a CPU UDF with a unit test and need to convert it to a native CUDA implementation. Prefer udf-convert-to-cudf unless a CUDA implementation is necessary for performance or correctness, or if requested by the user. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# Convert UDF to Native CUDA RapidsUDF + +## Workflow + +- [ ] Step 1: Copy CUDA add-on templates into the udf-gen-test project +- [ ] Step 2: Create the Java RapidsUDF/JNI bridge +- [ ] Step 3: Implement the CUDA/libcudf native function +- [ ] Step 4: Build with the `cuda-native-udf` Maven profile +- [ ] Step 5: Fill in the comparison test and iterate +- [ ] Step 6: Run judge subagent if requested +- [ ] Step 7: Review conversion + +**Before making any edits, create a visible TODO checklist for every workflow step in this skill and keep it updated.** Do not produce a final answer until every required checklist item is marked complete. + +## Prerequisites + +- Project directory from Step 1 (`udf-gen-test`) with a passing unit test +- Native build tools: CMake 3.30.4+, a CUDA-compatible C++ compiler, `git`, and `unzip` +- Docker is optional, but can be used for a stable native build environment + +Derive `` and `` from the UDF class name. + +> **Note:** Commands require access to `/tmp` (Spark temp storage) and `/dev` (GPU device). If commands fail due to sandbox restrictions, re-run them unsandboxed. + +## Step 1: Copy CUDA Add-On Templates + +Copy this skill's CUDA templates into the existing project: +```bash +cp -r templates/cuda/* // +chmod +x //native/scripts/extract-cudf-libs.sh +``` + +The `udf-gen-test` Maven template already contains an inactive `cuda-native-udf` profile. The native profile is activated only when you build with `-Pcuda-native-udf`. + +Read [NATIVE_BUILD_ENV.md](references/NATIVE_BUILD_ENV.md) before changing build configuration. +Read `examples/` for native RapidsUDF examples. + +## Step 2: Create the RapidsUDF/JNI Bridge + +Use `src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java` as a starting point: + +1. Rename it to `NativeRapidsUDF.java`. +2. Rename the class to `NativeRapidsUDF`. +3. Copy the original CPU UDF interface and row-by-row method onto the class. +4. Implement `evaluateColumnar` to validate column count/types and call the native method. +5. Rename the native method to a descriptive operation name, e.g. `cosineSimilarityNative`. + +For Scala projects, keep this Java wrapper under `src/main/java/com/udf/` and register it from the Scala test/project. JNI can be used from Scala, but the Java wrapper keeps native symbol names and examples simpler. +If the Java wrapper's CPU fallback needs to call a Scala object, direct references can fail before `scala-maven-plugin` compiles the Scala classes; use reflection in the row-by-row fallback only, and keep `evaluateColumnar` on the normal JNI path. + +Read [JNI_CUDA_GUIDE.md](references/JNI_CUDA_GUIDE.md) for the `evaluateColumnar` contract, type mapping, pointer ownership, `NativeDepsLoader`, and native memory rules. +**Note:** memory allocations must use the active RMM resource; avoid direct usage of ad hoc CUDA or Thrust allocators. + +## Step 3: Implement Native CUDA Code + +Rename and edit: +- `native/src/main/cpp/src/PlaceholderUDFNameJni.cpp` +- `native/src/main/cpp/src/placeholder_udf_name.cu` +- `native/src/main/cpp/src/placeholder_udf_name.hpp` + +Update `native/src/main/cpp/CMakeLists.txt` `SOURCE_FILES` to match the renamed files. If libcudf ABI/version compatibility is unclear, defer to the user. + +Read [JNI_CUDA_GUIDE.md](references/JNI_CUDA_GUIDE.md) before writing kernels. + +Verify cuDF header names before choosing includes or APIs. After dependency extraction, the active header tree will be cloned under `target/cudf-repo/cpp/include`. + +### Critical Requirements + +- **NEVER use `copyToHost()` or native methods that copy inputs from GPU to CPU.** This defeats the purpose of GPU acceleration +- **Do NOT hardcode test values.** The RapidsUDF must implement actual business logic for ANY potential input + +## Step 4: Build + +The native Maven profile uses the RAPIDS dependency already declared in `pom.xml`. + +```bash +mvn package -Pcuda-native-udf -DskipTests +``` + +To use the Docker build environment: +```bash +docker build -t cuda-udf-build . +mkdir -p "$HOME/.m2" +docker run --rm --gpus all \ + --user "$(id -u):$(id -g)" \ + -e HOME=/workspace \ + -v "$PWD":/workspace \ + -v "$HOME/.m2":/workspace/.m2 \ + -w /workspace \ + cuda-udf-build \ + -c "mvn -B -Dmaven.repo.local=/workspace/.m2/repository package -Pcuda-native-udf -DskipTests -Dnative.build.path=/workspace/target/native-build-docker" +``` + +If the build fails while resolving cuDF headers or RAPIDS CMake, check network access and the generated `cudf.git.branch` / `rapids.cmake.branch` properties. These properties may contain either a branch or a tag. + +## Step 5: Build and Test + +Fill in the target-specific TODOs in `src/test//com/udf/CudfComparisonTest.`: +- Register `NativeRapidsUDF` as the GPU implementation +- Replace placeholder UDF names + +Run: +```bash +# Java +mvn test -Dtest=CudfComparisonTest -Pcuda-native-udf + +# Scala project using a Java native RapidsUDF wrapper +mvn test -Dsuites=com.udf.CudfComparisonTest -Pcuda-native-udf +``` + +To run the tests inside the Docker build environment: + +```bash +docker run --rm --gpus all \ + --user "$(id -u):$(id -g)" \ + -e HOME=/workspace \ + -v "$PWD":/workspace \ + -v "$HOME/.m2":/workspace/.m2 \ + -v /etc/passwd:/etc/passwd:ro \ + -v /etc/group:/etc/group:ro \ + -w /workspace \ + cuda-udf-build \ + -c "mvn -B -Dmaven.repo.local=/workspace/.m2/repository test -Dtest=CudfComparisonTest -Pcuda-native-udf -Dnative.build.path=/workspace/target/native-build-docker -DskipCudfExtraction=true" +``` + +If tests fail, iterate on the Java bridge or native implementation. + +### Difficult Test Failures + +Treat the unit test as the CPU behavior specification. Do not weaken or remove test cases silently. + +- Tests that check for CPU errors may not be directly applicable to a columnar implementation: the GPU path typically evaluates a whole column and may produce nulls for invalid rows instead of throwing row-level exceptions. If this causes an unavoidable mismatch, add a clear comment in the test and a `TODO/NOTE` in the implementation explaining the mismatch. +- If a test case does not pass because of inherent CUDA/libcudf/API limitations or low-level GPU/CPU semantic differences, comment out the conflicting assertion/test only after documenting how you tried to make the behavior match and why those attempts failed. Add a note to the user. +- If the behavior is important, common, or part of the documented input domain, **always prefer fixing the implementation** over commenting out the test case. The exception is a performance-vs-correctness tradeoff that the user explicitly approves. + +## Step 6: Run Judge Subagent If Requested + +If the user explicitly asked for the judge, a judge subagent, or a review agent, treat that as an explicit request for delegation: you **MUST** launch a separate subagent with `model: inherit` and instruct it to use the **udf-judge-conversion** skill. Ask it to review the `UnitTest`, `CudfComparisonTest`, Java bridge, and JNI/CUDA sources. + +If the user did not request a judge/review agent, mark this step as skipped and continue to Step 7. If a required judge subagent is blocked by tool policy, stop and tell the user that explicit permission/instruction is needed. + +If you run the judge, wait for it to complete and review its report. If the judge finds any issues, 1) fix the issues, 2) re-run the tests, and 3) re-run the judge subagent. + +## Step 7: Review Conversion + +Review your own work to ensure: +- The test runs on the GPU and directly compares CPU-GPU outputs +- The implementation does not overfit to test cases +- No `copyToHost()` or row-by-row GPU-to-CPU copying is used for computation +- No debug statements (e.g., `TableDebug.get().debug(...)`) remain in final output + +## Output + +Upon successful completion: +- Native RapidsUDF file at `src/main/java/com/udf/NativeRapidsUDF.java` +- JNI/CUDA sources under `native/src/main/cpp/src/` +- Packaged native library in the generated UDF JAR +- Comparison test passes + +These outputs are required for **Step 3: Benchmark**. diff --git a/skills/udf-convert-to-cuda/examples/CosineSimilarityJni.cpp b/skills/udf-convert-to-cuda/examples/CosineSimilarityJni.cpp new file mode 100644 index 00000000000..39cc570d5e4 --- /dev/null +++ b/skills/udf-convert-to-cuda/examples/CosineSimilarityJni.cpp @@ -0,0 +1,67 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cosine_similarity.hpp" + +#include +#include +#include + +#include + +#include +#include + +namespace { + +constexpr char const* RUNTIME_ERROR_CLASS = "java/lang/RuntimeException"; +constexpr char const* ILLEGAL_ARG_CLASS = "java/lang/IllegalArgumentException"; + +void throw_java_exception(JNIEnv* env, char const* class_name, char const* message) +{ + jclass ex_class = env->FindClass(class_name); + if (ex_class != nullptr) { + env->ThrowNew(ex_class, message); + } +} + +} // namespace + +extern "C" { + +JNIEXPORT jlong JNICALL +Java_com_udf_CosineSimilarityNativeRapidsUDF_cosineSimilarity(JNIEnv* env, + jclass, + jlong j_view1, + jlong j_view2) +{ + try { + auto v1 = reinterpret_cast(j_view1); + auto v2 = reinterpret_cast(j_view2); + if (v1 == nullptr || v2 == nullptr) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, "input column view is null"); + return 0; + } + if (v1->type().id() != v2->type().id() || v1->type().id() != cudf::type_id::LIST) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, "inputs are not list columns"); + return 0; + } + + auto lv1 = cudf::lists_column_view(*v1); + auto lv2 = cudf::lists_column_view(*v2); + std::unique_ptr result = cosine_similarity(lv1, lv2); + return reinterpret_cast(result.release()); + } catch (std::bad_alloc const& e) { + auto message = std::string("Unable to allocate native memory: ") + e.what(); + throw_java_exception(env, RUNTIME_ERROR_CLASS, message.c_str()); + } catch (std::invalid_argument const& e) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, e.what()); + } catch (std::exception const& e) { + throw_java_exception(env, RUNTIME_ERROR_CLASS, e.what()); + } + return 0; +} + +} diff --git a/skills/udf-convert-to-cuda/examples/CosineSimilarityNativeRapidsUDF.java b/skills/udf-convert-to-cuda/examples/CosineSimilarityNativeRapidsUDF.java new file mode 100644 index 00000000000..af953a35516 --- /dev/null +++ b/skills/udf-convert-to-cuda/examples/CosineSimilarityNativeRapidsUDF.java @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import ai.rapids.cudf.ColumnVector; +import com.nvidia.spark.RapidsUDF; +import org.apache.spark.sql.api.java.UDF2; + +import scala.collection.mutable.WrappedArray; + +/** + * Native CUDA RapidsUDF example for cosine similarity over two LIST(FLOAT32) columns. + */ +public class CosineSimilarityNativeRapidsUDF + implements UDF2, WrappedArray, Float>, RapidsUDF { + @Override + public Float call(WrappedArray v1, WrappedArray v2) { + if (v1 == null || v2 == null) { + return null; + } + if (v1.length() != v2.length()) { + throw new IllegalArgumentException("Array lengths must match: " + + v1.length() + " != " + v2.length()); + } + + double dotProduct = 0; + double magnitude1 = 0; + double magnitude2 = 0; + for (int i = 0; i < v1.length(); i++) { + float f1 = v1.apply(i); + float f2 = v2.apply(i); + dotProduct += f1 * f2; + magnitude1 += f1 * f1; + magnitude2 += f2 * f2; + } + return (float) (dotProduct / (Math.sqrt(magnitude1) * Math.sqrt(magnitude2))); + } + + @Override + public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { + if (args.length != 2) { + throw new IllegalArgumentException("Unexpected argument count: " + args.length); + } + if (numRows != args[0].getRowCount() || numRows != args[1].getRowCount()) { + throw new IllegalArgumentException("Input row count mismatch"); + } + + NativeUDFLoader.ensureLoaded(); + return new ColumnVector(cosineSimilarity(args[0].getNativeView(), args[1].getNativeView())); + } + + private static native long cosineSimilarity(long vectorView1, long vectorView2); +} diff --git a/skills/udf-convert-to-cuda/examples/cosine_similarity.cu b/skills/udf-convert-to-cuda/examples/cosine_similarity.cu new file mode 100644 index 00000000000..e36e3c17cfc --- /dev/null +++ b/skills/udf-convert-to-cuda/examples/cosine_similarity.cu @@ -0,0 +1,119 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "cosine_similarity.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include + +namespace { + +struct cosine_similarity_functor { + float const* const v1; + float const* const v2; + int32_t const* const v1_offsets; + int32_t const* const v2_offsets; + + __device__ float operator()(cudf::size_type row_idx) + { + auto const v1_start_idx = v1_offsets[row_idx]; + auto const v1_num_elems = v1_offsets[row_idx + 1] - v1_start_idx; + auto const v2_start_idx = v2_offsets[row_idx]; + auto const v2_num_elems = v2_offsets[row_idx + 1] - v2_start_idx; + + double magnitude1 = 0; + double magnitude2 = 0; + double dot_product = 0; + for (auto i = 0; i < v1_num_elems; i++) { + float const f1 = v1[v1_start_idx + i]; + float const f2 = v2[v2_start_idx + i]; + magnitude1 += f1 * f1; + magnitude2 += f2 * f2; + dot_product += f1 * f2; + } + return static_cast(dot_product / (cuda::std::sqrt(magnitude1) * cuda::std::sqrt(magnitude2))); + } +}; + +} // namespace + +std::unique_ptr cosine_similarity(cudf::lists_column_view const& lv1, + cudf::lists_column_view const& lv2, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + if (!cudf::have_same_types(lv1.child(), lv2.child()) || + lv1.child().type().id() != cudf::type_id::FLOAT32) { + throw std::invalid_argument("inputs are not lists of floats"); + } + + auto const row_count = lv1.size(); + if (row_count != lv2.size()) { + throw std::invalid_argument("input row counts do not match"); + } + if (row_count == 0) { + return cudf::make_empty_column(cudf::data_type{cudf::type_id::FLOAT32}); + } + if (lv1.child().null_count() != 0 || lv2.child().null_count() != 0) { + throw std::invalid_argument("null floats are not supported"); + } + + auto const lv1_offsets_ptr = lv1.offsets().data(); + auto const lv2_offsets_ptr = lv2.offsets().data(); + auto const lv1_null_mask = lv1.parent().null_mask(); + auto const lv2_null_mask = lv2.parent().null_mask(); + + bool const are_offsets_equal = + thrust::all_of(rmm::exec_policy_nosync(stream), + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(row_count), + [lv1_offsets_ptr, lv2_offsets_ptr, lv1_null_mask, lv2_null_mask] + __device__(cudf::size_type idx) -> bool { + bool const lv1_is_null = + lv1_null_mask != nullptr && !cudf::bit_is_set(lv1_null_mask, idx); + bool const lv2_is_null = + lv2_null_mask != nullptr && !cudf::bit_is_set(lv2_null_mask, idx); + if (lv1_is_null || lv2_is_null) { + return true; + } + return (lv1_offsets_ptr[idx + 1] - lv1_offsets_ptr[idx]) == + (lv2_offsets_ptr[idx + 1] - lv2_offsets_ptr[idx]); + }); + if (!are_offsets_equal) { + throw std::invalid_argument("input list lengths do not match for every row"); + } + + rmm::device_uvector float_results(row_count, stream, mr); + thrust::transform(rmm::exec_policy_nosync(stream), + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(row_count), + float_results.data(), + cosine_similarity_functor({lv1.child().data(), + lv2.child().data(), + lv1.offsets().data(), + lv2.offsets().data()})); + + auto [null_mask, null_count] = + cudf::bitmask_and(cudf::table_view({lv1.parent(), lv2.parent()}), stream, mr); + return std::make_unique(cudf::data_type{cudf::type_id::FLOAT32}, + row_count, + float_results.release(), + std::move(null_mask), + null_count); +} diff --git a/skills/udf-convert-to-cuda/examples/cosine_similarity.hpp b/skills/udf-convert-to-cuda/examples/cosine_similarity.hpp new file mode 100644 index 00000000000..99b78ede0f7 --- /dev/null +++ b/skills/udf-convert-to-cuda/examples/cosine_similarity.hpp @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include + +#include + +std::unique_ptr cosine_similarity( + cudf::lists_column_view const& lv1, + cudf::lists_column_view const& lv2, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); diff --git a/skills/udf-convert-to-cuda/references/JNI_CUDA_GUIDE.md b/skills/udf-convert-to-cuda/references/JNI_CUDA_GUIDE.md new file mode 100644 index 00000000000..2137b79837b --- /dev/null +++ b/skills/udf-convert-to-cuda/references/JNI_CUDA_GUIDE.md @@ -0,0 +1,162 @@ + + +# JNI and CUDA RapidsUDF Guide + +## RapidsUDF Contract + +The RapidsUDF interface provides a way to run a CPU UDF on the GPU when using the RAPIDS Accelerator for Apache Spark. The interface provides a single method you need to override called `evaluateColumnar`. The CPU UDF method must remain on the native RapidsUDF class so Spark can fall back to the CPU if a surrounding plan cannot run on the GPU. + +`evaluateColumnar(int numRows, ColumnVector... args)` receives columnar forms of the same inputs as the CPU UDF. All input columns should have `numRows` rows. Scalar inputs may be expanded into full columns by the RAPIDS Accelerator, so do not rely on detecting scalar-vs-column input. + +The returned `ColumnVector` must have `numRows` rows and a cuDF type that matches the Spark return type: + +| Spark Type | cuDF Type | +|---|---| +| BooleanType | BOOL8 | +| ByteType | INT8 | +| ShortType | INT16 | +| IntegerType | INT32 | +| LongType | INT64 | +| FloatType | FLOAT32 | +| DoubleType | FLOAT64 | +| DecimalType | DECIMAL32, DECIMAL64, DECIMAL128 * | +| DateType | TIMESTAMP_DAYS | +| TimestampType | TIMESTAMP_MICROSECONDS | +| StringType | STRING | +| ArrayType | LIST of element type | +| MapType | LIST of STRUCT(key, value) | +| StructType | STRUCT of fields | + +For example, if the CPU UDF returns the Spark type ArrayType(MapType(StringType, StringType)) then evaluateColumnar must return a column of type LIST(LIST(STRUCT(STRING,STRING))). + +*Note: cuDF's DECIMAL32 corresponds to precision <= 9 digits, DECIMAL64 corresponds to 9 < precision <= 18 digits, and DECIMAL128 corresponds to 18 < precision <= 38 digits. Precision greater than 38 digits is unsupported. +Note that cuDF decimals use a negative scale relative to Spark DecimalType. For example, Spark DecimalType(precision=11, scale=2) would translate to cuDF type DECIMAL64(scale=-2). + +For `ArrayType(elementType, containsNull)`, the LIST parent null mask represents null arrays. Child nulls represent null array elements and must match the `containsNull` contract. Either preserve child nulls deliberately or reject them explicitly. + +## Java Wrapper + +Use `NativeDepsLoader.loadNativeDeps(new String[] {"rapidsudfjni"})` from a synchronized loader. Call it from `evaluateColumnar`, not a static initializer, because the Spark driver may not have the executor CUDA runtime. + +Pass input columns to JNI with `ColumnVector.getNativeView()`. Wrap the native result with `new ColumnVector(nativeHandle)`. + +Do not close input `ColumnVector`s. The RAPIDS Accelerator owns them. Closing inputs can cause double-close errors. + +## JNI and Native Ownership + +JNI arguments are non-owning pointers: +```cpp +auto input = reinterpret_cast(j_input); +``` + +The native function must allocate and return an owning `cudf::column`: +```cpp +std::unique_ptr result = compute(*input); +return reinterpret_cast(result.release()); +``` + +Never return a pointer to an input view, child view, stack object, or a column owned by a temporary that will be destroyed before Java wraps it. + +Catch `std::bad_alloc`, `std::invalid_argument`, and `std::exception`, then throw Java exceptions with `JNIEnv::ThrowNew`. + +## CUDA/libcudf Implementation + +Start with libcudf column APIs before writing custom kernels. Use custom CUDA kernels when the operation requires fused logic, custom reductions, or logic unavailable in cuDF Java/libcudf primitives. + +### Checklist + +- Validate input types and row counts in Java before crossing JNI when possible +- Validate libcudf types again in JNI for native safety +- Preserve Spark null semantics +- Prefer `cudf::column_view`/`cudf::lists_column_view` for input views +- Return `std::unique_ptr` +- Avoid host copies in the final implementation +- Prefer public libcudf APIs; avoid using `cudf::detail` +- Keep one native function focused on one UDF operation + +### Correctness Pitfalls + +- **Null values of fixed-width columns are undefined memory.** Check the null mask (`cudf::bit_is_set(...)` or `column_device_view::is_valid(...)`) before reading element values. +- **Empty list/string columns have no offsets.** Accessing the offsets child of an empty list or string column is undefined behavior. Handle the empty case early (e.g., return `cudf::make_empty_column(...)`). +- **Use `cudf::have_same_types(a, b)` for type comparison**, not `a.type() == b.type()` — equality misses differences such as decimal scale. +- **`cudf::size_type` is `int32_t`. LIST offsets are always `int32_t`.** String offsets may be `int32_t` or `int64_t` for large strings. +- **Nested column null masks must agree across levels.** When constructing LIST/STRUCT output yourself, ensure parent and child null masks are consistent. +- **`CUDF_EXPECTS` conditions must be pure predicates** — side effects inside the condition may only execute in debug builds. + +### Useful Patterns + +- `rmm::device_uvector`: temporary device output buffers that can be released into a `cudf::column` +- `rmm::exec_policy_nosync(stream)`: pass the intended CUDA stream to Thrust algorithms (prefer the `_nosync` variant unless you need an implicit host-device sync) +- `cudf::make_empty_column(...)`: return correctly typed empty outputs +- `cudf::make_numeric_column(...)`: allocate fixed-width output columns with a null mask +- `cudf::bitmask_and(cudf::table_view({...}))`: combine input validity masks for output null semantics +- `cudf::lists_column_view`: inspect list offsets, child columns, parent null masks, and nested list shapes +- `cudf::strings_column_view`: inspect string chars/offsets when implementing string kernels +- `cudf::create_null_mask(...)`: create all-valid, all-null, or uninitialized masks for new outputs +- CUB and Thrust APIs: useful for scans, reductions, transforms, selection, and sorting when libcudf does not provide the exact operation + +### Memory Allocation + +- All device allocations must go through the active RMM memory resource. +- Use libcudf factories or RMM types such as `rmm::device_uvector` and `rmm::device_buffer`; avoid direct calls to `cudaMalloc`, `cudaMallocAsync`, or other ad hoc device allocators. +- Use the output MR for returned columns when the API exposes one; use `cudf::get_current_device_resource_ref()` for short-lived temporary buffers. +- Use RMM pinned memory for large host buffers. Small CPU-only metadata may use normal C++ containers. + +Example allocating CUB scratch buffers through RMM: + +```cpp +size_t temp_storage_bytes = 0; +cub::DeviceScan::InclusiveSum(nullptr, temp_storage_bytes, in, out, n, stream.value()); +rmm::device_buffer temp_storage(temp_storage_bytes, stream, cudf::get_current_device_resource_ref()); +cub::DeviceScan::InclusiveSum(temp_storage.data(), temp_storage_bytes, in, out, n, stream.value()); +``` + +### Stream and MR Plumbing + +Top-level native functions should accept stream and MR as the last two parameters, with defaults: + +```cpp +std::unique_ptr my_op( + cudf::column_view const& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); +``` + +Use the passed-in `mr` for the returned column and `cudf::get_current_device_resource_ref()` for short-lived temporaries. Propagate `stream` to every libcudf call, Thrust call, and kernel launch — do not introduce `rmm::cuda_stream_default` inside the implementation. + +### Kernel Launch Discipline + +Always check kernel launches; silent launch failures cause downstream corruption. + +```cpp +my_kernel<<>>(args); +CUDF_CHECK_CUDA(stream.value()); +``` + +Prefer `cuda::std::` (e.g. `cuda::std::min`, `cuda::std::sqrt`, `cuda::std::numeric_limits`) over `std::` inside `__device__` and `CUDF_HOST_DEVICE` code. + +Avoid synchronizing in the hot path except when required to fetch output sizes or while debugging. + +### Output Construction + +For variable-size list outputs: +1. Compute per-row child sizes on device, using zero for null parent rows. +2. Prefix-sum sizes into an `INT32` offsets column of length `numRows + 1`. +3. Allocate the child column from the final offset, fill it on device, and set child nulls if `containsNull=true`. +4. Assemble the LIST column from offsets, child column, parent null mask, and parent null count. + +For string outputs, construct proper offsets, chars, and null masks. For scalar numeric outputs, prefer libcudf transforms/reductions where possible. + +## Debugging + +Rerun tests with `-Ddebug.memory.leaks=true` to enable Java refcount debugging; this catches leaked `ColumnVector`, `Table`, `Scalar`, and Java-owned buffer objects. +Note that it does **not** catch native memory leaks; use RMM RAII patterns to ensure all native allocations are freed. + +For native kernel memory errors, run the comparison test under Compute Sanitizer: + +```bash +compute-sanitizer --tool memcheck mvn test +``` diff --git a/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md b/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md new file mode 100644 index 00000000000..a45d912d9a5 --- /dev/null +++ b/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md @@ -0,0 +1,92 @@ + + +# Native CUDA UDF Build Environment + +## Dependency Model + +The native build uses the RAPIDS JAR already resolved by Maven. The `cuda-native-udf` profile asks Maven to copy `rapids-4-spark_--.jar` and `rapids-4-spark_-.jar` into `target/rapids-jar`. The `native/scripts/extract-cudf-libs.sh` script then extracts `libcudf.so*` and `libnvcomp.so*`, clones matching cuDF headers, builds `librapidsudfjni.so`, and packages it in the UDF JAR for `NativeDepsLoader`. + +No separate manual JAR download is required. Maven should resolve the RAPIDS dependency declared in `pom.xml`; the native profile reuses the same coordinates and copies the resolved JAR into `target/rapids-jar`. + +The profile first tries the CUDA-classified artifact (`-cuda12`) and then the unclassified artifact. If extraction fails, the selected JAR probably does not contain Linux native CUDA libraries or the Maven cache/repository is inconsistent with the generated version properties. + +## Required Tools + +- CUDA toolkit matching spark-rapids build and a compatible NVIDIA driver +- CMake 3.30.4+ +- C++ compiler compatible with the selected CUDA toolkit +- JDK 17 +- Maven +- `git` +- `unzip` + +## CUDA Toolkit Version + +The native build compiles against the prebuilt libcudf in the spark-rapids jar, so the local CUDA toolkit must match the version spark-rapids was built against. + +1. Get the CUDA version(s) spark-rapids is built against: + +```bash +curl -fsSL https://nvidia.github.io/spark-rapids/docs/download.html \ + | perl -0777 -ne 'while (/built against CUDA\s+(\d+\.\d+)(?:\s+or\s+CUDA\s+(\d+\.\d+))?/g) { print "$1\n"; print "$2\n" if defined $2 }' +``` + +2. Check the active toolkit (`nvcc --version`). CMake uses `$CUDACXX`, else `nvcc` on `PATH`, else `$CUDAToolkit_ROOT/bin/nvcc` — the default `PATH` `nvcc` may not be the one you want. + +3. If it doesn't match, point the build at a matching toolkit that's already installed; otherwise install one that matches: + +```bash +export CUDACXX=/usr/local/cuda-/bin/nvcc +export CUDAToolkit_ROOT=/usr/local/cuda- +export PATH="$CUDAToolkit_ROOT/bin:$PATH" +``` + +Docker is optional. Use it when local compiler/CMake/CUDA versions drift or when the build needs to be reproducible across machines. + +The provided Dockerfile installs JDK 17 and sets it via `/etc/profile.d/java17.sh`. If a modified Dockerfile or alternate entrypoint bypasses the login shell and `mvn` reports Java 8, export `JAVA_HOME=/usr/lib/jvm/java-17-openjdk` and prepend `$JAVA_HOME/bin` to `PATH` explicitly. + +Use the full Docker command listed in SKILL.md. It runs as the calling user to avoid root-owned artifacts, mounts the project and Maven cache, and uses a Docker-specific native build path so CMake cache paths do not conflict with host builds. + +If a previous root container run already wrote `target/` artifacts, fix ownership or clean them before rerunning as a non-root user. + +CMake stores absolute source and build paths in `CMakeCache.txt`. A host-generated `target/native-build` cannot be reused from `/workspace/target/native-build` inside Docker. Use `mvn clean`, remove the stale native build directory, or pass a Docker-specific path such as `-Dnative.build.path=/workspace/target/native-build-docker`. + +## Version Alignment + +Keep these values aligned: +- Spark version +- Scala binary version +- `rapids4spark.version` +- `cuda.version` +- `cudf.git.branch` +- `rapids.cmake.branch` +- JDK version + +The generated template maps RAPIDS `..` to the `v..00` cuDF and rapids-cmake tags. If building a snapshot, a custom RAPIDS JAR, or a patch release with known native ABI changes, verify the matching cuDF/RMM/CCCL versions with the user. + +## Fast Rebuilds and Verification + +After the first successful extraction, use `-DskipCudfExtraction=true` while iterating on Java/JNI/CUDA source: + +```bash +mvn package -Pcuda-native-udf -DskipCudfExtraction=true -DskipTests +``` + +Verify deployable packaging with: + +```bash +jar tf target/*.jar | grep librapidsudfjni.so +``` + +## Build Modes + +Default: `USE_PREBUILT_CUDF=ON`. + +This extracts `libcudf` from the RAPIDS JAR and builds only the UDF JNI/CUDA library. This is the stable, fast path. + +Escape hatch: `-DUSE_PREBUILT_CUDF=OFF`. + +This builds cuDF from source through RAPIDS CMake/CPM. It is slow and more sensitive to branch drift; ask the user before using it. diff --git a/skills/udf-convert-to-cuda/templates/cuda/Dockerfile b/skills/udf-convert-to-cuda/templates/cuda/Dockerfile new file mode 100644 index 00000000000..320dc2f1423 --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/Dockerfile @@ -0,0 +1,65 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Reproducible build image for native CUDA RapidsUDF code. +ARG CUDA_VERSION=12.8.0 +ARG LINUX_VERSION=rockylinux8 + +FROM nvidia/cuda:${CUDA_VERSION}-devel-${LINUX_VERSION} + +ARG TOOLSET_VERSION=14 +ARG CMAKE_VERSION=3.30.4 +ARG CMAKE_ARCH=x86_64 +ARG CCACHE_VERSION=4.11.2 +ARG PARALLEL_LEVEL=10 + +ENV TOOLSET_VERSION=${TOOLSET_VERSION} +ENV PARALLEL_LEVEL=${PARALLEL_LEVEL} +ENV JAVA_HOME=/usr/lib/jvm/java-17-openjdk + +RUN dnf --enablerepo=powertools install -y \ + gcc-toolset-${TOOLSET_VERSION} \ + git \ + java-17-openjdk-devel \ + maven \ + ninja-build \ + patch \ + python39 \ + scl-utils \ + tar \ + unzip \ + wget \ + zlib-devel \ + && alternatives --set python /usr/bin/python3 + +RUN cd /usr/local && \ + wget --quiet https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}.tar.gz && \ + tar zxf cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}.tar.gz && \ + rm cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}.tar.gz +ENV PATH=${JAVA_HOME}/bin:/usr/local/cmake-${CMAKE_VERSION}-linux-${CMAKE_ARCH}/bin:${PATH} + +# Bake the SCL activation and Java 17 environment into /etc/profile.d so they are restored by `bash -l` on every container start. +RUN printf 'source /opt/rh/gcc-toolset-%s/enable\n' "${TOOLSET_VERSION}" \ + > /etc/profile.d/scl-gcc-toolset.sh && \ + printf '%s\n%s\n' \ + 'export JAVA_HOME=/usr/lib/jvm/java-17-openjdk' \ + 'export PATH=$JAVA_HOME/bin:$PATH' \ + > /etc/profile.d/java17.sh + +RUN cd /tmp && \ + wget --quiet https://github.com/ccache/ccache/releases/download/v${CCACHE_VERSION}/ccache-${CCACHE_VERSION}.tar.gz && \ + tar zxf ccache-${CCACHE_VERSION}.tar.gz && \ + rm ccache-${CCACHE_VERSION}.tar.gz && \ + cd ccache-${CCACHE_VERSION} && \ + mkdir build && \ + cd build && \ + scl enable gcc-toolset-${TOOLSET_VERSION} \ + "cmake .. \ + -DCMAKE_BUILD_TYPE=Release \ + -DZSTD_FROM_INTERNET=ON \ + -DREDIS_STORAGE_BACKEND=OFF && \ + cmake --build . --parallel ${PARALLEL_LEVEL} --target install" && \ + cd ../.. && \ + rm -rf ccache-${CCACHE_VERSION} + +ENTRYPOINT ["bash", "-l"] diff --git a/skills/udf-convert-to-cuda/templates/cuda/native/scripts/extract-cudf-libs.sh b/skills/udf-convert-to-cuda/templates/cuda/native/scripts/extract-cudf-libs.sh new file mode 100644 index 00000000000..52020e71c3c --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/native/scripts/extract-cudf-libs.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(cd "${SCRIPT_DIR}/../.." && pwd)" +TARGET_DIR="${TARGET_DIR:-${PROJECT_DIR}/target}" +NATIVE_DEPS_DIR="${TARGET_DIR}/native-deps" +CUDF_REPO_DIR="${TARGET_DIR}/cudf-repo" +RAPIDS_JAR_DIR="${TARGET_DIR}/rapids-jar" + +SCALA_VERSION="${SCALA_VERSION:-2.12}" +RAPIDS4SPARK_VERSION="${RAPIDS4SPARK_VERSION:-26.04.0}" +CUDA_VERSION="${CUDA_VERSION:-cuda12}" +CUDF_BRANCH="${CUDF_BRANCH:-v26.04.00}" + +mkdir -p "${NATIVE_DEPS_DIR}" "${CUDF_REPO_DIR}" + +choose_rapids_jar() { + local candidates=( + "${RAPIDS_JAR_DIR}/rapids-4-spark_${SCALA_VERSION}-${RAPIDS4SPARK_VERSION}-${CUDA_VERSION}.jar" + "${RAPIDS_JAR_DIR}/rapids-4-spark_${SCALA_VERSION}-${RAPIDS4SPARK_VERSION}.jar" + "${HOME}/.m2/repository/com/nvidia/rapids-4-spark_${SCALA_VERSION}/${RAPIDS4SPARK_VERSION}/rapids-4-spark_${SCALA_VERSION}-${RAPIDS4SPARK_VERSION}-${CUDA_VERSION}.jar" + "${HOME}/.m2/repository/com/nvidia/rapids-4-spark_${SCALA_VERSION}/${RAPIDS4SPARK_VERSION}/rapids-4-spark_${SCALA_VERSION}-${RAPIDS4SPARK_VERSION}.jar" + ) + + for candidate in "${candidates[@]}"; do + if [[ -f "${candidate}" ]]; then + echo "${candidate}" + return 0 + fi + done + + echo "ERROR: Could not find a rapids-4-spark jar." >&2 + echo "Tried target/rapids-jar and ~/.m2 for version ${RAPIDS4SPARK_VERSION} (${CUDA_VERSION})." >&2 + echo "Run the build through Maven with -Pcuda-native-udf so the profile can copy the RAPIDS dependency first." >&2 + return 1 +} + +JAR_PATH="$(choose_rapids_jar)" + +echo "Using RAPIDS jar: ${JAR_PATH}" +echo "Using cuDF header ref: ${CUDF_BRANCH}" + +TEMP_DIR="${TARGET_DIR}/cudf-extract" +rm -rf "${TEMP_DIR}" +mkdir -p "${TEMP_DIR}" + +if ! unzip -o "${JAR_PATH}" "*/libcudf.so*" "*/libnvcomp.so*" -d "${TEMP_DIR}"; then + echo "ERROR: Failed to extract libcudf/libnvcomp from ${JAR_PATH}" >&2 + echo "The selected RAPIDS jar may not include native Linux CUDA libraries." >&2 + rm -rf "${TEMP_DIR}" + exit 1 +fi + +while IFS= read -r source_file; do + cp -f "${source_file}" "${NATIVE_DEPS_DIR}/$(basename "${source_file}")" +done < <(find "${TEMP_DIR}" -name "*.so*") +rm -rf "${TEMP_DIR}" + +if [[ ! -f "${NATIVE_DEPS_DIR}/libcudf.so" ]]; then + echo "ERROR: libcudf.so was not extracted into ${NATIVE_DEPS_DIR}" >&2 + exit 1 +fi + +if [[ ! -d "${CUDF_REPO_DIR}/.git" ]]; then + git clone --depth 1 --branch "${CUDF_BRANCH}" https://github.com/rapidsai/cudf.git "${CUDF_REPO_DIR}" +else + echo "Using existing cuDF headers at ${CUDF_REPO_DIR}" +fi + +if [[ ! -d "${CUDF_REPO_DIR}/cpp/include" ]]; then + echo "ERROR: cuDF headers not found at ${CUDF_REPO_DIR}/cpp/include" >&2 + exit 1 +fi + +echo "Native dependencies ready:" +echo " Libraries: ${NATIVE_DEPS_DIR}" +echo " Headers: ${CUDF_REPO_DIR}/cpp/include" diff --git a/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/CMakeLists.txt b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/CMakeLists.txt new file mode 100644 index 00000000000..a9398b4938c --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/CMakeLists.txt @@ -0,0 +1,119 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +cmake_minimum_required(VERSION 3.30.4 FATAL_ERROR) + +set(RAPIDS_CMAKE_BRANCH "v26.04.00" CACHE STRING "rapids-cmake branch or tag") +if(RAPIDS_CMAKE_BRANCH MATCHES "^v(.+)") + set(rapids-cmake-version "${CMAKE_MATCH_1}") + set(rapids-cmake-tag "${RAPIDS_CMAKE_BRANCH}") +else() + set(rapids-cmake-branch "${RAPIDS_CMAKE_BRANCH}") +endif() +set(NATIVE_LIBRARY_NAME "rapidsudfjni" CACHE STRING "JNI shared library target name") +set(NATIVE_DEPS_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../target/native-deps" CACHE PATH "Directory containing prebuilt libcudf") +set(CUDF_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../../target/cudf-repo/cpp" CACHE PATH "cuDF source directory for headers") +set(GPU_ARCHS "RAPIDS" CACHE STRING "CUDA architectures") + +file(DOWNLOAD + https://raw.githubusercontent.com/rapidsai/rapids-cmake/${RAPIDS_CMAKE_BRANCH}/RAPIDS.cmake + ${CMAKE_BINARY_DIR}/RAPIDS.cmake +) +include(${CMAKE_BINARY_DIR}/RAPIDS.cmake) + +include(rapids-cmake) +include(rapids-cpm) +include(rapids-cuda) + +if(DEFINED ENV{CXX} AND NOT "$ENV{CXX}" STREQUAL "") + set(CMAKE_CXX_COMPILER "$ENV{CXX}" CACHE FILEPATH "C++ compiler" FORCE) +endif() + +if(DEFINED GPU_ARCHS) + set(CMAKE_CUDA_ARCHITECTURES "${GPU_ARCHS}") +endif() +rapids_cuda_init_architectures(RAPIDSUDFJNI) + +project(RAPIDSUDFJNI VERSION 26.04.0 LANGUAGES C CXX CUDA) + +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) +set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_STANDARD 20) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) +set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -w --expt-extended-lambda --expt-relaxed-constexpr") + +option(USE_PREBUILT_CUDF "Use libcudf extracted from the rapids-4-spark jar" ON) +option(PER_THREAD_DEFAULT_STREAM "Build with per-thread default stream" ON) +option(CUDF_ENABLE_ARROW_S3 "Enable Arrow S3 support in source-build mode" OFF) + +if(USE_PREBUILT_CUDF) + if(NOT EXISTS "${NATIVE_DEPS_DIR}") + message(FATAL_ERROR "NATIVE_DEPS_DIR does not exist: ${NATIVE_DEPS_DIR}") + endif() + if(NOT EXISTS "${CUDF_SOURCE_DIR}/include") + message(FATAL_ERROR "CUDF_SOURCE_DIR headers not found: ${CUDF_SOURCE_DIR}/include") + endif() + + find_library(CUDF_LIBRARY NAMES cudf PATHS "${NATIVE_DEPS_DIR}" NO_DEFAULT_PATH REQUIRED) + + get_property(rapids-cmake-dir GLOBAL PROPERTY rapids-cmake-dir) + if(NOT rapids-cmake-dir) + set(rapids-cmake-dir "${CMAKE_BINARY_DIR}/_deps/rapids-cmake-src") + endif() + + rapids_cpm_init() + include("${rapids-cmake-dir}/cpm/cccl.cmake") + rapids_cpm_cccl() + include("${rapids-cmake-dir}/cpm/rmm.cmake") + rapids_cpm_rmm() + + if(NOT TARGET rmm::rmm) + message(FATAL_ERROR "rmm::rmm target was not created") + endif() + + get_target_property(RMM_INCLUDE_DIRS rmm::rmm INTERFACE_INCLUDE_DIRECTORIES) + + add_library(cudf_imported SHARED IMPORTED GLOBAL) + set_target_properties(cudf_imported PROPERTIES IMPORTED_LOCATION "${CUDF_LIBRARY}") + target_include_directories(cudf_imported INTERFACE + "${CUDF_SOURCE_DIR}/include" + ${RMM_INCLUDE_DIRS} + ) + target_link_libraries(cudf_imported INTERFACE rmm::rmm) + add_library(cudf::cudf ALIAS cudf_imported) +else() + rapids_cpm_init() + rapids_cpm_find(cudf 26.04.00 + CPM_ARGS + GIT_REPOSITORY https://github.com/rapidsai/cudf.git + GIT_TAG ${RAPIDS_CMAKE_BRANCH} + GIT_SHALLOW TRUE + SOURCE_SUBDIR cpp + OPTIONS "BUILD_TESTS OFF" + "BUILD_BENCHMARKS OFF" + "CUDF_ENABLE_ARROW_S3 ${CUDF_ENABLE_ARROW_S3}" + "CUDF_KVIKIO_REMOTE_IO OFF" + "DISABLE_DEPRECATION_WARNING ON" + "AUTO_DETECT_CUDA_ARCHITECTURES OFF" + ) +endif() + +find_package(JNI REQUIRED) + +set(SOURCE_FILES + "src/PlaceholderUDFNameJni.cpp" + "src/placeholder_udf_name.cu" +) + +add_library(${NATIVE_LIBRARY_NAME} SHARED ${SOURCE_FILES}) +set_target_properties(${NATIVE_LIBRARY_NAME} PROPERTIES BUILD_RPATH "\$ORIGIN") + +if(PER_THREAD_DEFAULT_STREAM) + target_compile_definitions(${NATIVE_LIBRARY_NAME} PRIVATE CUDA_API_PER_THREAD_DEFAULT_STREAM) +endif() + +target_include_directories(${NATIVE_LIBRARY_NAME} PRIVATE ${JNI_INCLUDE_DIRS}) +target_compile_definitions(${NATIVE_LIBRARY_NAME} PUBLIC SPDLOG_ACTIVE_LEVEL=SPDLOG_LEVEL_OFF) +target_link_libraries(${NATIVE_LIBRARY_NAME} cudf::cudf) diff --git a/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/PlaceholderUDFNameJni.cpp b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/PlaceholderUDFNameJni.cpp new file mode 100644 index 00000000000..1b07459b9fd --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/PlaceholderUDFNameJni.cpp @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "placeholder_udf_name.hpp" + +#include +#include + +#include + +#include +#include + +namespace { + +constexpr char const* RUNTIME_ERROR_CLASS = "java/lang/RuntimeException"; +constexpr char const* ILLEGAL_ARG_CLASS = "java/lang/IllegalArgumentException"; + +void throw_java_exception(JNIEnv* env, char const* class_name, char const* message) +{ + jclass ex_class = env->FindClass(class_name); + if (ex_class != nullptr) { + env->ThrowNew(ex_class, message); + } +} + +} // namespace + +extern "C" { + +JNIEXPORT jlong JNICALL +Java_com_udf_PlaceholderUDFNameNativeRapidsUDF_evaluateNative(JNIEnv* env, + jclass, + jlong input_view) +{ + try { + auto input = reinterpret_cast(input_view); + if (input == nullptr) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, "input column view is null"); + return 0; + } + + std::unique_ptr result = placeholder_udf_name(*input); + return reinterpret_cast(result.release()); + } catch (std::bad_alloc const& e) { + auto message = std::string("Unable to allocate native memory: ") + e.what(); + throw_java_exception(env, RUNTIME_ERROR_CLASS, message.c_str()); + } catch (std::invalid_argument const& e) { + throw_java_exception(env, ILLEGAL_ARG_CLASS, e.what()); + } catch (std::exception const& e) { + throw_java_exception(env, RUNTIME_ERROR_CLASS, e.what()); + } + return 0; +} + +} diff --git a/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.cu b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.cu new file mode 100644 index 00000000000..5e0de8f20ff --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.cu @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "placeholder_udf_name.hpp" + +#include +#include +#include + +std::unique_ptr placeholder_udf_name(cudf::column_view const& input) +{ + // TODO: Replace this placeholder with the actual CUDA/libcudf implementation. + auto null_mask = cudf::create_null_mask(input.size(), cudf::mask_state::ALL_NULL); + return cudf::make_numeric_column( + cudf::data_type{cudf::type_id::INT32}, input.size(), std::move(null_mask), input.size()); +} diff --git a/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.hpp b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.hpp new file mode 100644 index 00000000000..d34ac2f8828 --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/native/src/main/cpp/src/placeholder_udf_name.hpp @@ -0,0 +1,13 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include + +std::unique_ptr placeholder_udf_name(cudf::column_view const& input); diff --git a/skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/NativeUDFLoader.java b/skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/NativeUDFLoader.java new file mode 100644 index 00000000000..d5469882951 --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/NativeUDFLoader.java @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import ai.rapids.cudf.NativeDepsLoader; + +import java.io.IOException; + +/** Loads JNI libraries packaged in this UDF jar. */ +public final class NativeUDFLoader { + private static boolean loaded; + + private NativeUDFLoader() { + } + + public static synchronized void ensureLoaded() { + if (!loaded) { + try { + NativeDepsLoader.loadNativeDeps(new String[] {"rapidsudfjni"}); + loaded = true; + } catch (IOException e) { + throw new RuntimeException("Failed to load native CUDA UDF library", e); + } + } + } +} diff --git a/skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java b/skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java new file mode 100644 index 00000000000..8212de0b399 --- /dev/null +++ b/skills/udf-convert-to-cuda/templates/cuda/src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java @@ -0,0 +1,45 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import ai.rapids.cudf.ColumnVector; +import com.nvidia.spark.RapidsUDF; +// TODO: add imports for CPU UDF's base type, e.g.: +// import org.apache.hadoop.hive.ql.exec.UDF; +// import org.apache.spark.sql.api.java.UDFn; + +/** + * Template for a native CUDA RapidsUDF. + * + * 1. Rename this class and file to {@code NativeRapidsUDF}. + * 2. Match the CPU UDF's Spark contract: + * - Hive UDF : add {@code extends org.apache.hadoop.hive.ql.exec.UDF} + * - Java typed UDF : add {@code implements UDFn} alongside {@code RapidsUDF} + * - Scala CPU UDF : implement the equivalent {@code UDFn<...>} contract. + * Invoke the Scala UDF via reflection from {@code call(...)}. + * 3. Add the CPU evaluation method. + * 4. Update {@code evaluateColumnar} and {@code evaluateNative} as needed to match the signature. + */ +public class PlaceholderUDFNameNativeRapidsUDF implements RapidsUDF { + + // TODO: copy the original CPU evaluation method here (evaluate / call). + + @Override + public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { + if (args.length != 1) { + throw new IllegalArgumentException("Unexpected argument count: " + args.length); + } + if (numRows != args[0].getRowCount()) { + throw new IllegalArgumentException( + "Expected " + numRows + " rows, received " + args[0].getRowCount()); + } + + NativeUDFLoader.ensureLoaded(); + return new ColumnVector(evaluateNative(args[0].getNativeView())); + } + + private static native long evaluateNative(long inputView); +} diff --git a/skills/udf-convert-to-cudf/SKILL.md b/skills/udf-convert-to-cudf/SKILL.md new file mode 100644 index 00000000000..58d4b01d2b7 --- /dev/null +++ b/skills/udf-convert-to-cudf/SKILL.md @@ -0,0 +1,128 @@ +--- +name: udf-convert-to-cudf +description: Assists with converting an Apache Spark UDF to a GPU-accelerated RapidsUDF using cuDF Java APIs. This is step 2 of 3 in the UDF conversion workflow (udf-gen-test -> udf-convert-to-cudf -> udf-benchmark). Use this skill when you have a CPU UDF with a unit test and need to convert it to a RapidsUDF. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# Convert UDF to cuDF RapidsUDF + +## Workflow + +- [ ] Step 1: Create the RapidsUDF file +- [ ] Step 2: Implement the `evaluateColumnar` method +- [ ] Step 3: Build and test +- [ ] Step 4: Check for memory leaks +- [ ] Step 5: Run judge subagent if requested +- [ ] Step 6: Review conversion + +**Before making any edits, create a visible TODO checklist for every workflow step in this skill and keep it updated.** Do not produce a final answer until every required checklist item is marked complete. + +## Prerequisites + +- Project directory from Step 1 (udf-gen-test) with passing unit test + +Derive `` and `` from the UDF class name. + +> **Note:** Commands require access to `/tmp` (Spark temp storage) and `/dev` (GPU device). If commands fail due to sandbox restrictions, re-run them unsandboxed. + +## Step 1: Create the RapidsUDF File + +Create a copy of the original UDF file in the same source directory (`src/main//com/udf/`), then modify it: + +1. Add imports: + Java: `import ai.rapids.cudf.*;`, `import com.nvidia.spark.RapidsUDF;` + Scala: `import ai.rapids.cudf._`, `import com.nvidia.spark.RapidsUDF`, `import Arm.{withResource, closeOnExcept}` +2. Add `implements RapidsUDF` to the class declaration +3. Add the `evaluateColumnar` method stub: + Java: `public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { }` + Scala: `def evaluateColumnar(numRows: Int, args: ColumnVector*): ColumnVector = { }` +4. Rename the class and the file to `RapidsUDF` + +## Step 2: Implement the `evaluateColumnar` method + +### Background + +**Read `references/RAPIDS_UDF.md`** for detailed background on: +- How RapidsUDF and `evaluateColumnar` work +- Input ColumnVector types and output type mapping +- Debugging techniques and GPU memory management + +**Read `examples/` for example RapidsUDF implementations for the target language.** + +### Implementation + +1. Clone https://github.com/rapidsai/cudf (branch matching spark-rapids version) to `~/.cache/aether_agent/` if not already present. Explore `java/src//java/ai/rapids/cudf` for relevant methods and usage patterns. +2. Implement the `evaluateColumnar` method using cuDF APIs. + +### Critical Requirements + +- **NEVER use `copyToHost()` or methods that copy data GPU→CPU.** This defeats the purpose of GPU acceleration +- **Do NOT hardcode test values.** The RapidsUDF must implement actual business logic for ANY potential input + +## Step 3: Build and Test + +Fill in the target-specific TODOs in `src/test//com/udf/CudfComparisonTest.`: +- Implement `registerRapidsUDF` to register the new RapidsUDF class. +- Replace placeholders with the actual camel/snake UDF name + +Then run the test: +```bash +# Java +mvn test -Dtest=CudfComparisonTest + +# Scala +mvn test -Dsuites=com.udf.CudfComparisonTest +``` + +If the test fails, analyze the error and iterate on the RapidsUDF implementation. + +### Difficult Test Failures + +Treat the unit test as the CPU behavior specification. Do not weaken or remove test cases silently. + +- Tests that check for CPU errors may not be directly applicable to a columnar implementation: the GPU path typically evaluates a whole column and may produce nulls for invalid rows instead of throwing row-level exceptions. If this causes an unavoidable mismatch, add a clear comment in the test and a `TODO/NOTE` in the implementation explaining the mismatch. +- If a test case does not pass because of inherent cuDF/libcudf/API limitations or low-level GPU/CPU semantic differences, comment out the conflicting assertion/test only after documenting how you tried to make the behavior match and why those attempts failed. Add a note to the user. +- If the behavior is important, common, or part of the documented input domain, **always prefer fixing the implementation** over commenting out the test case. The exception is a performance-vs-correctness tradeoff that the user explicitly approves. + +## Step 4: Memory Leak Check + +Re-run with memory leak detection: +```bash +# Java +mvn test -Dtest=CudfComparisonTest -Ddebug.memory.leaks=true > /tmp/memleak.log 2>&1 + +# Scala +mvn test -Dsuites=com.udf.CudfComparisonTest -Ddebug.memory.leaks=true > /tmp/memleak.log 2>&1 + +# Check for leaks +grep "LEAKED" /tmp/memleak.log | head -5 +``` + +If leaks are found, ensure all GPU objects are properly closed. + +## Step 5: Run Judge Subagent If Requested + +If the user explicitly asked for the judge, a judge subagent, or a review agent, treat that as an explicit request for delegation: you **MUST** launch a separate subagent with `model: inherit` and instruct it to use the **udf-judge-conversion** skill. Ask it to review the `UnitTest`, `CudfComparisonTest`, and RapidsUDF implementation. + +If the user did not request a judge/review agent, mark this step as skipped and continue to Step 6. If a required judge subagent is blocked by tool policy, stop and tell the user that explicit permission/instruction is needed. + +If you run the judge, wait for it to complete and review its report. If the judge finds any issues, 1) fix the issues, 2) re-run the tests and leak checks, and 3) re-run the judge subagent. + +## Step 6: Review Conversion + +Review your own work to ensure: +- The test runs on the GPU and directly compares CPU-GPU outputs +- The implementation does not overfit to test cases +- No `copyToHost()` or row-by-row GPU-to-CPU copying is used for computation +- No debug statements (e.g., `TableDebug.get().debug(...)`) remain in final output + +## Output + +Upon successful completion: +- RapidsUDF file at `src/main//com/udf/RapidsUDF.` +- Comparison test passes with no memory leaks + +These outputs are required for **Step 3: Benchmark**. diff --git a/skills/udf-convert-to-cudf/examples/URLDecode.java b/skills/udf-convert-to-cudf/examples/URLDecode.java new file mode 100644 index 00000000000..7122e15a68e --- /dev/null +++ b/skills/udf-convert-to-cudf/examples/URLDecode.java @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import ai.rapids.cudf.*; +import com.nvidia.spark.RapidsUDF; +import org.apache.spark.sql.api.java.UDF1; + +import java.io.UnsupportedEncodingException; +import java.net.URLDecoder; + +/** Decode URL-encoded strings. */ +public class URLDecode implements UDF1, RapidsUDF { + /** Row-by-row implementation that executes on the CPU */ + @Override + public String call(String s) { + String result = null; + if (s != null) { + try { + result = URLDecoder.decode(s, "utf-8"); + } catch (IllegalArgumentException ignored) { + result = s; + } catch (UnsupportedEncodingException e) { + // utf-8 is a builtin, standard encoding, so this should never happen + throw new RuntimeException(e); + } + } + return result; + } + + /** Columnar implementation that runs on the GPU */ + @Override + public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { + // The CPU implementation takes a single string argument, so similarly + // there should only be one column argument of type STRING. + if (args.length != 1) { + throw new IllegalArgumentException("Unexpected argument count: " + args.length); + } + ColumnVector input = args[0]; + if (numRows != input.getRowCount()) { + throw new IllegalArgumentException("Expected " + numRows + " rows, received " + input.getRowCount()); + } + if (!input.getType().equals(DType.STRING)) { + throw new IllegalArgumentException("Argument type is not a string column: " + + input.getType()); + } + + // The cudf urlDecode does not convert '+' to a space, so do that as a pre-pass first. + // All intermediate results are closed to avoid leaking GPU resources. + try (Scalar plusScalar = Scalar.fromString("+"); + Scalar spaceScalar = Scalar.fromString(" "); + ColumnVector replaced = input.stringReplace(plusScalar, spaceScalar)) { + return replaced.urlDecode(); + } + } +} diff --git a/skills/udf-convert-to-cudf/examples/URLDecodeExtendsFunction.scala b/skills/udf-convert-to-cudf/examples/URLDecodeExtendsFunction.scala new file mode 100644 index 00000000000..4a5e4f086f0 --- /dev/null +++ b/skills/udf-convert-to-cudf/examples/URLDecodeExtendsFunction.scala @@ -0,0 +1,44 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import java.net.URLDecoder + +import ai.rapids.cudf._ +import com.nvidia.spark.RapidsUDF +import Arm.{withResource, closeOnExcept} + +/** Decode URL-encoded strings. */ +class URLDecode extends Function1[String, String] with RapidsUDF with Serializable { + /** Row-by-row implementation that executes on the CPU */ + override def apply(s: String): String = { + Option(s).map { s => + try { + URLDecoder.decode(s, "utf-8") + } catch { + case _: IllegalArgumentException => s + } + }.orNull + } + + /** Columnar implementation that runs on the GPU */ + override def evaluateColumnar(numRows: Int, args: ColumnVector*): ColumnVector = { + // The CPU implementation takes a single string argument, so similarly + // there should only be one column argument of type STRING. + require(args.length == 1, s"Unexpected argument count: ${args.length}") + val input = args.head + require(numRows == input.getRowCount, s"Expected $numRows rows, received ${input.getRowCount}") + require(input.getType == DType.STRING, s"Argument type is not a string: ${input.getType}") + + // The cudf urlDecode does not convert '+' to a space, so do that as a pre-pass first. + // All intermediate results are closed using withResource to avoid leaking GPU resources. + withResource(Scalar.fromString("+")) { plusScalar => + withResource(Scalar.fromString(" ")) { spaceScalar => + withResource(input.stringReplace(plusScalar, spaceScalar)) { replaced => + replaced.urlDecode() + } + } + } + } +} diff --git a/skills/udf-convert-to-cudf/examples/URLDecodeHive.java b/skills/udf-convert-to-cudf/examples/URLDecodeHive.java new file mode 100644 index 00000000000..d5b571e7085 --- /dev/null +++ b/skills/udf-convert-to-cudf/examples/URLDecodeHive.java @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import ai.rapids.cudf.*; +import com.nvidia.spark.RapidsUDF; +import org.apache.hadoop.hive.ql.exec.UDF; + +import java.io.UnsupportedEncodingException; +import java.net.URLDecoder; + +/** Decode URL-encoded strings. */ +public class URLDecode extends UDF implements RapidsUDF { + + /** Row-by-row implementation that executes on the CPU */ + public String evaluate(String s) { + String result = null; + if (s != null) { + try { + result = URLDecoder.decode(s, "utf-8"); + } catch (IllegalArgumentException ignored) { + result = s; + } catch (UnsupportedEncodingException e) { + // utf-8 is a builtin, standard encoding, so this should never happen + throw new RuntimeException(e); + } + } + return result; + } + + /** Columnar implementation that runs on the GPU */ + @Override + public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { + // The CPU implementation takes a single string argument, so similarly + // there should only be one column argument of type STRING. + if (args.length != 1) { + throw new IllegalArgumentException("Unexpected argument count: " + args.length); + } + ColumnVector input = args[0]; + if (numRows != input.getRowCount()) { + throw new IllegalArgumentException("Expected " + numRows + " rows, received " + input.getRowCount()); + } + if (!input.getType().equals(DType.STRING)) { + throw new IllegalArgumentException("Argument type is not a string column: " + + input.getType()); + } + + // The cudf urlDecode does not convert '+' to a space, so do that as a pre-pass first. + // All intermediate results are closed to avoid leaking GPU resources. + try (Scalar plusScalar = Scalar.fromString("+"); + Scalar spaceScalar = Scalar.fromString(" "); + ColumnVector replaced = input.stringReplace(plusScalar, spaceScalar)) { + return replaced.urlDecode(); + } + } +} diff --git a/skills/udf-convert-to-cudf/examples/URLDecodeWithField.scala b/skills/udf-convert-to-cudf/examples/URLDecodeWithField.scala new file mode 100644 index 00000000000..7dc4f122dcd --- /dev/null +++ b/skills/udf-convert-to-cudf/examples/URLDecodeWithField.scala @@ -0,0 +1,48 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import java.net.URLDecoder + +import ai.rapids.cudf._ +import com.nvidia.spark.RapidsUDF +import Arm.{withResource, closeOnExcept} + +/** Decode URL-encoded strings. */ +object URLDecode { + val myUDF = udf( + new Function1[String, String] with RapidsUDF with Serializable { + /** Row-by-row implementation that executes on the CPU */ + override def apply(s: String): String = { + Option(s).map { s => + try { + URLDecoder.decode(s, "utf-8") + } catch { + case _: IllegalArgumentException => s + } + }.orNull + } + + /** Columnar implementation that runs on the GPU */ + override def evaluateColumnar(numRows: Int, args: ColumnVector*): ColumnVector = { + // The CPU implementation takes a single string argument, so similarly + // there should only be one column argument of type STRING. + require(args.length == 1, s"Unexpected argument count: ${args.length}") + val input = args.head + require(numRows == input.getRowCount, s"Expected $numRows rows, received ${input.getRowCount}") + require(input.getType == DType.STRING, s"Argument type is not a string: ${input.getType}") + + // The cudf urlDecode does not convert '+' to a space, so do that as a pre-pass first. + // All intermediate results are closed using withResource to avoid leaking GPU resources. + withResource(Scalar.fromString("+")) { plusScalar => + withResource(Scalar.fromString(" ")) { spaceScalar => + withResource(input.stringReplace(plusScalar, spaceScalar)) { replaced => + replaced.urlDecode() + } + } + } + } + } + ) +} diff --git a/skills/udf-convert-to-cudf/references/RAPIDS_UDF.md b/skills/udf-convert-to-cudf/references/RAPIDS_UDF.md new file mode 100644 index 00000000000..737aa640d41 --- /dev/null +++ b/skills/udf-convert-to-cudf/references/RAPIDS_UDF.md @@ -0,0 +1,111 @@ + + +# Background: RAPIDS Accelerated UDFs + +These instructions document how to implement a GPU version of an existing CPU UDF using the RapidsUDF interface. The RapidsUDF interface provides a way to run a CPU UDF on the GPU when using the RAPIDS Accelerator for Apache Spark. + +## Implementation + +The original CPU implementation is in the `evaluate` method. To make a UDF run on the GPU, you must implement the RapidsUDF interface, which provides a single method you need to override called `evaluateColumnar`. The `evaluateColumnar` function should use pre-existing cuDF methods from the [Java APIs of RAPIDS cudf](https://docs.rapids.ai/api/cudf-java/legacy) to perform the UDF computation by operating on cudF ColumnVectors. + +Note that you must keep both CPU and GPU evaluate methods, so that the UDF will still work if a higher-level operation involving the Rapids UDF falls back to the CPU. + +Refer to examples/ for example RapidsUDF implementations. + +## Interpreting Inputs + +The RAPIDS Accelerator will pass columnar forms of the same inputs for the CPU version of the UDF into the `args` array. For example, if the CPU UDF expects two inputs, a String and an Integer, then the evaluateColumnar method will be invoked with an array of two cuDF ColumnVector instances of type STRING and INT32 respectively. + +Note that passing scalar inputs to a RAPIDS accelerated UDF is supported with limitations. The scalar value will be replicated into a full column before being passed to evaluateColumnar. Therefore the UDF implementation cannot easily detect the difference between a scalar input and a columnar input. + +The implementation of evaluateColumnar must return a column with the specified numRows, equal to the input number of rows. All input columns will contain the same number of rows. + +## Generating output + +evaluateColumnar must return a ColumnVector of an appropriate cuDF type to match the result type of the original UDF. + +The following table shows the mapping of Spark types to equivalent cuDF columnar types: + +| Spark Type | cuDF Type | +|---------------|--------------------------------------------| +| BooleanType | BOOL8 | +| ByteType | INT8 | +| ShortType | INT16 | +| IntegerType | INT32 | +| LongType | INT64 | +| FloatType | FLOAT32 | +| DoubleType | FLOAT64 | +| DecimalType | DECIMAL32, DECIMAL64, DECIMAL128 * | +| DateType | TIMESTAMP_DAYS | +| TimestampType | TIMESTAMP_MICROSECONDS | +| StringType | STRING | +| NullType | INT8 | +| ArrayType | LIST of the underlying element type | +| MapType | LIST of STRUCT of the key and value types | +| StructType | STRUCT of all the field types | + +For example, if the CPU UDF returns the Spark type `ArrayType(MapType(StringType, StringType))` then evaluateColumnar must return a column of type `LIST(LIST(STRUCT(STRING,STRING)))`. + +*Note: cuDF's DECIMAL32 corresponds to precision <= 9 digits, DECIMAL64 corresponds to 9 < precision <= 18 digits, and DECIMAL128 corresponds to 18 < precision <= 38 digits. Precision greater than 38 digits is unsupported. + +Note that cuDF decimals use a negative scale relative to Spark DecimalType. For example, Spark DecimalType(precision=11, scale=2) would translate to cuDF type DECIMAL64(scale=-2). + +## Debugging + +When debugging, it may be helpful to print data type information about cuDF objects. For example, to get information about a ColumnVector: + +```java +System.out.println("Param 1 info:" + param1Column); +``` + +Example output: + +```text +Param 1 info: ColumnVector{rows=10, type=INT32, nullCount=Optional.empty, offHeap=(ID: 880 7d1d4c5951e0)} +``` + +To print the actual values in a column or table, use `TableDebug`: + +```java +TableDebug debugger = TableDebug.get(); +debugger.debug("Param 1 data:", param1Column); +``` + +Note that you should NEVER call this from production code, since it causes a device-to-host copy. + +## Managing Memory + +The Java memory model is not friendly for doing GPU operations because the JVM makes the assumption that everything we're trying to do is in heap memory. **Therefore, you must free the GPU resources in a timely manner with try-finally blocks**, calling `close()` to release GPU resources and `incRefCount()` to increment reference counts. + +The JVM's garbage collector is generally triggered when the JVM heap runs out of free space, but not necessarily when the GPU memory runs out. +To prevent these GPU memory leaks, the cuDF Java code tracks these objects, and if the garbage collector causes the memory to be freed instead of a proper close, it will output a warning like the following: + +```text +ERROR ColumnVector: A DEVICE COLUMN VECTOR WAS LEAKED (ID: 15 7fb5f94d8fa0) +``` + +These messages are an indication that an object on the GPU was not properly closed. Once a leak is detected, the Spark driver/executor `extraJavaOptions` can be set to `-Dai.rapids.refcount.debug=true -ea` to get a stack trace for the leak. + +The user will run the unit test and provide tracebacks if memory leaks occur to help you debug the issue. + +For Scala, use `withResource` and `closeOnExcept` from the `Arm` object for resource management. + +**Note:** Avoid placing the input ColumnVectors (those passed in `args`) in try-finally or try-with-resources blocks. The RAPIDS Accelerator will close the input columns for you. For example, avoid doing this: + +```java +ColumnVector param1 = args[0]; +try { + // Do something with param1 +} finally { + param1.close(); +} +``` + +This will result in a double-close error: + +```text +java.lang.IllegalStateException: Close called too many times ColumnVector{rows=10, type=INT32, nullCount=Optional.empty, offHeap=(ID: 637 0)} +``` diff --git a/skills/udf-convert-to-sql/SKILL.md b/skills/udf-convert-to-sql/SKILL.md new file mode 100644 index 00000000000..a55f464e555 --- /dev/null +++ b/skills/udf-convert-to-sql/SKILL.md @@ -0,0 +1,87 @@ +--- +name: udf-convert-to-sql +description: Assists with converting an Apache Spark UDF to a functionally equivalent Spark SQL expression. This is step 2 of 3 in the UDF conversion workflow (udf-gen-test -> udf-convert-to-sql -> udf-benchmark). Use this skill when you have a CPU UDF with a unit test and need to convert it to SQL for GPU acceleration. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# Convert UDF to Spark SQL + +## Workflow + +- [ ] Step 1: Implement the SQL expression +- [ ] Step 2: Fill in the comparison test and iterate +- [ ] Step 3: Run judge subagent if requested +- [ ] Step 4: Review conversion + +**Before making any edits, create a visible TODO checklist for every workflow step in this skill and keep it updated.** Do not produce a final answer until every required checklist item is marked complete. + +## Prerequisites + +- Project directory from Step 1 (udf-gen-test) with passing unit test + +Derive `` and `` from the UDF class name. + +> **Note:** Commands require access to `/tmp` (Spark temp storage) and `/dev` (GPU device). If commands fail due to sandbox restrictions, re-run them unsandboxed. + +## Step 1: Implement the SQL Expression + +Implement the SQL expression in a file at `src/main/resources/.sql`. + +**Read `examples/` for example UDF-to-SQL conversions for the target language.** + +### Guidelines + +- Focus on correctness FIRST, then GPU compatibility — the test will report which operators are not GPU-compatible +- Avoid expensive joins; prefer window functions, CTEs, and built-in array/map functions over explode-and-aggregate patterns + +**Do NOT hardcode test sample values or outputs.** The SQL expression must work correctly for ANY potential input. + +## Step 2: Fill in test and iterate + +Update `src/test//com/udf/SqlComparisonTest.`: +- Update the SQL file path to point to your `src/main/resources/.sql` file +- Replace placeholders with the actual camel/snake UDF name + +Then run the test: +```bash +# Java +mvn test -Dtest=SqlComparisonTest + +# Scala +mvn test -Dsuites=com.udf.SqlComparisonTest +``` + +If the test fails, analyze the error and iterate on the SQL expression. + +### Difficult Test Failures + +Treat the unit test as the CPU behavior specification. Do not weaken or remove test cases silently. + +- Tests that check for CPU errors may not be directly applicable to SQL operators: Spark RAPIDS typically evaluates a whole column/batch and may produce nulls for invalid rows instead of throwing one row-level exception. Make an explicit judgment call about the UDF contract. Add a clear comment in the test and a `TODO/NOTE` in the SQL statement explaining the mismatch. +- In rare cases, the Spark RAPIDS Plugin has known discrepancies in certain SQL operators. If a test case does not pass because of these discrepancies, notify the user and comment out the conflicting assertion/test only after documenting how you tried to make the behavior match and why those attempts failed. +- If the behavior is important, common, or part of the documented input domain, **always prefer fixing the SQL expression** over commenting out the test case. The exception is a performance-vs-correctness tradeoff that the user explicitly approves. + +## Step 3: Run Judge Subagent If Requested + +If the user explicitly asked for the judge, a judge subagent, or a review agent, treat that as an explicit request for delegation: you **MUST** launch a separate subagent with `model: inherit` and instruct it to use the **udf-judge-conversion** skill. Ask it to review the `UnitTest`, `SqlComparisonTest`, and SQL expression. + +If the user did not request a judge/review agent, mark this step as skipped and continue to Step 4. If a required judge subagent is blocked by tool policy, stop and tell the user that explicit permission/instruction is needed. + +If you run the judge, wait for it to complete and review its report. If the judge finds any issues, 1) fix the issues, 2) re-run the tests, and 3) re-run the judge subagent. + +## Step 4: Review Conversion + +Review your own work to ensure: +- The test runs on the GPU and directly compares CPU-SQL outputs +- The implementation does not overfit to test cases + +## Output + +Upon successful completion: +- SQL file at `src/main/resources/.sql` +- Comparison test passes + +These outputs are required for **Step 3: Benchmark**. diff --git a/skills/udf-convert-to-sql/examples/FormatPhone.java b/skills/udf-convert-to-sql/examples/FormatPhone.java new file mode 100644 index 00000000000..1b0d227199a --- /dev/null +++ b/skills/udf-convert-to-sql/examples/FormatPhone.java @@ -0,0 +1,27 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import org.apache.spark.sql.api.java.UDF1; + +/** + * Strip non-digit characters and format as (XXX) XXX-XXXX. + * See format_phone.sql for equivalent SQL expression. + */ +public class FormatPhone implements UDF1 { + @Override + public String call(String phone) throws Exception { + if (phone == null) { + return null; + } + String digits = phone.replaceAll("[^0-9]", ""); + if (digits.length() != 10) { + return null; + } + return String.format("(%s) %s-%s", + digits.substring(0, 3), + digits.substring(3, 6), + digits.substring(6)); + } +} diff --git a/skills/udf-convert-to-sql/examples/FormatPhone.scala b/skills/udf-convert-to-sql/examples/FormatPhone.scala new file mode 100644 index 00000000000..0aebe0ef11d --- /dev/null +++ b/skills/udf-convert-to-sql/examples/FormatPhone.scala @@ -0,0 +1,22 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import org.apache.spark.sql.functions.udf + +/** + * Strip non-digit characters and format as (XXX) XXX-XXXX. + * See format_phone.sql for equivalent SQL expression. + */ +object FormatPhone { + val formatPhone = udf((phone: String) => { + Option(phone).flatMap { p => + val digits = p.replaceAll("[^0-9]", "") + if (digits.length == 10) + Some(s"($${digits.substring(0, 3)}) $${digits.substring(3, 6)}-$${digits.substring(6)}") + else + None + }.orNull + }) +} diff --git a/skills/udf-convert-to-sql/examples/FormatPhoneHive.java b/skills/udf-convert-to-sql/examples/FormatPhoneHive.java new file mode 100644 index 00000000000..4609b7254ee --- /dev/null +++ b/skills/udf-convert-to-sql/examples/FormatPhoneHive.java @@ -0,0 +1,26 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import org.apache.hadoop.hive.ql.exec.UDF; + +/** + * Strip non-digit characters and format as (XXX) XXX-XXXX. + * See format_phone.sql for equivalent SQL expression. + */ +public class FormatPhone extends UDF { + public String evaluate(String phone) { + if (phone == null) { + return null; + } + String digits = phone.replaceAll("[^0-9]", ""); + if (digits.length() != 10) { + return null; + } + return String.format("(%s) %s-%s", + digits.substring(0, 3), + digits.substring(3, 6), + digits.substring(6)); + } +} diff --git a/skills/udf-convert-to-sql/examples/NormalizeTags.java b/skills/udf-convert-to-sql/examples/NormalizeTags.java new file mode 100644 index 00000000000..152d63bb480 --- /dev/null +++ b/skills/udf-convert-to-sql/examples/NormalizeTags.java @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import org.apache.spark.sql.api.java.UDF1; +import scala.collection.Seq; +import scala.collection.Iterator; + +import java.util.ArrayList; +import java.util.List; +import java.util.TreeSet; + +/** + * Lowercase, deduplicate, and sort a variable-length tag array. + * See normalize_tags.sql for equivalent SQL expression. + */ +public class NormalizeTags implements UDF1, List> { + @Override + public List call(Seq tags) throws Exception { + if (tags == null) { + return null; + } + TreeSet result = new TreeSet<>(); + Iterator it = tags.iterator(); + while (it.hasNext()) { + String tag = it.next(); + if (tag != null) { + String stripped = tag.replaceAll("^ +| +$", "").toLowerCase(); + if (!stripped.isEmpty()) { + result.add(stripped); + } + } + } + return result.isEmpty() ? null : new ArrayList<>(result); + } +} diff --git a/skills/udf-convert-to-sql/examples/NormalizeTags.scala b/skills/udf-convert-to-sql/examples/NormalizeTags.scala new file mode 100644 index 00000000000..92a2eee6954 --- /dev/null +++ b/skills/udf-convert-to-sql/examples/NormalizeTags.scala @@ -0,0 +1,25 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import org.apache.spark.sql.expressions.UserDefinedFunction +import org.apache.spark.sql.functions.udf + +/** + * Lowercase, deduplicate, and sort a variable-length tag array. + * See normalize_tags.sql for equivalent SQL expression. + */ +object NormalizeTags { + val normalizeTags: UserDefinedFunction = udf((tags: Seq[String]) => { + Option(tags).flatMap { ts => + val cleaned = ts + .filter(_ != null) + .map(_.replaceAll("^ +| +$", "").toLowerCase) + .filter(_.nonEmpty) + .distinct + .sorted + if (cleaned.isEmpty) None else Some(cleaned) + }.orNull + }) +} diff --git a/skills/udf-convert-to-sql/examples/NormalizeTagsHive.java b/skills/udf-convert-to-sql/examples/NormalizeTagsHive.java new file mode 100644 index 00000000000..058bd210c5e --- /dev/null +++ b/skills/udf-convert-to-sql/examples/NormalizeTagsHive.java @@ -0,0 +1,32 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import org.apache.hadoop.hive.ql.exec.UDF; + +import java.util.ArrayList; +import java.util.List; +import java.util.TreeSet; + +/** + * Lowercase, deduplicate, and sort a variable-length tag array. + * See normalize_tags.sql for equivalent SQL expression. + */ +public class NormalizeTags extends UDF { + public List evaluate(List tags) { + if (tags == null) { + return null; + } + TreeSet result = new TreeSet<>(); + for (String tag : tags) { + if (tag != null) { + String stripped = tag.replaceAll("^ +| +$", "").toLowerCase(); + if (!stripped.isEmpty()) { + result.add(stripped); + } + } + } + return result.isEmpty() ? null : new ArrayList<>(result); + } +} diff --git a/skills/udf-convert-to-sql/examples/format_phone.sql b/skills/udf-convert-to-sql/examples/format_phone.sql new file mode 100644 index 00000000000..6a35040c0e7 --- /dev/null +++ b/skills/udf-convert-to-sql/examples/format_phone.sql @@ -0,0 +1,17 @@ +-- SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +-- SPDX-License-Identifier: Apache-2.0 + +SELECT + CASE + WHEN phone IS NULL THEN NULL + WHEN LENGTH(REGEXP_REPLACE(phone, '[^0-9]', '')) != 10 THEN NULL + ELSE CONCAT( + '(', + SUBSTR(REGEXP_REPLACE(phone, '[^0-9]', ''), 1, 3), + ') ', + SUBSTR(REGEXP_REPLACE(phone, '[^0-9]', ''), 4, 3), + '-', + SUBSTR(REGEXP_REPLACE(phone, '[^0-9]', ''), 7, 4) + ) + END AS result +FROM __table__ diff --git a/skills/udf-convert-to-sql/examples/normalize_tags.sql b/skills/udf-convert-to-sql/examples/normalize_tags.sql new file mode 100644 index 00000000000..385d6adfca8 --- /dev/null +++ b/skills/udf-convert-to-sql/examples/normalize_tags.sql @@ -0,0 +1,15 @@ +-- SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +-- SPDX-License-Identifier: Apache-2.0 + +SELECT + CASE + WHEN tags IS NULL THEN NULL + WHEN SIZE(FILTER(tags, x -> x IS NOT NULL AND TRIM(x) != '')) = 0 THEN NULL + ELSE ARRAY_SORT(ARRAY_DISTINCT( + TRANSFORM( + FILTER(tags, x -> x IS NOT NULL AND TRIM(x) != ''), + x -> LOWER(TRIM(x)) + ) + )) + END AS result +FROM __table__ diff --git a/skills/udf-gen-test/SKILL.md b/skills/udf-gen-test/SKILL.md new file mode 100644 index 00000000000..44b668dfd12 --- /dev/null +++ b/skills/udf-gen-test/SKILL.md @@ -0,0 +1,148 @@ +--- +name: udf-gen-test +description: Assists with generating a unit test for an Apache Spark UDF. This is step 1 of 3 in the UDF conversion workflow (udf-gen-test -> udf-convert-to-* -> udf-benchmark). Use this skill when you have a CPU UDF and need to create a unit test for the UDF before converting it into a GPU-compatible implementation. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# UDF Unit Test Generation + +## Workflow + +- [ ] Step 1: Set up project (copy template, add UDF source) +- [ ] Step 2: Implement the unit test (fill in TODO methods) +- [ ] Step 3: Compile and test until passing +- [ ] Step 4: Run coverage and inspect gaps +- [ ] Step 5: Verify outputs + +**Before making any edits, create a visible TODO checklist for every workflow step in this skill and keep it updated.** Do not produce a final answer until every required checklist item is marked complete. + +## Prerequisites + +- Path to the input UDF file (Java or Scala) + +Derive `` and `` from the UDF class name. + +> **Note:** Commands require access to `/tmp` (Spark temp storage) and `/dev` (GPU device). If commands fail due to sandbox restrictions, re-run them unsandboxed. + +## Step 1: Set Up the Project + +### 1a. Copy the template project + +The project can be found under this skill's templates directory. +```bash +cp -r templates/ // +``` + +This provides a complete Maven project with all test and benchmark infrastructure. + +### 1b. Copy or extract the UDF source + +Before copying code, decide whether the input UDF is already self-contained: +- If the UDF file contains only the target UDF and local helpers it directly needs, copy it as-is. +- If the UDF is part of a larger project or a file containing unrelated UDFs/classes, extract only the target UDF class/object and all local helper classes/methods required for that UDF to compile and run (modifying package declarations as needed). + +The template project should contain the smallest self-contained implementation of the target CPU UDF. + +Place the resulting source file(s) in the source directory: +- Java: `/src/main/java/com/udf/` +- Scala: `/src/main/scala/com/udf/` + +Set the package declaration to `com.udf`: +- Java: `package com.udf;` +- Scala: `package com.udf` + +## Step 2: Implement the Unit Test + +Read `src/test//com/udf/UnitTest.`. Replace placeholders with the actual camel/snake UDF name. + +Fill in the TODO methods following the docstrings. Include diverse edge cases in `createTestData` (nulls, empty strings, malformed inputs, varying lengths). + +### Test Data Coverage + +The generated tests should serve as a strong specification of the CPU UDF behavior over a documented input domain, and are intended to prove that a GPU or SQL implementation preserves the CPU UDF behavior. +For each input type and visible UDF branch, include applicable examples from these coverage dimensions: +- null inputs and null elements +- empty strings, arrays, maps, or structs +- malformed or unparsable inputs +- edges of input boundaries, such as min/max valid values, string length, or array length +- numeric sign/identity cases, such as negative, zero, and positive values +- string variety, such as unicode, ASCII, and encoding-sensitive inputs +- date/time boundaries, such as epoch, end-of-day/month/year, leap day, and DST/timezone transitions +- decimal precision and scale +- duplicate rows and repeated values +- mixed valid/invalid rows in the same DataFrame +- nested empty and nested null values + +Assertions should verify schema, row count, deterministic ordering, output values, null propagation, and exception/default behavior. Every visible UDF branch should be covered by the unit test or explicitly documented as out of scope. + +### Critical Requirements + +- Do NOT hardcode the UDF name; use the provided `udfName` argument. This ensures the correct registered UDF is exercised. +- Assume the user's UDF implementation is correct; the assertions should reflect its actual behavior. + +## Step 3: Compile and Test + +```bash +# Java +mvn test -Dtest=UnitTest + +# Scala +mvn test -Dsuites=com.udf.UnitTest +``` + +If it fails, analyze the error output (stdout/stderr) and fix the test code. Continue iterating until the test passes. + +## Step 4: Coverage Report + +The template projects use JaCoCo (Java) / scoverage (Scala) code coverage tools. + +```bash +# Java +mvn -Pcoverage test jacoco:report -Dtest=UnitTest + +# Scala +mvn -Pcoverage scoverage:report -Dsuites=com.udf.UnitTest +``` + +For Java, read `target/site/jacoco/jacoco.csv` and inspect LINE, BRANCH, and METHOD counters for the target CPU UDF class and local helper classes. In `jacoco.xml`, counters appear as `` elements, and source-line misses appear under ``. + +For Scala, read `target/scoverage.xml` and inspect statement, branch, and method-level coverage for the target CPU UDF class/object and local helper classes/objects. scoverage XML stores package/class/method `statement-rate` and `branch-rate` attributes, and each executable statement has `line`, `branch`, and `invocation-count` attributes. + +Use the coverage report as actionable feedback: +1. Inspect missed Java line, branch, and method coverage, or missed Scala statement, branch, and method-level coverage. +2. Add test cases and assertions that exercise those paths. +3. Re-run the unit test and coverage report. +4. Repeat until important CPU UDF branches are covered. + +If a missed line, statement, branch, or method path cannot or should not be tested, add a clear comment explaining why. Examples include: +- unreachable defensive code +- unsupported input domains +- unrelated template infrastructure + +Report the relevant counters for the target CPU UDF and local helper classes/objects: +- Java: LINE, BRANCH, and METHOD counters from JaCoCo. +- Scala: statement and branch coverage from scoverage, plus method-level statement/branch rates from `` elements. + +NOTE: JaCoCo and scoverage will not track source-level coverage in external JARs. If the UDF relies on external JAR business logic, make a note of this residual coverage gap. + +## Step 5: Verify Outputs + +After the test passes, verify that: +1. The test data covers various edge cases and reflects realistic input formats +2. The assertions reflect actual UDF behavior (no "cheating" by hardcoding values) +3. The coverage report shows strong coverage of the target CPU UDF and local helper logic +4. Any uncovered lines, branches, or methods are explicitly explained +5. Any external JAR logic invoked by the UDF is called out as outside the coverage scope + +If any quality checks fail, revise the test code and re-run. + +## Output + +Upon successful completion: +- Project directory: `//` +- Unit test: `src/test//com/udf/UnitTest.` + +These outputs are required for **Step 2: Convert UDF**. diff --git a/skills/udf-gen-test/templates/java/.mvn/jvm.config b/skills/udf-gen-test/templates/java/.mvn/jvm.config new file mode 100644 index 00000000000..9d7bf1a15b2 --- /dev/null +++ b/skills/udf-gen-test/templates/java/.mvn/jvm.config @@ -0,0 +1,16 @@ +-Xmx5g +-ea +--add-opens=java.base/java.lang=ALL-UNNAMED +--add-opens=java.base/java.lang.invoke=ALL-UNNAMED +--add-opens=java.base/java.lang.reflect=ALL-UNNAMED +--add-opens=java.base/java.io=ALL-UNNAMED +--add-opens=java.base/java.net=ALL-UNNAMED +--add-opens=java.base/java.nio=ALL-UNNAMED +--add-opens=java.base/java.util=ALL-UNNAMED +--add-opens=java.base/java.util.concurrent=ALL-UNNAMED +--add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED +--add-opens=java.base/sun.nio.ch=ALL-UNNAMED +--add-opens=java.base/sun.nio.cs=ALL-UNNAMED +--add-opens=java.base/sun.security.action=ALL-UNNAMED +--add-opens=java.base/sun.util.calendar=ALL-UNNAMED +--add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED diff --git a/skills/udf-gen-test/templates/java/pom.xml b/skills/udf-gen-test/templates/java/pom.xml new file mode 100644 index 00000000000..6925eb2f55f --- /dev/null +++ b/skills/udf-gen-test/templates/java/pom.xml @@ -0,0 +1,342 @@ + + + + 4.0.0 + com.udf + aether-agent-udfs + 1.0.0 + Aether UDF Conversion + This project contains UDFs that will be converted from CPU to GPU. + jar + + + 17 + 17 + UTF-8 + UTF-8 + UTF-8 + 2.12 + + 3.5.5 + 26.04.0 + 0.8.14 + + cuda12 + v26.04.00 + v26.04.00 + + + false + off + + ON + RAPIDS + 10 + ON + OFF + false + rapidsudfjni + ${project.build.directory}/native-build + + + -Xmx5g -ea + -Dai.rapids.refcount.debug=${debug.memory.leaks} + -Dorg.slf4j.simpleLogger.defaultLogLevel=off + -Dorg.slf4j.simpleLogger.log.ai.rapids.cudf=${cudf.log.level} + --add-opens=java.base/java.lang=ALL-UNNAMED + --add-opens=java.base/java.lang.invoke=ALL-UNNAMED + --add-opens=java.base/java.lang.reflect=ALL-UNNAMED + --add-opens=java.base/java.io=ALL-UNNAMED + --add-opens=java.base/java.net=ALL-UNNAMED + --add-opens=java.base/java.nio=ALL-UNNAMED + --add-opens=java.base/java.util=ALL-UNNAMED + --add-opens=java.base/java.util.concurrent=ALL-UNNAMED + --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED + --add-opens=java.base/sun.nio.ch=ALL-UNNAMED + --add-opens=java.base/sun.nio.cs=ALL-UNNAMED + --add-opens=java.base/sun.security.action=ALL-UNNAMED + --add-opens=java.base/sun.util.calendar=ALL-UNNAMED + --add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED + + + + + debug-leaks + + + debug.memory.leaks + true + + + + error + + + + coverage + + + + org.jacoco + jacoco-maven-plugin + ${jacoco.version} + + + HTML + XML + CSV + + + com/udf/bench/* + com/udf/SparkUtils* + + + + + prepare-agent + + prepare-agent + + + jacoco.agent.argLine + + + + report + verify + + report + + + + + + + + + cuda-native-udf + + + + org.apache.maven.plugins + maven-dependency-plugin + 3.6.1 + + + copy-rapids-jar-with-classifier + generate-sources + + copy + + + + + com.nvidia + rapids-4-spark_${scala.binary.version} + ${rapids4spark.version} + ${cuda.version} + jar + false + ${project.build.directory}/rapids-jar + + + true + + + + copy-rapids-jar-no-classifier + generate-sources + + copy + + + + + com.nvidia + rapids-4-spark_${scala.binary.version} + ${rapids4spark.version} + jar + false + ${project.build.directory}/rapids-jar + + + true + + + + + + org.apache.maven.plugins + maven-antrun-plugin + 3.1.0 + + + extract-cuda-native-dependencies + generate-sources + + ${skipCudfExtraction} + + + + + + + + + + + + + run + + + + cmake-cuda-native-udf + compile + + + + + + + + + + + + + + + + + + + + + + + + + + run + + + + + + org.apache.maven.plugins + maven-resources-plugin + 3.3.1 + + + copy-cuda-native-library-to-classes + process-classes + + copy-resources + + + true + ${project.build.outputDirectory}/${os.arch}/${os.name} + + + ${native.build.path} + + lib${native.library.name}.so + + + + + + + + + + + + + + + + org.apache.spark + spark-hive_${scala.binary.version} + ${spark.version} + provided + + + + com.nvidia + rapids-4-spark_${scala.binary.version} + ${rapids4spark.version} + provided + + + + junit + junit + 4.13.2 + test + + + + org.slf4j + slf4j-simple + 1.7.36 + test + + + + + + + + org.apache.maven.plugins + maven-surefire-plugin + 3.1.2 + + + **/*Test.java + + @{jacoco.agent.argLine} ${test.jvm.args} + + + + + org.codehaus.mojo + exec-maven-plugin + 3.1.0 + + + + org.apache.maven.plugins + maven-shade-plugin + 3.2.4 + + + package + + shade + + + false + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + + + + diff --git a/skills/udf-gen-test/templates/java/run_gen_data.sh b/skills/udf-gen-test/templates/java/run_gen_data.sh new file mode 100644 index 00000000000..44c802c1864 --- /dev/null +++ b/skills/udf-gen-test/templates/java/run_gen_data.sh @@ -0,0 +1,73 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Generate or validate benchmark data + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +print_usage() { + echo "Usage: $0 --rows NUM [--validate] [--output-path PATH] [--mvn-arg ARG]..." +} + +ROWS="" +VALIDATE="" +OUTPUT_PATH="" +MAVEN_ARGS=() + +while [[ $# -gt 0 ]]; do + case $1 in + --rows) ROWS="$2"; shift 2;; + --validate) VALIDATE="true"; shift;; + --output-path) OUTPUT_PATH="$2"; shift 2;; + --mvn-arg) MAVEN_ARGS+=("$2"); shift 2;; + *) + echo "Unknown option: $1" + print_usage + exit 1 + ;; + esac +done + +if [ -z "$ROWS" ]; then + echo "Error: --rows is required" + print_usage + exit 1 +fi + +SPARK_CONFS=( + --spark-conf spark.master="local[*]" + --spark-conf spark.driver.memory="16g" + --spark-conf spark.rapids.sql.enabled="true" + --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" + --spark-conf spark.locality.wait="0s" + --spark-conf spark.sql.cache.serializer="com.nvidia.spark.ParquetCachedBatchSerializer" + --spark-conf spark.rapids.sql.format.parquet.reader.type="MULTITHREADED" + --spark-conf spark.rapids.sql.reader.batchSizeBytes="1000MB" + --spark-conf spark.sql.files.maxPartitionBytes="512MB" + --spark-conf spark.rapids.sql.metrics.level="DEBUG" +) + +EXEC_ARGS="--rows $ROWS --partitions 32" +for arg in "${SPARK_CONFS[@]}"; do + EXEC_ARGS="$EXEC_ARGS $arg" +done + +if [ -n "$VALIDATE" ]; then + EXEC_ARGS="$EXEC_ARGS --validate" + echo "Running GenData in validation mode with $ROWS rows..." +else + if [ -z "$OUTPUT_PATH" ]; then + OUTPUT_PATH="data/bench_data_${ROWS}_rows.parquet" + fi + EXEC_ARGS="$EXEC_ARGS --output-path $OUTPUT_PATH" + echo "Running GenData to generate $ROWS rows -> $OUTPUT_PATH..." +fi + +mvn "${MAVEN_ARGS[@]}" compile exec:java \ + -Dexec.mainClass="com.udf.bench.GenData" \ + -Dexec.classpathScope=compile \ + -Dexec.args="$EXEC_ARGS" diff --git a/skills/udf-gen-test/templates/java/run_micro_benchmark.sh b/skills/udf-gen-test/templates/java/run_micro_benchmark.sh new file mode 100644 index 00000000000..6cac6d3e8c9 --- /dev/null +++ b/skills/udf-gen-test/templates/java/run_micro_benchmark.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Run in-memory microbenchmark for RapidsUDFs. + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +print_usage() { + echo "Usage: $0 --mode cpu|gpu|all --data-path PATH [--rows N] [--warmup N] [--measured N] [--pool-fraction F] [--profile] [--mvn-arg ARG]..." +} + +MODE="" +DATA_PATH="" +PROFILE="" +MAVEN_ARGS=() +RUNNER_ARGS=() + +while [[ $# -gt 0 ]]; do + case $1 in + --mode) MODE="$2"; RUNNER_ARGS+=("$1" "$2"); shift 2;; + --data-path) DATA_PATH="$2"; RUNNER_ARGS+=("$1" "$2"); shift 2;; + --profile) PROFILE="true"; RUNNER_ARGS+=("$1"); shift;; + --mvn-arg) MAVEN_ARGS+=("$2"); shift 2;; + *) RUNNER_ARGS+=("$1"); shift;; + esac +done + +if [ -z "$MODE" ] || [ -z "$DATA_PATH" ]; then + echo "Error: --mode and --data-path are required" + print_usage + exit 1 +fi + +MVN_CMD=( + mvn "${MAVEN_ARGS[@]}" compile exec:java + -Dexec.mainClass=com.udf.bench.MicroBenchRunner + -Dexec.classpathScope=compile + "-Dexec.args=${RUNNER_ARGS[*]}" +) + +if [ -n "$PROFILE" ]; then + REPORT_PATH="results/microbench_$(date +%Y%m%d_%H%M%S)" + mkdir -p results + echo "Running microbenchmark (mode=$MODE) on $DATA_PATH with nsys profiling..." + echo "nsys report will be saved to: ${REPORT_PATH}.nsys-rep" + nsys profile \ + -c cudaProfilerApi \ + --capture-range-end=stop \ + --trace=cuda,nvtx \ + --nvtx-domain-include="libcudf" \ + -o "$REPORT_PATH" \ + "${MVN_CMD[@]}" +else + echo "Running microbenchmark (mode=$MODE) on $DATA_PATH..." + "${MVN_CMD[@]}" +fi diff --git a/skills/udf-gen-test/templates/java/run_spark_benchmark.sh b/skills/udf-gen-test/templates/java/run_spark_benchmark.sh new file mode 100644 index 00000000000..5ce799d5413 --- /dev/null +++ b/skills/udf-gen-test/templates/java/run_spark_benchmark.sh @@ -0,0 +1,70 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Run CPU or GPU Spark benchmark. + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +print_usage() { + echo "Usage: $0 --mode cpu|gpu --data-path PATH [--result-path PATH] [--mvn-arg ARG]..." +} + +MODE="" +DATA_PATH="" +RESULT_PATH="" +MAVEN_ARGS=() + +while [[ $# -gt 0 ]]; do + case $1 in + --mode) MODE="$2"; shift 2;; + --data-path) DATA_PATH="$2"; shift 2;; + --result-path) RESULT_PATH="$2"; shift 2;; + --mvn-arg) MAVEN_ARGS+=("$2"); shift 2;; + *) + echo "Unknown option: $1" + print_usage + exit 1 + ;; + esac +done + +if [ -z "$MODE" ] || [ -z "$DATA_PATH" ]; then + echo "Error: --mode and --data-path are required" + print_usage + exit 1 +fi + +DATA_BASENAME=$(basename "$DATA_PATH" .parquet) +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +if [ -z "$RESULT_PATH" ]; then + RESULT_PATH="results/${MODE}_${DATA_BASENAME}_${TIMESTAMP}_result.json" +fi + +SPARK_CONFS=( + --spark-conf spark.master="local[*]" + --spark-conf spark.driver.memory="16g" + --spark-conf spark.rapids.sql.enabled="true" + --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" + --spark-conf spark.locality.wait="0s" + --spark-conf spark.sql.cache.serializer="com.nvidia.spark.ParquetCachedBatchSerializer" + --spark-conf spark.rapids.sql.format.parquet.reader.type="MULTITHREADED" + --spark-conf spark.rapids.sql.reader.batchSizeBytes="1000MB" + --spark-conf spark.sql.files.maxPartitionBytes="512MB" + --spark-conf spark.rapids.sql.metrics.level="DEBUG" +) + +EXEC_ARGS="--mode $MODE --data-path $DATA_PATH --result-path $RESULT_PATH" +for arg in "${SPARK_CONFS[@]}"; do + EXEC_ARGS="$EXEC_ARGS $arg" +done +EXEC_ARGS="$EXEC_ARGS --spark-conf spark.app.name=${MODE}_${DATA_BASENAME}_${TIMESTAMP}" + +echo "Running $MODE benchmark on $DATA_PATH..." +mvn "${MAVEN_ARGS[@]}" compile exec:java \ + -Dexec.mainClass="com.udf.bench.SparkBenchRunner" \ + -Dexec.classpathScope=compile \ + -Dexec.args="$EXEC_ARGS" diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/SparkUtils.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/SparkUtils.java new file mode 100644 index 00000000000..d50816e6fdf --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/SparkUtils.java @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import com.nvidia.spark.rapids.ExplainPlan; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +/** + * Spark utility methods. + */ +public class SparkUtils { + + /** + * Apply key=value Spark configs to a builder. + * + * @param builder the SparkSession builder to configure + * @param sparkConfs "spark.key=value" config strings + * @return the same builder, for chaining + */ + public static SparkSession.Builder applySparkConfs( + SparkSession.Builder builder, List sparkConfs) { + for (String conf : sparkConfs) { + String[] kv = conf.split("=", 2); + if (kv.length == 2) builder.config(kv[0], kv[1]); + } + return builder; + } + + /** + * Get a required argument from a parsed argument map, or throw. + * + * @param parsed the parsed argument map + * @param key the argument key (without "--" prefix) + * @return the argument value + * @throws IllegalArgumentException if the key is missing + */ + public static String requireArg(Map parsed, String key) { + String val = parsed.get(key); + if (val == null) { + throw new IllegalArgumentException("--" + key + " is required"); + } + return val; + } + + /** + * Ops that cause fallback but can be ignored, since they are strictly used for testing: + * - RDDScanExec/LocalTableScanExec: surfaces due to spark.createDataFrame() + * - CollectLimitExec: surfaces during dataframe collection (e.g. df.show()) + * - ToPrettyString: surfaces due to df.show() + */ + private static final Set IGNORE_OPERATIONS = new HashSet<>( + Arrays.asList("RDDScanExec", "LocalTableScanExec", "CollectLimitExec", "ToPrettyString") + ); + + /** + * Assert that the DataFrame's plan can run on GPU. + * NOTE: This is only reliable in explainOnly mode, with AQE disabled. + * + * @param df the DataFrame to check + * @throws RuntimeException if any operations cannot run on GPU + */ + public static void assertPlanRunsOnGpu(Dataset df) { + assertPlanRunsOnGpu(df, false); + } + + /** + * Assert that the DataFrame's plan can run on GPU. + * NOTE: This is only reliable in explainOnly mode, with AQE disabled. + * + * @param df the DataFrame to check + * @param returnFullPlan if true, include the full plan in the error message + * @throws RuntimeException if any operations cannot run on GPU + */ + public static void assertPlanRunsOnGpu(Dataset df, boolean returnFullPlan) { + String plan = getGpuPlan(df); + List unsupportedOps = getUnsupportedOps(plan); + if (!unsupportedOps.isEmpty()) { + StringBuilder sb = new StringBuilder(); + sb.append("Some operations cannot run on GPU.\nFound the following unsupported ops:\n"); + for (String op : unsupportedOps) { + sb.append("- ").append(op).append("\n"); + } + if (returnFullPlan) { + sb.append("\nFull physical plan:\n").append(plan); + } + throw new RuntimeException(sb.toString()); + } + } + + /** Get the potential GPU plan using the RAPIDS ExplainPlan API. */ + private static String getGpuPlan(Dataset df) { + return ExplainPlan.explainPotentialGpuPlan(df, "NOT_ON_GPU"); + } + + /** Parse the plan for unsupported operations (lines starting with '!'). */ + private static List getUnsupportedOps(String plan) { + List result = new ArrayList<>(); + for (String line : plan.split("\n")) { + // Each unsupported line looks like: ![Exec] cannot run on GPU + String trimmed = line.trim(); + if (trimmed.startsWith("!")) { + int start = trimmed.indexOf('<'); + int end = trimmed.indexOf('>'); + if (start >= 0 && end > start) { + String op = trimmed.substring(start + 1, end); + if (!IGNORE_OPERATIONS.contains(op)) { + result.add(trimmed); + } + } + } + } + return result; + } +} diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/BenchUtils.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/BenchUtils.java new file mode 100644 index 00000000000..8fbf2fa2fa5 --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/BenchUtils.java @@ -0,0 +1,112 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import static org.apache.spark.sql.functions.*; + +/** + * Benchmark utilities. + * - generateSyntheticData: Create benchmark data for the UDF + * - executeCpu: Register and run the CPU UDF + * - executeGpu: Register and run the GPU implementation + */ +public class BenchUtils { + + // --------------------------------------------------------------------------- + // Data generation + // --------------------------------------------------------------------------- + + /** + * TODO: Generate a synthetic DataFrame matching the unit test schema. + * + * Use {@code spark.range(0, numRows, 1, numPartitions)} as the base, then apply + * randomized column generators to produce data matching the UDF's expected input. + * + * Requirements: + * - Column names and types MUST match the unit test dataset schema + * - Data should be realistic and varied (different lengths, edge cases, etc.) + * - For variable-length inputs, generate sizable rows representative of + * enterprise-scale data + * + * Example: + * 
{@code
+     *   Dataset baseDF = spark.range(0, numRows, 1, numPartitions).toDF("id");
+     *   return baseDF.select(
+     *       col("id"),
+     *       expr("CAST(rand() * 850 AS INT)").alias("credit_score")
+     *   );
+     * }
+ * + * @param spark active SparkSession + * @param numRows number of rows to generate + * @param numPartitions number of output partitions + * @return DataFrame with the same schema as the unit test data + */ + public static Dataset generateSyntheticData( + SparkSession spark, long numRows, int numPartitions) { + return null; // TODO + } + + // --------------------------------------------------------------------------- + // Execution + // --------------------------------------------------------------------------- + + /** + * TODO: Execute the CPU UDF on the benchmark DataFrame. + * 1. Register the CPU UDF with Spark + * 2. Execute it on {@code df} + * 3. Return the result DataFrame + * + * Example: + * 
{@code
+     *   df.createOrReplaceTempView("bench_table");
+     *   spark.sql("CREATE TEMPORARY FUNCTION calculate_risk AS 'com.udf.CalculateRiskUDF'");
+     *   return spark.sql("SELECT *, calculate_risk(credit_score) AS risk_level FROM bench_table");
+     * }
+ * + * @param spark active SparkSession + * @param df input benchmark DataFrame + * @return result DataFrame after applying the CPU UDF + */ + public static Dataset executeCpu(SparkSession spark, Dataset df) { + return null; // TODO + } + + /** + * TODO: Execute the GPU implementation on the benchmark DataFrame. + * + * For RapidsUDF - register the RapidsUDF and run the same query as executeCpu: + * 
{@code
+     *   df.createOrReplaceTempView("bench_table");
+     *   spark.sql("CREATE TEMPORARY FUNCTION calculate_risk_rapids AS 'com.udf.CalculateRiskRapidsUDF'");
+     *   return spark.sql("SELECT *, calculate_risk_rapids(credit_score) AS risk_level FROM bench_table");
+     * }
+ * + * For SQL - read the SQL file from src/main/resources/ and adapt it for + * benchmarking. The SQL was written for the unit test, so you must: + * 1. Replace "test_table" with "bench_table" + * 2. Replace the SELECT column list with "SELECT *" to avoid referencing + * columns that may not exist in the benchmark DataFrame + * 
{@code
+     *   df.createOrReplaceTempView("bench_table");
+     *   String sqlContent = new String(Files.readAllBytes(Paths.get("src/main/resources/calculate_risk.sql")));
+     *   String benchSql = sqlContent.replace("test_table", "bench_table");
+     *   // Also replace the SELECT column list with SELECT * if needed
+     *   return spark.sql(benchSql);
+     * }
+ * + * @param spark active SparkSession + * @param df input benchmark DataFrame + * @return result DataFrame after applying the GPU implementation + */ + public static Dataset executeGpu(SparkSession spark, Dataset df) { + return null; // TODO + } +} diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/GenData.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/GenData.java new file mode 100644 index 00000000000..94a22eea753 --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/GenData.java @@ -0,0 +1,110 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import com.udf.SparkUtils; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; + +/** + * Generates benchmark data and optionally validates by running + * BenchUtils.executeCpu and BenchUtils.executeGpu. + * + * Usage: + * mvn exec:java -Dexec.mainClass=com.udf.bench.GenData \ + * -Dexec.args="--rows 1000 --validate --spark-conf k=v ..." + */ +public class GenData { + + public static void main(String[] args) { + Map argMap = new HashMap<>(); + List sparkConfs = new ArrayList<>(); + parseArgs(args, argMap, sparkConfs); + + long rows = Long.parseLong(SparkUtils.requireArg(argMap, "rows")); + int partitions = Integer.parseInt(argMap.getOrDefault("partitions", "32")); + boolean validate = argMap.containsKey("validate"); + String outputPath = argMap.get("output-path"); + + // Build Spark session + SparkSession.Builder builder = SparkSession.builder().appName("GenData"); + SparkUtils.applySparkConfs(builder, sparkConfs); + SparkSession spark = builder.enableHiveSupport().getOrCreate(); + + try { + // Generate synthetic data + Dataset df = BenchUtils.generateSyntheticData(spark, rows, partitions); + + // Verify row count + long actualRows = df.count(); + if (actualRows != rows) { + System.err.println("Row count mismatch: expected=" + rows + + ", actual=" + actualRows); + System.exit(1); + } + System.out.println("Generated " + actualRows + " rows across " + + partitions + " partitions"); + + if (validate) { + // Validation mode — run both CPU and GPU execute, don't write + for (String label : new String[]{"cpu", "gpu"}) { + try { + if ("cpu".equals(label)) { + BenchUtils.executeCpu(spark, df).collect(); + } else { + BenchUtils.executeGpu(spark, df).collect(); + } + System.out.println("Validation (" + label + ") passed."); + } catch (Exception e) { + System.err.println("Validation (" + label + ") failed: " + + e.getClass().getSimpleName() + ": " + e.getMessage()); + e.printStackTrace(System.err); + System.exit(1); + } + } + } else { + // Generation mode — write to output path + if (outputPath == null) { + throw new IllegalArgumentException( + "--output-path is required when not in validation mode"); + } + df.write().mode("overwrite").parquet(outputPath); + System.err.println("Successfully generated dataset and saved to: " + outputPath); + } + } catch (Exception e) { + System.err.println("Failed to generate dataset: " + + e.getClass().getSimpleName()); + e.printStackTrace(System.err); + System.exit(1); + } finally { + spark.stop(); + } + + System.exit(0); + } + + /** Parse CLI arguments. */ + private static void parseArgs(String[] args, Map map, List sparkConfs) { + int i = 0; + while (i < args.length) { + switch (args[i]) { + case "--rows": map.put("rows", args[i + 1]); i += 2; break; + case "--partitions": map.put("partitions", args[i + 1]); i += 2; break; + case "--validate": map.put("validate", "true"); i += 1; break; + case "--output-path": map.put("output-path", args[i + 1]); i += 2; break; + case "--spark-conf": sparkConfs.add(args[i + 1]); i += 2; break; + default: + throw new IllegalArgumentException("Unknown argument: " + args[i]); + } + } + } +} diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java new file mode 100644 index 00000000000..d258f94334d --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java @@ -0,0 +1,313 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench; + +import java.io.File; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; + +import ai.rapids.cudf.ColumnVector; +import ai.rapids.cudf.Cuda; +import ai.rapids.cudf.CudaMemInfo; +import ai.rapids.cudf.HostColumnVector; +import ai.rapids.cudf.Rmm; +import ai.rapids.cudf.RmmAllocationMode; +import ai.rapids.cudf.Table; + +/** + * Microbenchmark runner for CPU vs. RapidsUDF. Measures UDF execution time on in-memory dataset. + * + * Reads Parquet file (produced by GenData) via cuDF Table.readParquet. + * Benchmarks CPU (row-by-row evaluate) and GPU (evaluateColumnar) paths. + * Data loading and host/device transfers are not part of timing. + * + * Usage: + * mvn exec:java -Dexec.mainClass=com.udf.bench.MicroBenchRunner \ + * -Dexec.args="--mode all --data-path data/bench_data --rows 1000000" + */ +public class MicroBenchRunner { + + private static final int DEFAULT_WARMUP = 2; + private static final int DEFAULT_MEASURED = 4; + private static final float DEFAULT_RMM_ALLOC_FRACTION = 0.9f; + + /** + * TODO: Extract column data from host memory into Java objects. + * + * Called once before CPU timing loop. Convert HostColumnVectors to + * array of Java objects for executeCpu. + * Use hostColumns[i].getJavaString(row), .getInt(row), .getDouble(row), + * .getStruct(row), .getList(row), etc. to extract values into typed arrays. + * + * This is outside of the timing loop due to overhead of extracting/boxing + * Java types from cuDF. + * + * Example for a UDF that takes (String, int): + * 
{@code
+     *   String[] col0 = new String[numRows];
+     *   int[] col1 = new int[numRows];
+     *   for (int i = 0; i < numRows; i++) {
+     *       col0[i] = hostColumns[0].getJavaString(i);
+     *       col1[i] = hostColumns[1].getInt(i);
+     *   }
+     *   return new Object[] { col0, col1 };
+     * }
+ * + * @param hostColumns all columns copied to host memory + * @param numRows number of rows in the dataset + * @return array of typed Java arrays, one per UDF input column + */ + public static Object[] prepareCpuData(HostColumnVector[] hostColumns, int numRows) { + // TODO: Extract columns to Java arrays + return null; // TODO + } + + /** + * TODO: Execute the CPU UDF on Java data row-by-row. + * + * Example: + * 
{@code
+     *   import com.udf.PlaceholderUDFName;
+     *   String[] col0 = (String[]) data[0];
+     *   int[] col1 = (int[]) data[1];
+     *   PlaceholderUDFName udf = new PlaceholderUDFName();
+     *   for (int i = 0; i < numRows; i++) {
+     *       udf.evaluate(col0[i], col1[i]);
+     *   }
+     * }
+ * + * @param data Java arrays from {@link #prepareCpuData} + * @param numRows number of rows in the dataset + */ + public static void executeCpu(Object[] data, int numRows) { + // TODO: Cast arrays and call CPU UDF evaluate() per row + } + + /** + * TODO: Execute the GPU UDF via evaluateColumnar. + * + * Example: + * 
{@code
+     *   import com.udf.PlaceholderRapidsUDFName;
+     *   PlaceholderRapidsUDFName udf = new PlaceholderRapidsUDFName();
+     *   return udf.evaluateColumnar(numRows,
+     *       table.getColumn(0), table.getColumn(1));
+     * }
+ * + * @param table the dataset loaded on GPU + * @param numRows number of rows in the dataset + * @return result ColumnVector (NOTE: caller must close) + */ + public static ColumnVector executeGpu(Table table, int numRows) { + // TODO: Instantiate RapidsUDF and call evaluateColumnar() + return null; // TODO + } + + public static void main(String[] args) { + Map argMap = new HashMap<>(); + parseArgs(args, argMap); + + String dataPath = argMap.get("data-path"); + if (dataPath == null) { + throw new IllegalArgumentException("--data-path is required"); + } + String mode = argMap.getOrDefault("mode", "all"); + int maxRows = Integer.parseInt(argMap.getOrDefault("rows", "-1")); + float rmmAllocFraction = Float.parseFloat(argMap.getOrDefault("pool-fraction", String.valueOf(DEFAULT_RMM_ALLOC_FRACTION))); + int warmup = Integer.parseInt(argMap.getOrDefault("warmup", String.valueOf(DEFAULT_WARMUP))); + int measured = Integer.parseInt(argMap.getOrDefault("measured", String.valueOf(DEFAULT_MEASURED))); + boolean profile = argMap.containsKey("profile"); + + // Resolve execution mode + if (!"cpu".equals(mode) && !"gpu".equals(mode) && !"all".equals(mode)) { + throw new IllegalArgumentException( + "Unknown mode: '" + mode + "'. Must be 'cpu', 'gpu', or 'all'."); + } + boolean runCpu = "cpu".equals(mode) || "all".equals(mode); + boolean runGpu = "gpu".equals(mode) || "all".equals(mode); + + // Initialize RMM pool + if (!Rmm.isInitialized()) { + CudaMemInfo memInfo = Cuda.memGetInfo(); + long poolSize = (long) (memInfo.free * rmmAllocFraction) & ~255L; + Rmm.initialize(RmmAllocationMode.POOL, null, poolSize); + } + + // Read Parquet data into cuDF table + try (Table table = readParquetData(dataPath, maxRows)) { + int numRows = (int) table.getRowCount(); + int numCols = table.getNumberOfColumns(); + double mb = getTableSizeMB(table); + System.out.printf("Loaded %,d rows x %d columns (%.1f MB) from: %s%n", + numRows, numCols, mb, dataPath); + System.out.printf("Microbenchmark: mode=%s, warmup=%d, measured=%d%n", + mode, warmup, measured); + + double cpuMinMs = Double.NaN; + double gpuMinMs = Double.NaN; + + // --- CPU Benchmark --- + if (runCpu) { + HostColumnVector[] hostColumns = copyAllToHost(table); + try { + Object[] cpuData = prepareCpuData(hostColumns, numRows); + long[] times = runBenchmark(warmup, measured, false, () -> + executeCpu(cpuData, numRows)); + double medianMs = times[times.length / 2] / 1e6; + cpuMinMs = times[0] / 1e6; + System.out.printf(" CPU | %,14d rows | median %10.1f ms | min %10.1f ms%n", + numRows, medianMs, cpuMinMs); + } catch (Exception e) { + System.err.printf("CPU benchmark failed: %s%n", e.getMessage()); + e.printStackTrace(System.err); + } finally { + closeAll(hostColumns); + } + } + + // --- GPU Benchmark --- + if (runGpu) { + try { + long[] times = runBenchmark(warmup, measured, profile, () -> { + try (ColumnVector result = executeGpu(table, numRows)) {} + }); + double medianMs = times[times.length / 2] / 1e6; + gpuMinMs = times[0] / 1e6; + System.out.printf(" GPU | %,14d rows | median %10.1f ms | min %10.1f ms%n", + numRows, medianMs, gpuMinMs); + } catch (Exception e) { + System.err.printf("GPU benchmark failed: %s%n", e.getMessage()); + e.printStackTrace(System.err); + } + } + + // --- Speedup --- + if (!Double.isNaN(cpuMinMs) && !Double.isNaN(gpuMinMs)) { + double speedup = cpuMinMs / gpuMinMs; + System.out.printf(">> Speedup: %.2fx (CPU/GPU best)%n", speedup); + } + } + + System.exit(0); + } + + /** + * Run warmup + measured iterations. Profile the measured iterations if enabled. + * @return sorted array of measured elapsed times in nanoseconds + */ + private static long[] runBenchmark(int warmup, int measured, boolean profile, Runnable block) { + for (int i = 0; i < warmup; i++) { + block.run(); + } + long[] times = new long[measured]; + for (int i = 0; i < measured; i++) { + if (profile) Cuda.profilerStart(); + long start = System.nanoTime(); + block.run(); + times[i] = System.nanoTime() - start; + if (profile) Cuda.profilerStop(); + } + Arrays.sort(times); + return times; + } + + /** + * Read Parquet partition files from a directory into a cuDF Table. + * Reads files in sorted order, stopping once maxRows is reached. + * @param maxRows stop after accumulating this many rows; -1 means read all. + */ + private static Table readParquetData(String dataPath, int maxRows) { + File[] partFiles = new File(dataPath).listFiles((dir, name) -> name.endsWith(".parquet")); + if (partFiles == null || partFiles.length == 0) { + throw new IllegalArgumentException("No .parquet files found in: " + dataPath); + } + Arrays.sort(partFiles); + + Table[] tables = new Table[partFiles.length]; + int count = 0; + long totalRows = 0; + try { + for (int i = 0; i < partFiles.length; i++) { + tables[i] = Table.readParquet(partFiles[i]); + count++; + totalRows += tables[i].getRowCount(); + if (maxRows > 0 && totalRows >= maxRows) break; + } + Table combined = (count == 1) ? tables[0] : Table.concatenate(Arrays.copyOf(tables, count)); + try (Table src = combined) { + return limitTable(src, maxRows); + } + } finally { + if (count > 1) { + closeAll(tables); + } + } + } + + /** Return a new Table with at most numRows rows. */ + private static Table limitTable(Table table, int numRows) { + int n = (numRows <= 0) + ? (int) table.getRowCount() + : (int) Math.min(numRows, table.getRowCount()); + ColumnVector[] cols = new ColumnVector[table.getNumberOfColumns()]; + try { + for (int i = 0; i < cols.length; i++) { + cols[i] = table.getColumn(i).subVector(0, n); + } + return new Table(cols); + } finally { + closeAll(cols); + } + } + + /** Get the size of the table in MB. */ + private static double getTableSizeMB(Table table) { + long bytes = 0; + for (int i = 0; i < table.getNumberOfColumns(); i++) { + bytes += table.getColumn(i).getDeviceMemorySize(); + } + return bytes / (1024.0 * 1024.0); + } + + /** Copy all device columns to host memory. */ + private static HostColumnVector[] copyAllToHost(Table table) { + HostColumnVector[] hostCols = new HostColumnVector[table.getNumberOfColumns()]; + for (int i = 0; i < hostCols.length; i++) { + hostCols[i] = table.getColumn(i).copyToHost(); + } + return hostCols; + } + + /** Close all resources in an array. */ + private static void closeAll(AutoCloseable[] resources) { + if (resources != null) { + for (AutoCloseable r : resources) { + if (r != null) { + try { r.close(); } catch (Exception ignore) {} + } + } + } + } + + /** Parse CLI arguments. */ + private static void parseArgs(String[] args, Map map) { + int i = 0; + while (i < args.length) { + switch (args[i]) { + case "--mode": map.put("mode", args[i + 1]); i += 2; break; + case "--data-path": map.put("data-path", args[i + 1]); i += 2; break; + case "--warmup": map.put("warmup", args[i + 1]); i += 2; break; + case "--measured": map.put("measured", args[i + 1]); i += 2; break; + case "--rows": map.put("rows", args[i + 1]); i += 2; break; + case "--pool-fraction": map.put("pool-fraction", args[i + 1]); i += 2; break; + case "--profile": map.put("profile", "true"); i += 1; break; + default: + throw new IllegalArgumentException("Unknown argument: " + args[i]); + } + } + } +} diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/SparkBenchRunner.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/SparkBenchRunner.java new file mode 100644 index 00000000000..a165a38d5ac --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/SparkBenchRunner.java @@ -0,0 +1,163 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench; + +import java.io.File; +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import com.fasterxml.jackson.core.util.DefaultIndenter; +import com.fasterxml.jackson.core.util.DefaultPrettyPrinter; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.SerializationFeature; +import com.udf.SparkUtils; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; + +/** + * UDF benchmark runner. Measures the end-to-end runtime of: + * Read Parquet -> Execute (CPU or GPU) -> Write no-op sink + * + * Produces a JSON file with the benchmark results. + * On error, also produces a separate error log file. + * + * Usage: + * mvn exec:java -Dexec.mainClass=com.udf.bench.SparkBenchRunner \ + * -Dexec.args="--mode cpu --data-path data/bench_data_10M_rows.parquet ..." + */ +public class SparkBenchRunner { + + private static final String DEFAULT_SPARK_LOG_LEVEL = "ERROR"; + + public static void main(String[] args) { + Map argMap = new HashMap<>(); + List sparkConfs = new ArrayList<>(); + parseArgs(args, argMap, sparkConfs); + + String mode = SparkUtils.requireArg(argMap, "mode"); + String dataPath = SparkUtils.requireArg(argMap, "data-path"); + String resultPath = SparkUtils.requireArg(argMap, "result-path"); + String sparkLogLevel = argMap.getOrDefault("spark-log-level", DEFAULT_SPARK_LOG_LEVEL); + + // Validate mode + if (!"cpu".equals(mode) && !"gpu".equals(mode)) { + throw new IllegalArgumentException( + "Unknown mode: '" + mode + "'. Must be 'cpu' or 'gpu'."); + } + + // Build Spark session + SparkSession.Builder builder = SparkSession.builder(); + SparkUtils.applySparkConfs(builder, sparkConfs); + SparkSession spark = builder.enableHiveSupport().getOrCreate(); + spark.sparkContext().setLogLevel(sparkLogLevel); + + try { + // --- START JOB --- + long startTime = System.nanoTime(); + Dataset df = spark.read().parquet(dataPath); + Dataset resultDf = "cpu".equals(mode) + ? BenchUtils.executeCpu(spark, df) + : BenchUtils.executeGpu(spark, df); + resultDf.write().format("noop").mode("overwrite").save(); + double elapsed = (System.nanoTime() - startTime) / 1e9; + // --- END JOB --- + + System.err.printf("E2E Runtime (s): %.2f%n", elapsed); + + writeReport(resultPath, mode, dataPath, elapsed, "success", args, null, null); + + } catch (Exception e) { + System.err.println("Benchmark run failed: " + e.getClass().getSimpleName()); + e.printStackTrace(System.err); + + // Error stack trace is written to a separate error log file. + String errorLogPath = resultPath.replace("_result.json", "_error.log"); + writeErrorLog(errorLogPath, e); + + writeReport(resultPath, mode, dataPath, -1, "error", args, + e.getMessage(), errorLogPath); + + System.exit(1); + } finally { + spark.stop(); + } + + System.exit(0); + } + + /** Write a JSON benchmark report containing the result and args. */ + private static void writeReport( + String path, String mode, String dataPath, double elapsed, + String status, String[] cliArgs, + String errorMessage, String errorLogFile) { + File resultDir = new File(path).getParentFile(); + if (resultDir != null) resultDir.mkdirs(); + + try { + Map report = new LinkedHashMap<>(); + report.put("mode", mode); + report.put("data_path", dataPath); + report.put("status", status); + report.put("e2e_runtime", elapsed); + report.put("cli_args", Arrays.asList(cliArgs)); + if (errorMessage != null) { + Map error = new LinkedHashMap<>(); + error.put("error_message", errorMessage); + if (errorLogFile != null) { + error.put("error_log_file", errorLogFile); + } + report.put("error", error); + } + + ObjectMapper mapper = new ObjectMapper(); + mapper.enable(SerializationFeature.INDENT_OUTPUT); + DefaultPrettyPrinter printer = new DefaultPrettyPrinter(); + printer.indentArraysWith(DefaultIndenter.SYSTEM_LINEFEED_INSTANCE); + mapper.writer(printer).writeValue(new File(path), report); + System.err.println("Report written to: " + path); + } catch (Exception e) { + System.err.println("Failed to write report: " + e.getMessage()); + } + } + + /** Write an exception to an error log file. */ + private static void writeErrorLog(String path, Exception e) { + File logDir = new File(path).getParentFile(); + if (logDir != null) logDir.mkdirs(); + + try (PrintWriter pw = new PrintWriter(path)) { + StringWriter sw = new StringWriter(); + e.printStackTrace(new PrintWriter(sw)); + pw.print(sw.toString()); + } catch (Exception writeErr) { + System.err.println("Failed to write error log: " + writeErr.getMessage()); + } + System.err.println("Error details written to: " + path); + } + + /** Parse CLI arguments. */ + private static void parseArgs(String[] args, Map map, List sparkConfs) { + int i = 0; + while (i < args.length) { + switch (args[i]) { + case "--mode": map.put("mode", args[i + 1]); i += 2; break; + case "--data-path": map.put("data-path", args[i + 1]); i += 2; break; + case "--result-path": map.put("result-path", args[i + 1]); i += 2; break; + case "--spark-log-level": map.put("spark-log-level", args[i + 1]); i += 2; break; + case "--spark-conf": sparkConfs.add(args[i + 1]); i += 2; break; + default: + throw new IllegalArgumentException("Unknown argument: " + args[i]); + } + } + } +} diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java new file mode 100644 index 00000000000..8217963425d --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java @@ -0,0 +1,70 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +public class CudfComparisonTest { + + private static SparkSession spark; + private static ClassLoader origContextClassLoader; + + @BeforeClass + public static void setUp() { + origContextClassLoader = TestUtils.installMutableClassLoader(); + spark = SparkSession.builder() + .appName("UDF vs. RapidsUDF Comparison Test") + .master("local[*]") + .config("spark.plugins", "com.nvidia.spark.SQLPlugin") + .config("spark.rapids.memory.gpu.pool", "NONE") + .config("spark.rapids.sql.explain", "NONE") + .enableHiveSupport() + .getOrCreate(); + } + + @AfterClass + public static void tearDown() { + if (spark != null) spark.stop(); + if (origContextClassLoader != null) { + Thread.currentThread().setContextClassLoader(origContextClassLoader); + } + } + + /** TODO: Register the RapidsUDF with Spark. */ + public static void registerRapidsUDF(SparkSession spark, String udfName) { } + + @Test + public void testCpuVsRapidsUDF() { + Dataset testDF = UnitTest.createTestData(spark).repartition(1); + + // Run CPU UDF + UnitTest.registerUDF(spark, "placeholder_udf_name"); + Dataset cpuResultDF = UnitTest.executeUDF( + spark, "placeholder_udf_name", testDF); + UnitTest.verifyUDFResults(cpuResultDF, testDF); + + // Run RapidsUDF + registerRapidsUDF(spark, "placeholder_rapids_udf_name"); + Dataset gpuResultDF = UnitTest.executeUDF( + spark, "placeholder_rapids_udf_name", testDF); + UnitTest.verifyUDFResults(gpuResultDF, testDF); + + // Compare + TestUtils.assertDataFrameEquals(gpuResultDF, cpuResultDF); + } + + /** + * TODO: If UnitTest adds extra @Test methods beyond the main result checks, + * add corresponding comparison tests here. Each case should run the same input + * through the CPU UDF and the RapidsUDF, apply equivalent assertions to both + * outputs, and compare the RapidsUDF output against the CPU output. + */ +} diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java new file mode 100644 index 00000000000..d335432d920 --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java @@ -0,0 +1,76 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +public class SqlComparisonTest { + + private static SparkSession spark; + private static ClassLoader origContextClassLoader; + + @BeforeClass + public static void setUp() { + origContextClassLoader = TestUtils.installMutableClassLoader(); + spark = SparkSession.builder() + .appName("UDF vs. SQL Comparison Test") + .master("local[*]") + .config("spark.plugins", "com.nvidia.spark.SQLPlugin") + .config("spark.rapids.skipGpuArchitectureCheck", "true") + .config("spark.rapids.sql.mode", "explainOnly") + .config("spark.sql.adaptive.enabled", "false") + .enableHiveSupport() + .getOrCreate(); + } + + @AfterClass + public static void tearDown() { + if (spark != null) spark.stop(); + if (origContextClassLoader != null) { + Thread.currentThread().setContextClassLoader(origContextClassLoader); + } + } + + @Test + public void testUdfVsSqlExpression() throws IOException { + Dataset testDF = UnitTest.createTestData(spark).repartition(1); + + // Run CPU UDF + UnitTest.registerUDF(spark, "placeholder_udf_name"); + Dataset udfResultDF = UnitTest.executeUDF( + spark, "placeholder_udf_name", testDF); + UnitTest.verifyUDFResults(udfResultDF, testDF); + + // Read and execute SQL expression + testDF.createOrReplaceTempView("test_table"); + String sqlContent = new String( + Files.readAllBytes(Paths.get("src/main/resources/placeholder_udf_name.sql"))); + Dataset sqlResultDF = spark.sql(sqlContent); + UnitTest.verifyUDFResults(sqlResultDF, testDF); + + // Compare + TestUtils.assertDataFrameEquals(sqlResultDF, udfResultDF); + + // Verify GPU compatibility + SparkUtils.assertPlanRunsOnGpu(sqlResultDF); + } + + /** + * TODO: If UnitTest adds extra @Test methods beyond the main result checks, + * add corresponding comparison tests here. Each case should run the same input + * through the CPU UDF and the SQL expression, apply equivalent assertions to + * both outputs, and compare the SQL output against the CPU output. + */ +} diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/TestUtils.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/TestUtils.java new file mode 100644 index 00000000000..3beef6e6424 --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/TestUtils.java @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import java.net.URL; +import java.net.URLClassLoader; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.junit.Assert; + +/** + * Shared test utilities. + */ +public class TestUtils { + + /** + * Install a URLClassLoader as the thread context classloader so that + * RAPIDS ShimLoader.findURLClassLoader() can discover and mutate it. + * https://github.com/NVIDIA/spark-rapids/blob/main/sql-plugin-api/src/main/scala/com/nvidia/spark/rapids/ShimLoader.scala + * + * On Java 17 in a forked Surefire JVM the only classloader is + * AppClassLoader, which is not a URLClassLoader. Without a URL CL + * the RAPIDS ShimLoader will throw since it will fail to install + * shim classes, e.g. https://github.com/NVIDIA/spark-rapids/issues/13915. + * + * Must be called before plugin initialization, i.e., before SparkSession.getOrCreate(). + * Returns the original context classloader for the caller to restore on tearDown. + */ + public static ClassLoader installMutableClassLoader() { + ClassLoader original = Thread.currentThread().getContextClassLoader(); + if (original instanceof URLClassLoader) { + return original; + } + // Create a child URLClassLoader of original AppClassLoader with empty search path. + // ShimLoader will populate w/shim directories via addURL(). + URLClassLoader wrapper = new URLClassLoader(new URL[0], original); + Thread.currentThread().setContextClassLoader(wrapper); + return original; + } + + /** Compare two DataFrames row-by-row, reporting per-column mismatches. */ + public static void assertDataFrameEquals(Dataset actual, Dataset expected) { + Assert.assertEquals("Schema mismatch", expected.schema(), actual.schema()); + + Row[] actualRows = (Row[]) actual.collect(); + Row[] expectedRows = (Row[]) expected.collect(); + Arrays.sort(actualRows, (a, b) -> a.toString().compareTo(b.toString())); + Arrays.sort(expectedRows, (a, b) -> a.toString().compareTo(b.toString())); + + Assert.assertEquals("Row count mismatch", expectedRows.length, actualRows.length); + + List mismatches = new ArrayList<>(); + String[] fields = actual.schema().fieldNames(); + for (int i = 0; i < actualRows.length; i++) { + for (String field : fields) { + Object aVal = actualRows[i].getAs(field); + Object eVal = expectedRows[i].getAs(field); + boolean eq = (aVal == null && eVal == null) + || (aVal != null && aVal.equals(eVal)); + if (!eq) { + mismatches.add(String.format(" [row %d] %s: actual=%s, expected=%s", + i, field, aVal, eVal)); + } + } + } + if (!mismatches.isEmpty()) { + Assert.fail("\nFound " + mismatches.size() + " column-level mismatches:\n" + + String.join("\n", mismatches) + "\n"); + } + } +} diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java new file mode 100644 index 00000000000..c6fb7a26c02 --- /dev/null +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf; + +import java.util.Arrays; +import java.util.List; + +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.AfterClass; +import org.junit.Assert; +import org.junit.BeforeClass; +import org.junit.Test; + +public class UnitTest { + + private static SparkSession spark; + private static ClassLoader origContextClassLoader; + + @BeforeClass + public static void setUp() { + origContextClassLoader = TestUtils.installMutableClassLoader(); + spark = SparkSession.builder() + .appName("UDF Unit Test") + .master("local[*]") + .config("spark.plugins", "com.nvidia.spark.SQLPlugin") + .config("spark.rapids.skipGpuArchitectureCheck", "true") + .config("spark.rapids.sql.mode", "explainOnly") + .config("spark.sql.adaptive.enabled", "false") + .enableHiveSupport() + .getOrCreate(); + } + + @AfterClass + public static void tearDown() { + if (spark != null) spark.stop(); + if (origContextClassLoader != null) { + Thread.currentThread().setContextClassLoader(origContextClassLoader); + } + } + + /** + * TODO: Create a test DataFrame with diverse test cases including edge cases. + * + * Example: + * 
{@code
+     *   StructType schema = new StructType(new StructField[]{
+     *       DataTypes.createStructField("id", DataTypes.IntegerType, false),
+     *       DataTypes.createStructField("credit_score", DataTypes.IntegerType, true)
+     *   });
+     *   List data = Arrays.asList(
+     *       RowFactory.create(1, 800),
+     *       RowFactory.create(2, 550),
+     *       RowFactory.create(3, null)
+     *   );
+     *   return spark.createDataFrame(data, schema);
+     * }
+ */ + public static Dataset createTestData(SparkSession spark) { + return null; // TODO + } + + /** + * TODO: Register the UDF with Spark. + * + * Example (Hive UDF): + * 
{@code
+     *   spark.sql("CREATE TEMPORARY FUNCTION " + udfName
+     *       + " AS 'com.udf.CalculateRiskUDF'");
+     * }
+ */ + public static void registerUDF(SparkSession spark, String udfName) { + // TODO + } + + /** + * TODO: Execute the UDF on the test DataFrame and return the result. + * + * Example: + * 
{@code
+     *   testDF.createOrReplaceTempView("test_table");
+     *   return spark.sql("SELECT *, " + udfName
+     *       + "(credit_score) AS risk_level FROM test_table");
+     * }
+ */ + public static Dataset executeUDF(SparkSession spark, String udfName, Dataset testDF) { + return null; // TODO + } + + /** + * TODO: Verify UDF results using Assert statements. + * + * Example: + * 
{@code
+     *   Row[] results = (Row[]) resultDF.sort("id").collect();
+     *   Assert.assertEquals("LOW", results[0].getAs("risk_level"));
+     *   Assert.assertEquals("MEDIUM", results[1].getAs("risk_level"));
+     *   Assert.assertEquals("UNKNOWN", results[2].getAs("risk_level"));
+     * }
+ */ + public static void verifyUDFResults(Dataset resultDF, Dataset testDF) { + // TODO + } + + @Test + public void testUDFProducesCorrectResults() { + Dataset testDF = createTestData(spark).repartition(1); + + registerUDF(spark, "placeholder_udf_name"); + Dataset resultDF = executeUDF(spark, "placeholder_udf_name", testDF); + + verifyUDFResults(resultDF, testDF); + } +} diff --git a/skills/udf-gen-test/templates/scala/.mvn/jvm.config b/skills/udf-gen-test/templates/scala/.mvn/jvm.config new file mode 100644 index 00000000000..4d641a99f65 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/.mvn/jvm.config @@ -0,0 +1,15 @@ +-Xmx5g +--add-opens=java.base/java.lang=ALL-UNNAMED +--add-opens=java.base/java.lang.invoke=ALL-UNNAMED +--add-opens=java.base/java.lang.reflect=ALL-UNNAMED +--add-opens=java.base/java.io=ALL-UNNAMED +--add-opens=java.base/java.net=ALL-UNNAMED +--add-opens=java.base/java.nio=ALL-UNNAMED +--add-opens=java.base/java.util=ALL-UNNAMED +--add-opens=java.base/java.util.concurrent=ALL-UNNAMED +--add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED +--add-opens=java.base/sun.nio.ch=ALL-UNNAMED +--add-opens=java.base/sun.nio.cs=ALL-UNNAMED +--add-opens=java.base/sun.security.action=ALL-UNNAMED +--add-opens=java.base/sun.util.calendar=ALL-UNNAMED +--add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED diff --git a/skills/udf-gen-test/templates/scala/pom.xml b/skills/udf-gen-test/templates/scala/pom.xml new file mode 100644 index 00000000000..132034c8421 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/pom.xml @@ -0,0 +1,382 @@ + + + + 4.0.0 + com.udf + aether-agent-udfs + 1.0.0 + Aether UDF Conversion + This project contains UDFs that will be converted from CPU to GPU. + jar + + + 17 + 17 + UTF-8 + UTF-8 + UTF-8 + 2.12 + 2.12.15 + + 3.5.5 + 26.04.0 + 2.1.0 + 2.0.11 + ${test.jvm.args} + cuda12 + v26.04.00 + v26.04.00 + + + false + off + + ON + RAPIDS + 10 + ON + OFF + false + rapidsudfjni + ${project.build.directory}/native-build + + + -Xmx5g -ea + -Dai.rapids.refcount.debug=${debug.memory.leaks} + -Dorg.slf4j.simpleLogger.defaultLogLevel=off + -Dorg.slf4j.simpleLogger.log.ai.rapids.cudf=${cudf.log.level} + --add-opens=java.base/java.lang=ALL-UNNAMED + --add-opens=java.base/java.lang.invoke=ALL-UNNAMED + --add-opens=java.base/java.lang.reflect=ALL-UNNAMED + --add-opens=java.base/java.io=ALL-UNNAMED + --add-opens=java.base/java.net=ALL-UNNAMED + --add-opens=java.base/java.nio=ALL-UNNAMED + --add-opens=java.base/java.util=ALL-UNNAMED + --add-opens=java.base/java.util.concurrent=ALL-UNNAMED + --add-opens=java.base/java.util.concurrent.atomic=ALL-UNNAMED + --add-opens=java.base/sun.nio.ch=ALL-UNNAMED + --add-opens=java.base/sun.nio.cs=ALL-UNNAMED + --add-opens=java.base/sun.security.action=ALL-UNNAMED + --add-opens=java.base/sun.util.calendar=ALL-UNNAMED + --add-opens=java.security.jgss/sun.security.krb5=ALL-UNNAMED + + + + + debug-leaks + + + debug.memory.leaks + true + + + + error + + + + coverage + + + + org.scoverage + scoverage-maven-plugin + ${scoverage.plugin.version} + + ${scala.version} + ${scoverage.scalac.plugin.version} + true + com\.udf\.bench\..*;com\.udf\.SparkUtils.*;com\.udf\.Arm.* + + + + + + + cuda-native-udf + + + + org.apache.maven.plugins + maven-dependency-plugin + 3.6.1 + + + copy-rapids-jar-with-classifier + generate-sources + + copy + + + + + com.nvidia + rapids-4-spark_${scala.binary.version} + ${rapids4spark.version} + ${cuda.version} + jar + false + ${project.build.directory}/rapids-jar + + + true + + + + copy-rapids-jar-no-classifier + generate-sources + + copy + + + + + com.nvidia + rapids-4-spark_${scala.binary.version} + ${rapids4spark.version} + jar + false + ${project.build.directory}/rapids-jar + + + true + + + + + + org.apache.maven.plugins + maven-antrun-plugin + 3.1.0 + + + extract-cuda-native-dependencies + generate-sources + + ${skipCudfExtraction} + + + + + + + + + + + + + run + + + + cmake-cuda-native-udf + compile + + + + + + + + + + + + + + + + + + + + + + + + + + run + + + + + + org.apache.maven.plugins + maven-resources-plugin + 3.3.1 + + + copy-cuda-native-library-to-classes + process-classes + + copy-resources + + + true + ${project.build.outputDirectory}/${os.arch}/${os.name} + + + ${native.build.path} + + lib${native.library.name}.so + + + + + + + + + + + + + + + + org.apache.spark + spark-hive_${scala.binary.version} + ${spark.version} + provided + + + + org.scala-lang + scala-library + ${scala.version} + + + + com.nvidia + rapids-4-spark_${scala.binary.version} + ${rapids4spark.version} + provided + + + + org.scalatest + scalatest_${scala.binary.version} + 3.2.17 + test + + + + org.slf4j + slf4j-simple + 1.7.36 + test + + + + + src/main/scala + src/test/scala + + + + org.codehaus.mojo + build-helper-maven-plugin + 3.5.0 + + + add-java-source + generate-sources + + add-source + + + + src/main/java + + + + + + + + net.alchim31.maven + scala-maven-plugin + 4.3.0 + + + + compile + testCompile + + + + + ${scala.version} + + + + + + org.apache.maven.plugins + maven-surefire-plugin + 3.1.2 + + true + + + + + org.scalatest + scalatest-maven-plugin + 2.2.0 + + ${project.build.directory}/surefire-reports + ${scalatest.jvm.args} + + + + test + + test + + + + + + + org.codehaus.mojo + exec-maven-plugin + 3.1.0 + + + + org.apache.maven.plugins + maven-shade-plugin + 3.2.4 + + + package + + shade + + + false + + + *:* + + META-INF/*.SF + META-INF/*.DSA + META-INF/*.RSA + + + + + + + + + + diff --git a/skills/udf-gen-test/templates/scala/run_gen_data.sh b/skills/udf-gen-test/templates/scala/run_gen_data.sh new file mode 100644 index 00000000000..44c802c1864 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/run_gen_data.sh @@ -0,0 +1,73 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Generate or validate benchmark data + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +print_usage() { + echo "Usage: $0 --rows NUM [--validate] [--output-path PATH] [--mvn-arg ARG]..." +} + +ROWS="" +VALIDATE="" +OUTPUT_PATH="" +MAVEN_ARGS=() + +while [[ $# -gt 0 ]]; do + case $1 in + --rows) ROWS="$2"; shift 2;; + --validate) VALIDATE="true"; shift;; + --output-path) OUTPUT_PATH="$2"; shift 2;; + --mvn-arg) MAVEN_ARGS+=("$2"); shift 2;; + *) + echo "Unknown option: $1" + print_usage + exit 1 + ;; + esac +done + +if [ -z "$ROWS" ]; then + echo "Error: --rows is required" + print_usage + exit 1 +fi + +SPARK_CONFS=( + --spark-conf spark.master="local[*]" + --spark-conf spark.driver.memory="16g" + --spark-conf spark.rapids.sql.enabled="true" + --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" + --spark-conf spark.locality.wait="0s" + --spark-conf spark.sql.cache.serializer="com.nvidia.spark.ParquetCachedBatchSerializer" + --spark-conf spark.rapids.sql.format.parquet.reader.type="MULTITHREADED" + --spark-conf spark.rapids.sql.reader.batchSizeBytes="1000MB" + --spark-conf spark.sql.files.maxPartitionBytes="512MB" + --spark-conf spark.rapids.sql.metrics.level="DEBUG" +) + +EXEC_ARGS="--rows $ROWS --partitions 32" +for arg in "${SPARK_CONFS[@]}"; do + EXEC_ARGS="$EXEC_ARGS $arg" +done + +if [ -n "$VALIDATE" ]; then + EXEC_ARGS="$EXEC_ARGS --validate" + echo "Running GenData in validation mode with $ROWS rows..." +else + if [ -z "$OUTPUT_PATH" ]; then + OUTPUT_PATH="data/bench_data_${ROWS}_rows.parquet" + fi + EXEC_ARGS="$EXEC_ARGS --output-path $OUTPUT_PATH" + echo "Running GenData to generate $ROWS rows -> $OUTPUT_PATH..." +fi + +mvn "${MAVEN_ARGS[@]}" compile exec:java \ + -Dexec.mainClass="com.udf.bench.GenData" \ + -Dexec.classpathScope=compile \ + -Dexec.args="$EXEC_ARGS" diff --git a/skills/udf-gen-test/templates/scala/run_micro_benchmark.sh b/skills/udf-gen-test/templates/scala/run_micro_benchmark.sh new file mode 100644 index 00000000000..6cac6d3e8c9 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/run_micro_benchmark.sh @@ -0,0 +1,60 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Run in-memory microbenchmark for RapidsUDFs. + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +print_usage() { + echo "Usage: $0 --mode cpu|gpu|all --data-path PATH [--rows N] [--warmup N] [--measured N] [--pool-fraction F] [--profile] [--mvn-arg ARG]..." +} + +MODE="" +DATA_PATH="" +PROFILE="" +MAVEN_ARGS=() +RUNNER_ARGS=() + +while [[ $# -gt 0 ]]; do + case $1 in + --mode) MODE="$2"; RUNNER_ARGS+=("$1" "$2"); shift 2;; + --data-path) DATA_PATH="$2"; RUNNER_ARGS+=("$1" "$2"); shift 2;; + --profile) PROFILE="true"; RUNNER_ARGS+=("$1"); shift;; + --mvn-arg) MAVEN_ARGS+=("$2"); shift 2;; + *) RUNNER_ARGS+=("$1"); shift;; + esac +done + +if [ -z "$MODE" ] || [ -z "$DATA_PATH" ]; then + echo "Error: --mode and --data-path are required" + print_usage + exit 1 +fi + +MVN_CMD=( + mvn "${MAVEN_ARGS[@]}" compile exec:java + -Dexec.mainClass=com.udf.bench.MicroBenchRunner + -Dexec.classpathScope=compile + "-Dexec.args=${RUNNER_ARGS[*]}" +) + +if [ -n "$PROFILE" ]; then + REPORT_PATH="results/microbench_$(date +%Y%m%d_%H%M%S)" + mkdir -p results + echo "Running microbenchmark (mode=$MODE) on $DATA_PATH with nsys profiling..." + echo "nsys report will be saved to: ${REPORT_PATH}.nsys-rep" + nsys profile \ + -c cudaProfilerApi \ + --capture-range-end=stop \ + --trace=cuda,nvtx \ + --nvtx-domain-include="libcudf" \ + -o "$REPORT_PATH" \ + "${MVN_CMD[@]}" +else + echo "Running microbenchmark (mode=$MODE) on $DATA_PATH..." + "${MVN_CMD[@]}" +fi diff --git a/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh b/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh new file mode 100644 index 00000000000..5ce799d5413 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh @@ -0,0 +1,70 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Run CPU or GPU Spark benchmark. + +set -e + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +print_usage() { + echo "Usage: $0 --mode cpu|gpu --data-path PATH [--result-path PATH] [--mvn-arg ARG]..." +} + +MODE="" +DATA_PATH="" +RESULT_PATH="" +MAVEN_ARGS=() + +while [[ $# -gt 0 ]]; do + case $1 in + --mode) MODE="$2"; shift 2;; + --data-path) DATA_PATH="$2"; shift 2;; + --result-path) RESULT_PATH="$2"; shift 2;; + --mvn-arg) MAVEN_ARGS+=("$2"); shift 2;; + *) + echo "Unknown option: $1" + print_usage + exit 1 + ;; + esac +done + +if [ -z "$MODE" ] || [ -z "$DATA_PATH" ]; then + echo "Error: --mode and --data-path are required" + print_usage + exit 1 +fi + +DATA_BASENAME=$(basename "$DATA_PATH" .parquet) +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +if [ -z "$RESULT_PATH" ]; then + RESULT_PATH="results/${MODE}_${DATA_BASENAME}_${TIMESTAMP}_result.json" +fi + +SPARK_CONFS=( + --spark-conf spark.master="local[*]" + --spark-conf spark.driver.memory="16g" + --spark-conf spark.rapids.sql.enabled="true" + --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" + --spark-conf spark.locality.wait="0s" + --spark-conf spark.sql.cache.serializer="com.nvidia.spark.ParquetCachedBatchSerializer" + --spark-conf spark.rapids.sql.format.parquet.reader.type="MULTITHREADED" + --spark-conf spark.rapids.sql.reader.batchSizeBytes="1000MB" + --spark-conf spark.sql.files.maxPartitionBytes="512MB" + --spark-conf spark.rapids.sql.metrics.level="DEBUG" +) + +EXEC_ARGS="--mode $MODE --data-path $DATA_PATH --result-path $RESULT_PATH" +for arg in "${SPARK_CONFS[@]}"; do + EXEC_ARGS="$EXEC_ARGS $arg" +done +EXEC_ARGS="$EXEC_ARGS --spark-conf spark.app.name=${MODE}_${DATA_BASENAME}_${TIMESTAMP}" + +echo "Running $MODE benchmark on $DATA_PATH..." +mvn "${MAVEN_ARGS[@]}" compile exec:java \ + -Dexec.mainClass="com.udf.bench.SparkBenchRunner" \ + -Dexec.classpathScope=compile \ + -Dexec.args="$EXEC_ARGS" diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala new file mode 100644 index 00000000000..e65ed1d6a51 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala @@ -0,0 +1,63 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf + +/** + * Automatic resource management (ARM). + */ +object Arm { + /** + * Helper to auto-close GPU resources after use. + * + * @param resource The AutoCloseable resource + * @param f The function to execute with the resource + * @return The result of the function + */ + def withResource[T <: AutoCloseable, R](resource: T)(f: T => R): R = { + try { + f(resource) + } finally { + if (resource != null) { + resource.close() + } + } + } + + /** + * Helper to auto-close GPU resources on an exception. + * + * @param resource The AutoCloseable resource + * @param f The function to execute with the resource + * @return The result of the function + */ + def closeOnExcept[T <: AutoCloseable, R](resource: T)(f: T => R): R = { + try { + f(resource) + } catch { + case e: Exception => + if (resource != null) { + try { + resource.close() + } catch { + case closeException: Exception => + e.addSuppressed(closeException) + } + } + throw e + } + } + + /** + * Close all resources in an array, skipping nulls. + * + * @param resources The array of resources to close + */ + def closeAll[T <: AutoCloseable](resources: Array[T]): Unit = { + if (resources != null) { + resources.foreach(r => if (r != null) r.close()) + } + } +} diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/SparkUtils.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/SparkUtils.scala new file mode 100644 index 00000000000..8684af60cef --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/SparkUtils.scala @@ -0,0 +1,84 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf + +import com.nvidia.spark.rapids.ExplainPlan +import org.apache.spark.sql.{DataFrame, SparkSession} + +/** + * Spark utility methods. + */ +object SparkUtils { + + /** + * Apply key=value Spark configs to a builder. + * + * @param builder the SparkSession builder to configure + * @param sparkConfs "spark.key=value" config strings + * @return the same builder, for chaining + */ + def applySparkConfs( + builder: SparkSession.Builder, + sparkConfs: Seq[String] + ): SparkSession.Builder = { + for (conf <- sparkConfs) { + val kv = conf.split("=", 2) + if (kv.length == 2) builder.config(kv(0), kv(1)) + } + builder + } + + /** + * Ops that cause fallback but can be ignored, since they are strictly used for testing: + * - RDDScanExec/LocalTableScanExec: surfaces due to spark.createDataFrame() + * - CollectLimitExec: surfaces during dataframe collection (e.g. df.show()) + * - ToPrettyString: surfaces due to df.show() + */ + private val IgnoreOperations = Set( + "RDDScanExec", "LocalTableScanExec", "CollectLimitExec", "ToPrettyString" + ) + + /** + * Assert that the DataFrame's plan can run on GPU. + * NOTE: This is only reliable in explainOnly mode, with AQE disabled. + * + * @param df the DataFrame to check + * @param returnFullPlan if true, include the full plan in the error message + * @throws RuntimeException if any operations cannot run on GPU + */ + def assertPlanRunsOnGpu(df: DataFrame, returnFullPlan: Boolean = false): Unit = { + val plan = getGpuPlan(df) + val unsupportedOps = getUnsupportedOps(plan) + if (unsupportedOps.nonEmpty) { + val opsList = unsupportedOps.map(op => s"- $op").mkString("\n") + var errorMsg = s"Some operations cannot run on GPU.\nFound the following unsupported ops:\n$opsList" + if (returnFullPlan) { + errorMsg += s"\n\nFull physical plan:\n$plan" + } + throw new RuntimeException(errorMsg) + } + } + + /** Get the potential GPU plan using the RAPIDS ExplainPlan API. */ + private def getGpuPlan(df: DataFrame): String = { + ExplainPlan.explainPotentialGpuPlan(df, "NOT_ON_GPU") + } + + /** Parse the plan for unsupported operations (lines starting with '!'). */ + private def getUnsupportedOps(plan: String): Seq[String] = { + plan.split("\n").filter(_.trim.startsWith("!")).flatMap { line => + // Each unsupported line looks like: ![Exec] cannot run on GPU + val start = line.indexOf('<') + val end = line.indexOf('>') + if (start >= 0 && end > start) { + val op = line.substring(start + 1, end) + if (!IgnoreOperations.contains(op)) Some(line.trim) else None + } else { + None + } + }.toSeq + } +} diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/BenchUtils.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/BenchUtils.scala new file mode 100644 index 00000000000..e4d5c3471ed --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/BenchUtils.scala @@ -0,0 +1,109 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench + +import org.apache.spark.sql.{DataFrame, SparkSession} +import org.apache.spark.sql.functions._ +import org.apache.spark.sql.types._ + +/** + * Benchmark utilities. + * - generateSyntheticData: Create benchmark data for the UDF + * - executeCpu: Register and run the CPU UDF + * - executeGpu: Register and run the GPU implementation + */ +object BenchUtils { + + // --------------------------------------------------------------------------- + // Data generation + // --------------------------------------------------------------------------- + + /** + * TODO: Generate a synthetic DataFrame matching the unit test schema. + * + * Use `spark.range(0, numRows, 1, numPartitions)` as the base, then apply + * randomized column generators to produce data matching the UDF's expected input. + * + * Requirements: + * - Column names and types MUST match the unit test dataset schema + * - Data should be realistic and varied (different lengths, edge cases, etc.) + * - For variable-length inputs, generate sizable rows representative of + * enterprise-scale data + * + * Example: + * {{{ + * val baseDF = spark.range(0, numRows, 1, numPartitions) + * baseDF.select( + * col("id"), + * (rand() * 850).cast(IntegerType).alias("credit_score") + * ) + * }}} + * + * @param spark active SparkSession + * @param numRows number of rows to generate + * @param numPartitions number of output partitions + * @return DataFrame with the same schema as the unit test data + */ + def generateSyntheticData( + spark: SparkSession, + numRows: Long, + numPartitions: Int + ): DataFrame = ??? + + // --------------------------------------------------------------------------- + // Execution + // --------------------------------------------------------------------------- + + /** + * TODO: Execute the CPU UDF on the benchmark DataFrame. + * 1. Register the CPU UDF with Spark + * 2. Execute it on `df` + * 3. Return the result DataFrame + * + * Example: + * {{{ + * import com.udf.CalculateRiskUDF + * spark.udf.register("calculate_risk", new CalculateRiskUDF()) + * df.createOrReplaceTempView("bench_table") + * spark.sql("SELECT *, calculate_risk(credit_score) AS risk_level FROM bench_table") + * }}} + * + * @param spark active SparkSession + * @param df input benchmark DataFrame + * @return result DataFrame after applying the CPU UDF + */ + def executeCpu(spark: SparkSession, df: DataFrame): DataFrame = ??? + + /** + * TODO: Execute the GPU implementation on the benchmark DataFrame. + * + * For RapidsUDF - register RapidsUDF and run the same query as executeCpu: + * {{{ + * import com.udf.CalculateRiskRapidsUDF + * spark.udf.register("calculate_risk_rapids", new CalculateRiskRapidsUDF()) + * df.createOrReplaceTempView("bench_table") + * spark.sql("SELECT *, calculate_risk_rapids(credit_score) AS risk_level FROM bench_table") + * }}} + * + * For SQL - read the SQL file from src/main/resources/ and adapt it for + * benchmarking. The SQL was written for the unit test, so you must: + * 1. Replace "test_table" with "bench_table" + * 2. Replace the SELECT column list with "SELECT *" to avoid referencing + * columns that may not exist in the benchmark DataFrame + * {{{ + * df.createOrReplaceTempView("bench_table") + * val sqlContent = scala.io.Source.fromFile("src/main/resources/calculate_risk.sql").mkString + * val benchSql = sqlContent.replace("test_table", "bench_table") + * // Also replace the SELECT column list with SELECT * if needed + * spark.sql(benchSql) + * }}} + * + * @param spark active SparkSession + * @param df input benchmark DataFrame + * @return result DataFrame after applying the GPU implementation + */ + def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = ??? +} diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/GenData.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/GenData.scala new file mode 100644 index 00000000000..31377658c07 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/GenData.scala @@ -0,0 +1,101 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench + +import com.udf.SparkUtils +import org.apache.spark.sql.SparkSession + +/** + * Generates benchmark data and optionally validates by running + * BenchUtils.executeCpu and BenchUtils.executeGpu. + * + * Usage: + * mvn exec:java -Dexec.mainClass=com.udf.bench.GenData \ + * -Dexec.args="--rows 1000 --validate --spark-conf k=v ..." + */ +object GenData { + + def main(args: Array[String]): Unit = { + val (parsed, sparkConfs) = parseArgs(args) + + val rows = parsed.getOrElse("rows", + throw new IllegalArgumentException("--rows is required")).toLong + val partitions = parsed.getOrElse("partitions", "32").toInt + val validate = parsed.contains("validate") + val outputPath = parsed.get("output-path") + + // Build Spark session + val builder = SparkSession.builder().appName("GenData") + SparkUtils.applySparkConfs(builder, sparkConfs) + val spark = builder.getOrCreate() + + try { + // Generate synthetic data + val df = BenchUtils.generateSyntheticData(spark, rows, partitions) + + // Verify row count + val actualRows = df.count() + if (actualRows != rows) { + System.err.println(s"Row count mismatch: expected=$rows, actual=$actualRows") + sys.exit(1) + } + println(s"Generated $actualRows rows across $partitions partitions") + + if (validate) { + // Validation mode — run both CPU and GPU execute, don't write + for ((label, executeFn) <- Seq( + ("cpu", BenchUtils.executeCpu _), + ("gpu", BenchUtils.executeGpu _) + )) { + try { + executeFn(spark, df).collect() + println(s"Validation ($label) passed.") + } catch { + case e: Exception => + System.err.println( + s"Validation ($label) failed: ${e.getClass.getSimpleName}: ${e.getMessage}") + e.printStackTrace(System.err) + sys.exit(1) + } + } + } else { + // Generation mode — write to output path + val path = outputPath.getOrElse( + throw new IllegalArgumentException("--output-path is required when not in validation mode")) + df.write.mode("overwrite").parquet(path) + System.err.println(s"Successfully generated dataset and saved to: $path") + } + } catch { + case e: Exception => + System.err.println(s"Failed to generate dataset: ${e.getClass.getSimpleName}") + e.printStackTrace(System.err) + sys.exit(1) + } finally { + spark.stop() + } + + sys.exit(0) + } + + /** Parse CLI arguments. */ + private def parseArgs(args: Array[String]): (Map[String, String], Seq[String]) = { + var map = Map.empty[String, String] + var sparkConfs = Seq.empty[String] + var i = 0 + while (i < args.length) { + args(i) match { + case "--rows" => map += ("rows" -> args(i + 1)); i += 2 + case "--partitions" => map += ("partitions" -> args(i + 1)); i += 2 + case "--validate" => map += ("validate" -> "true"); i += 1 + case "--output-path" => map += ("output-path" -> args(i + 1)); i += 2 + case "--spark-conf" => sparkConfs :+= args(i + 1); i += 2 + case other => + throw new IllegalArgumentException(s"Unknown argument: $other") + } + } + (map, sparkConfs) + } +} diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala new file mode 100644 index 00000000000..3cf59d07fe1 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala @@ -0,0 +1,281 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench + +import java.io.File + +import scala.collection.mutable.ArrayBuffer + +import ai.rapids.cudf.{ + ColumnVector, + Cuda, + CudaMemInfo, + HostColumnVector, + Rmm, + RmmAllocationMode, + Table +} +import com.udf.Arm.{closeAll, withResource} + +/** + * Microbenchmark runner for CPU vs. RapidsUDF. Measures UDF execution time on in-memory dataset. + * + * Reads Parquet file (produced by GenData) via cuDF Table.readParquet. + * Benchmarks CPU (row-by-row evaluate) and GPU (evaluateColumnar) paths. + * Data loading and host/device transfers are not part of timing. + * + * Usage: + * mvn exec:java -Dexec.mainClass=com.udf.bench.MicroBenchRunner \ + * -Dexec.args="--mode all --data-path data/bench_data --rows 1000000" + */ +object MicroBenchRunner { + + private val DefaultWarmup = 2 + private val DefaultMeasured = 4 + private val DefaultRmmAllocFraction = 0.9f + + /** + * TODO: Extract column data from host memory into Scala objects. + * + * Called once before CPU timing loop. Convert HostColumnVectors to + * array of Scala objects for executeCpu. + * Use hostColumns(i).getJavaString(row), .getInt(row), .getDouble(row), + * .getStruct(row), .getList(row), etc. to extract values into typed arrays. + * + * This is outside of the timing loop due to overhead of extracting/boxing + * Java types from cuDF. + * + * Example for a UDF that takes (String, Int): + * {{{ + * val col0 = Array.tabulate(numRows)(i => hostColumns(0).getJavaString(i)) + * val col1 = Array.tabulate(numRows)(i => hostColumns(1).getInt(i)) + * Array[AnyRef](col0, col1.asInstanceOf[AnyRef]) + * }}} + * + * @param hostColumns all columns copied to host memory + * @param numRows number of rows in the dataset + * @return array of typed arrays, one per UDF input column + */ + def prepareCpuData( + hostColumns: Array[HostColumnVector], + numRows: Int + ): Array[AnyRef] = ??? + + /** + * TODO: Execute the CPU UDF on Scala data row-by-row. + * + * Example: + * {{{ + * val col0 = data(0).asInstanceOf[Array[String]] + * val col1 = data(1).asInstanceOf[Array[Int]] + * val udf = new com.udf.PlaceholderUDFName() + * var i = 0 + * while (i < numRows) { + * udf.apply(col0(i), col1(i)) + * i += 1 + * } + * }}} + * + * @param data typed arrays from [[prepareCpuData]] + * @param numRows number of rows in the dataset + */ + def executeCpu(data: Array[AnyRef], numRows: Int): Unit = ??? + + /** + * TODO: Execute the GPU UDF via evaluateColumnar. + * + * Example: + * {{{ + * val udf = new com.udf.PlaceholderRapidsUDFName() + * udf.evaluateColumnar(numRows, + * table.getColumn(0), table.getColumn(1)) + * }}} + * + * @param table the dataset loaded on GPU + * @param numRows number of rows in the dataset + * @return result ColumnVector (NOTE: caller must close) + */ + def executeGpu(table: Table, numRows: Int): ColumnVector = ??? + + def main(args: Array[String]): Unit = { + val parsed = parseArgs(args) + + val dataPath = parsed.getOrElse("data-path", + throw new IllegalArgumentException("--data-path is required")) + val mode = parsed.getOrElse("mode", "all") + val maxRows = parsed.getOrElse("rows", "-1").toInt + val rmmAllocFraction = parsed.getOrElse("pool-fraction", DefaultRmmAllocFraction.toString).toFloat + val warmup = parsed.getOrElse("warmup", DefaultWarmup.toString).toInt + val measured = parsed.getOrElse("measured", DefaultMeasured.toString).toInt + val profile = parsed.contains("profile") + + mode match { + case "cpu" | "gpu" | "all" => + case other => throw new IllegalArgumentException( + s"Unknown mode: '$other'. Must be 'cpu', 'gpu', or 'all'.") + } + val runCpu = mode == "cpu" || mode == "all" + val runGpu = mode == "gpu" || mode == "all" + + // Initialize RMM pool + if (!Rmm.isInitialized()) { + val memInfo = Cuda.memGetInfo() + val poolSize = (memInfo.free * rmmAllocFraction).toLong & ~255L + Rmm.initialize(RmmAllocationMode.POOL, null, poolSize) + } + + // Read Parquet data into cuDF table + withResource(readParquetData(dataPath, maxRows)) { table => + val numRows = table.getRowCount.toInt + val numCols = table.getNumberOfColumns + val mb = getTableSizeMB(table) + println(f"Loaded $numRows%,d rows x $numCols columns ($mb%.1f MB) from: $dataPath") + println(s"Microbenchmark: mode=$mode, warmup=$warmup, measured=$measured") + + var cpuMinMs: Option[Double] = None + var gpuMinMs: Option[Double] = None + + // --- CPU Benchmark --- + if (runCpu) { + val hostColumns = copyAllToHost(table) + try { + val cpuData = prepareCpuData(hostColumns, numRows) + val times = runBenchmark(warmup, measured) { + executeCpu(cpuData, numRows) + } + val medianMs = times(times.length / 2) / 1e6 + val minMs = times(0) / 1e6 + cpuMinMs = Some(minMs) + println( + f" CPU | $numRows%,14d rows | median $medianMs%10.1f ms | min $minMs%10.1f ms") + } catch { + case e: Exception => + System.err.println(s"CPU benchmark failed: ${e.getMessage}") + e.printStackTrace(System.err) + } finally { + closeAll(hostColumns) + } + } + + // --- GPU Benchmark --- + if (runGpu) { + try { + val times = runBenchmark(warmup, measured, profile = profile) { + withResource(executeGpu(table, numRows)) { _ => } + } + val medianMs = times(times.length / 2) / 1e6 + val minMs = times(0) / 1e6 + gpuMinMs = Some(minMs) + println( + f" GPU | $numRows%,14d rows | median $medianMs%10.1f ms | min $minMs%10.1f ms") + } catch { + case e: Exception => + System.err.println(s"GPU benchmark failed: ${e.getMessage}") + e.printStackTrace(System.err) + } + } + + // --- Speedup --- + for (cpu <- cpuMinMs; gpu <- gpuMinMs) { + val speedup = cpu / gpu + println(f">> Speedup: $speedup%.2fx (CPU/GPU best)") + } + } + + System.exit(0) + } + + /** + * Run warmup + measured iterations. Profile the measured iterations if enabled. + * @return sorted array of measured elapsed times in nanoseconds + */ + private def runBenchmark(warmup: Int, measured: Int, profile: Boolean = false) + (block: => Unit): Array[Long] = { + for (_ <- 0 until warmup) block + (0 until measured).map { i => + if (profile) Cuda.profilerStart() + val start = System.nanoTime() + block + val elapsed = System.nanoTime() - start + if (profile) Cuda.profilerStop() + elapsed + }.toArray.sorted + } + + /** + * Read Parquet partition files from a directory into a cuDF Table. + * Reads files in sorted order, stopping once maxRows is reached. + * @param maxRows stop after accumulating this many rows; -1 means read all. + */ + private def readParquetData(dataPath: String, maxRows: Int): Table = { + val partFiles = new File(dataPath).listFiles((_, name) => name.endsWith(".parquet")) + if (partFiles == null || partFiles.isEmpty) { + throw new IllegalArgumentException(s"No .parquet files found in: $dataPath") + } + + val tables = ArrayBuffer.empty[Table] + var totalRows = 0L + try { + for (f <- partFiles.sorted if maxRows <= 0 || totalRows < maxRows) { + val t = Table.readParquet(f) + tables += t + totalRows += t.getRowCount + } + val combined = if (tables.length == 1) tables(0) + else Table.concatenate(tables.toArray: _*) + withResource(combined) { src => limitTable(src, maxRows) } + } finally { + if (tables.length > 1) closeAll(tables.toArray) + } + } + + /** Return a new Table with at most numRows rows. */ + private def limitTable(table: Table, numRows: Int): Table = { + val n = if (numRows <= 0) table.getRowCount.toInt + else Math.min(numRows, table.getRowCount).toInt + val cols = new Array[ColumnVector](table.getNumberOfColumns) + try { + for (i <- cols.indices) { + cols(i) = table.getColumn(i).subVector(0, n) + } + new Table(cols: _*) + } finally { + closeAll(cols) + } + } + + /** Get the size of the table in MB. */ + private def getTableSizeMB(table: Table): Double = { + (0 until table.getNumberOfColumns) + .map(i => table.getColumn(i).getDeviceMemorySize) + .sum / (1024.0 * 1024.0) + } + + /** Copy all device columns to host memory. */ + private def copyAllToHost(table: Table): Array[HostColumnVector] = { + Array.tabulate(table.getNumberOfColumns)(i => table.getColumn(i).copyToHost()) + } + + /** Parse CLI arguments. */ + private def parseArgs(args: Array[String]): Map[String, String] = { + var map = Map.empty[String, String] + var i = 0 + while (i < args.length) { + args(i) match { + case "--mode" => map += ("mode" -> args(i + 1)); i += 2 + case "--data-path" => map += ("data-path" -> args(i + 1)); i += 2 + case "--warmup" => map += ("warmup" -> args(i + 1)); i += 2 + case "--measured" => map += ("measured" -> args(i + 1)); i += 2 + case "--rows" => map += ("rows" -> args(i + 1)); i += 2 + case "--pool-fraction" => map += ("pool-fraction" -> args(i + 1)); i += 2 + case "--profile" => map += ("profile" -> "true"); i += 1 + case other => + throw new IllegalArgumentException(s"Unknown argument: $other") + } + } + map + } +} diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/SparkBenchRunner.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/SparkBenchRunner.scala new file mode 100644 index 00000000000..3eafe49a201 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/SparkBenchRunner.scala @@ -0,0 +1,176 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf.bench + +import java.io.{File, PrintWriter, StringWriter} +import com.fasterxml.jackson.core.util.{DefaultIndenter, DefaultPrettyPrinter} +import com.fasterxml.jackson.databind.{ObjectMapper, SerializationFeature} +import com.udf.SparkUtils +import org.apache.spark.sql.{DataFrame, SparkSession} + +/** + * UDF benchmark runner. Measures the end-to-end runtime of: + * Read Parquet -> Execute (CPU or GPU) -> Write no-op sink + * + * Produces a JSON file with the benchmark results. + * On error, also produces separate error log file. + * + * Usage: + * mvn exec:java -Dexec.mainClass=com.udf.bench.SparkBenchRunner \ + * -Dexec.args="--mode cpu --data-path data/bench_data_10M_rows.parquet ..." + */ +object SparkBenchRunner { + + private val DefaultSparkLogLevel = "ERROR" + + def main(args: Array[String]): Unit = { + val (parsed, sparkConfs) = parseArgs(args) + + val mode = parsed.getOrElse("mode", + throw new IllegalArgumentException("--mode is required (cpu or gpu)")) + val dataPath = parsed.getOrElse("data-path", + throw new IllegalArgumentException("--data-path is required")) + val resultPath = parsed.getOrElse("result-path", + throw new IllegalArgumentException("--result-path is required")) + val sparkLogLevel = parsed.getOrElse("spark-log-level", DefaultSparkLogLevel) + + // Resolve execution mode + val executeFn: (SparkSession, DataFrame) => DataFrame = mode match { + case "cpu" => BenchUtils.executeCpu + case "gpu" => BenchUtils.executeGpu + case other => + throw new IllegalArgumentException( + s"Unknown mode: '$other'. Must be 'cpu' or 'gpu'.") + } + + // Build Spark session + val builder = SparkSession.builder() + SparkUtils.applySparkConfs(builder, sparkConfs) + val spark = builder.getOrCreate() + spark.sparkContext.setLogLevel(sparkLogLevel) + + try { + // --- START JOB --- + val startTime = System.nanoTime() + val df = spark.read.parquet(dataPath) + val resultDf = executeFn(spark, df) + resultDf.write.format("noop").mode("overwrite").save() + val elapsed = (System.nanoTime() - startTime) / 1e9 + // --- END JOB --- + + System.err.println(s"E2E Runtime (s): ${f"$elapsed%.2f"}") + + writeReport( + path = resultPath, + mode = mode, + dataPath = dataPath, + elapsed = elapsed, + status = "success", + cliArgs = args) + + } catch { + case e: Exception => + System.err.println(s"Benchmark run failed: ${e.getClass.getSimpleName}") + e.printStackTrace(System.err) + + // Error stack trace is written to a separate error log file. + val errorLogPath = resultPath.replace("_result.json", "_error.log") + writeErrorLog(errorLogPath, e) + + writeReport( + path = resultPath, + mode = mode, + dataPath = dataPath, + elapsed = -1, + status = "error", + cliArgs = args, + errorMessage = Option(e.getMessage), + errorLogFile = Some(errorLogPath)) + + sys.exit(1) + } finally { + spark.stop() + } + + sys.exit(0) + } + + /** Write a JSON benchmark report containing the result and args. */ + private def writeReport( + path: String, + mode: String, + dataPath: String, + elapsed: Double, + status: String, + cliArgs: Array[String], + errorMessage: Option[String] = None, + errorLogFile: Option[String] = None + ): Unit = { + val resultDir = new File(path).getParentFile + if (resultDir != null) resultDir.mkdirs() + + try { + import java.util.{LinkedHashMap => JLinkedHashMap, Arrays => JArrays} + val report = new JLinkedHashMap[String, AnyRef]() + report.put("mode", mode) + report.put("data_path", dataPath) + report.put("status", status) + report.put("e2e_runtime", java.lang.Double.valueOf(elapsed)) + report.put("cli_args", JArrays.asList(cliArgs: _*)) + errorMessage.foreach { msg => + val error = new JLinkedHashMap[String, String]() + error.put("error_message", msg) + errorLogFile.foreach(f => error.put("error_log_file", f)) + report.put("error", error) + } + + val mapper = new ObjectMapper() + mapper.enable(SerializationFeature.INDENT_OUTPUT) + val printer = new DefaultPrettyPrinter() + printer.indentArraysWith(DefaultIndenter.SYSTEM_LINEFEED_INSTANCE) + mapper.writer(printer).writeValue(new File(path), report) + System.err.println(s"Report written to: $path") + } catch { + case e: Exception => + System.err.println(s"Failed to write report: ${e.getMessage}") + } + } + + /** Write an exception to an error log file. */ + private def writeErrorLog(path: String, e: Exception): Unit = { + val logDir = new File(path).getParentFile + if (logDir != null) logDir.mkdirs() + + val pw = new PrintWriter(path) + try { + val sw = new StringWriter() + e.printStackTrace(new java.io.PrintWriter(sw)) + pw.print(sw.toString) + } finally { + pw.close() + } + System.err.println(s"Error details written to: $path") + } + + /** Parse CLI arguments. */ + private def parseArgs(args: Array[String]): (Map[String, String], Seq[String]) = { + var map = Map.empty[String, String] + var sparkConfs = Seq.empty[String] + var i = 0 + while (i < args.length) { + args(i) match { + case "--mode" => map += ("mode" -> args(i + 1)); i += 2 + case "--data-path" => map += ("data-path" -> args(i + 1)); i += 2 + case "--result-path" => map += ("result-path" -> args(i + 1)); i += 2 + case "--spark-log-level" => map += ("spark-log-level" -> args(i + 1)); i += 2 + case "--spark-conf" => sparkConfs :+= args(i + 1); i += 2 + case other => + throw new IllegalArgumentException(s"Unknown argument: $other") + } + } + (map, sparkConfs) + } +} diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala new file mode 100644 index 00000000000..ddc481c28b4 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf + +import org.apache.spark.sql.SparkSession +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.BeforeAndAfterAll + +class CudfComparisonTest extends AnyFunSuite with BeforeAndAfterAll { + + var spark: SparkSession = _ + + override def beforeAll(): Unit = { + spark = SparkSession.builder() + .appName("UDF vs. RapidsUDF Comparison Test") + .master("local[*]") + .config("spark.plugins", "com.nvidia.spark.SQLPlugin") + .config("spark.rapids.memory.gpu.pool", "NONE") + .config("spark.rapids.sql.explain", "NONE") + .getOrCreate() + } + + override def afterAll(): Unit = { + if (spark != null) spark.stop() + } + + /** TODO: Register the RapidsUDF with Spark. */ + def registerRapidsUDF(spark: SparkSession, udfName: String): Unit = ??? + + test("UDF vs RapidsUDF") { + val testDF = UnitTest.createTestData(spark).repartition(1) + + // Run CPU UDF + UnitTest.registerUDF(spark, "placeholder_udf_name") + val cpuResultDF = UnitTest.executeUDF(spark, "placeholder_udf_name", testDF) + UnitTest.verifyUDFResults(cpuResultDF, testDF) + + // Run RapidsUDF + registerRapidsUDF(spark, "placeholder_rapids_udf_name") + val gpuResultDF = UnitTest.executeUDF(spark, "placeholder_rapids_udf_name", testDF) + UnitTest.verifyUDFResults(gpuResultDF, testDF) + + // Compare + TestUtils.assertDataFrameEquals(actual = gpuResultDF, expected = cpuResultDF) + } + + /** + * TODO: If UnitTest adds extra tests beyond the main result checks, add + * corresponding comparison tests here. Each case should run the same input + * through the CPU UDF and the RapidsUDF, apply equivalent assertions to both + * outputs, and compare the RapidsUDF output against the CPU output. + */ +} diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala new file mode 100644 index 00000000000..bb83e405665 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala @@ -0,0 +1,59 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf + +import org.apache.spark.sql.SparkSession +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.BeforeAndAfterAll + +class SqlComparisonTest extends AnyFunSuite with BeforeAndAfterAll { + + var spark: SparkSession = _ + + override def beforeAll(): Unit = { + spark = SparkSession.builder() + .appName("UDF vs. SQL Comparison Test") + .master("local[*]") + .config("spark.plugins", "com.nvidia.spark.SQLPlugin") + .config("spark.rapids.skipGpuArchitectureCheck", "true") + .config("spark.rapids.sql.mode", "explainOnly") + .config("spark.sql.adaptive.enabled", "false") + .getOrCreate() + } + + override def afterAll(): Unit = { + if (spark != null) spark.stop() + } + + test("UDF vs SQL expression") { + val testDF = UnitTest.createTestData(spark).repartition(1) + + // Run CPU UDF + UnitTest.registerUDF(spark, "placeholder_udf_name") + val udfResultDF = UnitTest.executeUDF(spark, "placeholder_udf_name", testDF) + UnitTest.verifyUDFResults(udfResultDF, testDF) + + // Read and execute SQL expression + testDF.createOrReplaceTempView("test_table") + val sqlSource = scala.io.Source.fromFile("src/main/resources/placeholder_udf_name.sql") + val sqlContent = try sqlSource.mkString finally sqlSource.close() + val sqlResultDF = spark.sql(sqlContent) + UnitTest.verifyUDFResults(sqlResultDF, testDF) + + // Compare results + TestUtils.assertDataFrameEquals(actual = sqlResultDF, expected = udfResultDF) + + // Verify GPU compatibility + SparkUtils.assertPlanRunsOnGpu(sqlResultDF) + } + + /** + * TODO: If UnitTest adds extra tests beyond the main result checks, add + * corresponding comparison tests here. Each case should run the same input + * through the CPU UDF and the SQL expression, apply equivalent assertions to + * both outputs, and compare the SQL output against the CPU output. + */ +} diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/TestUtils.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/TestUtils.scala new file mode 100644 index 00000000000..a6d87d5e1b1 --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/TestUtils.scala @@ -0,0 +1,47 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf + +import org.apache.spark.sql.DataFrame + +/** + * Shared test utilities. + */ +object TestUtils { + + /** Compare two DataFrames row-by-row, reporting per-column mismatches. */ + def assertDataFrameEquals( + actual: DataFrame, + expected: DataFrame + ): Unit = { + assert(actual.schema == expected.schema, + s"Schema mismatch:\n actual: ${actual.schema}\n expected: ${expected.schema}") + + val actualRows = actual.collect().sortBy(_.toString) + val expectedRows = expected.collect().sortBy(_.toString) + + assert(actualRows.length == expectedRows.length, + s"Row count mismatch: actual=${actualRows.length}, expected=${expectedRows.length}") + + val mismatches = scala.collection.mutable.ArrayBuffer.empty[String] + for (i <- actualRows.indices) { + val aRow = actualRows(i) + val eRow = expectedRows(i) + for (field <- actual.schema.fieldNames) { + val aVal = Option(aRow.getAs[Any](field)) + val eVal = Option(eRow.getAs[Any](field)) + if (aVal != eVal) { + mismatches += s" [row $i] $field: actual=$aVal, expected=$eVal" + } + } + } + + if (mismatches.nonEmpty) { + throw new AssertionError( + s"\nFound ${mismatches.length} column-level mismatches:\n${mismatches.mkString("\n")}\n") + } + } +} diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala new file mode 100644 index 00000000000..eceffabca6e --- /dev/null +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package com.udf + +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.types._ +import org.scalatest.Assertions +import org.scalatest.funsuite.AnyFunSuite +import org.scalatest.BeforeAndAfterAll + +object UnitTest extends Assertions { + /** + * TODO: Create a test DataFrame with diverse test cases including edge cases. + * + * Example: + * {{{ + * val schema = StructType(Seq( + * StructField("id", IntegerType, nullable = false), + * StructField("credit_score", IntegerType, nullable = true) + * )) + * val testData = Seq( + * Row(1, 800), + * Row(2, 550), + * Row(3, null) + * ) + * spark.createDataFrame(spark.sparkContext.parallelize(testData), schema) + * }}} + */ + def createTestData(spark: SparkSession): DataFrame = ??? + + /** + * TODO: Register the UDF with Spark. + * + * Example: + * {{{ + * spark.udf.register(udfName, new CalculateRiskUDF()) + * }}} + */ + def registerUDF(spark: SparkSession, udfName: String): Unit = ??? + + /** + * TODO: Execute the UDF on the test DataFrame and return the result. + * + * Example: + * {{{ + * testDF.createOrReplaceTempView("test_table") + * spark.sql(s"SELECT *, $udfName(credit_score) AS risk_level FROM test_table") + * }}} + */ + def executeUDF(spark: SparkSession, udfName: String, testDF: DataFrame): DataFrame = ??? + + /** + * TODO: Verify UDF results using assert statements. + * + * Example: + * {{{ + * val results = resultDF.collect().sortBy(_.getAs[Int]("id")) + * assert(results(0).getAs[String]("risk_level") === "LOW") + * assert(results(1).getAs[String]("risk_level") === "MEDIUM") + * assert(results(2).getAs[String]("risk_level") === "UNKNOWN") + * }}} + */ + def verifyUDFResults(resultDF: DataFrame, testDF: DataFrame): Unit = ??? +} + +class UnitTest extends AnyFunSuite with BeforeAndAfterAll { + + var spark: SparkSession = _ + + override def beforeAll(): Unit = { + spark = SparkSession.builder() + .appName("UDF Unit Test") + .master("local[*]") + .config("spark.plugins", "com.nvidia.spark.SQLPlugin") + .config("spark.rapids.skipGpuArchitectureCheck", "true") + .config("spark.rapids.sql.mode", "explainOnly") + .config("spark.sql.adaptive.enabled", "false") + .getOrCreate() + } + + override def afterAll(): Unit = { + if (spark != null) spark.stop() + } + + test("UDF produces correct results") { + val testDF = UnitTest.createTestData(spark).repartition(1) + + UnitTest.registerUDF(spark, "placeholder_udf_name") + val resultDF = UnitTest.executeUDF(spark, "placeholder_udf_name", testDF) + + UnitTest.verifyUDFResults(resultDF, testDF) + } +} diff --git a/skills/udf-judge-conversion/SKILL.md b/skills/udf-judge-conversion/SKILL.md new file mode 100644 index 00000000000..91784aeb9ad --- /dev/null +++ b/skills/udf-judge-conversion/SKILL.md @@ -0,0 +1,92 @@ +--- +name: udf-judge-conversion +description: Reviews generated UDF tests and GPU/SQL implementations for robustness, anti-cheating, and GPU execution integrity. Use when the user requests a judge/review-agent pass, or when manually reviewing a completed conversion. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# Judge UDF Conversion + +## Purpose + +Review a completed UDF conversion and its tests as a skeptical QA/code-review subagent. +Your job is to review whether the GPU/SQL implementation is a properly validated functional replacement for the CPU UDF. + +## Inputs + +Review the files that exist in the generated project: +- CPU UDF source under `src/main//com/udf/` +- `src/test//com/udf/UnitTest.` +- `src/test//com/udf/CudfComparisonTest.` or `SqlComparisonTest.` +- GPU/SQL implementation files +- coverage reports, test output, or comments documenting accepted discrepancies if present + +## Workflow + +- [ ] Step 1: Read the unit test and comparison test. +- [ ] Step 2: Judge whether the tests are strong enough to specify CPU behavior. +- [ ] Step 3: Judge whether the implementation cheats or silently falls back to CPU logic. +- [ ] Step 4: Report actionable findings. + +## Unit Test Checks + +The unit test should be a strong specification of the CPU UDF behavior over its documented input domain. + +Check that: +- Test data covers applicable edge cases such as nulls, empty values, malformed inputs, boundaries, duplicates, mixed valid/invalid rows, nested empties/nulls, unicode, timestamps/timezones, and decimal scale. +- Assertions verify schema, row count, deterministic ordering, output values, null propagation, and exception/default behavior where applicable. +- The test exercises visible CPU UDF branches. Coverage reports should support this when available. +- Assertions reflect the CPU UDF's actual behavior and do not merely assert weak properties such as non-null output. +- Extra unit tests outside the shared `verifyUDFResults` path are mirrored in the comparison test and run against both CPU and GPU/SQL paths. + +## Comparison Test Checks + +The comparison test should provide strong evidence that the converted implementation preserves the CPU UDF behavior. + +Check that: +- The CPU path and GPU/SQL path run on the same input data. +- The CPU result and GPU/SQL result are compared directly. +- The comparison test actually runs on the GPU with the Spark RAPIDS plugin enabled. +- The converted path is also validated with the same result assertions used for the CPU path. +- Additional unit test cases are converted into CPU-vs-GPU/SQL comparison cases, not left as CPU-only tests. +- Commented-out tests or assertions include a clear explanation and a user-facing note. Documented deviations are acceptable only if the reason is explicit. + - Note: you should not accept a documented deviation that removes coverage of the UDF's core logic. + +## Implementation Checks + +Fail the review if the implementation is tailored to the tests instead of implementing the UDF generally. Look for: +- Hardcoded test inputs, IDs, row counts, or expected outputs. +- Conditional branches that only handle exact values from the tests. +- Literal lookup tables derived from test data. + +Fail the review if the implementation silently performs logic row-by-row on the CPU. Look for: +- `copyToHost()`, `cudaMemcpyDeviceToHost`, or row-by-row scalar copies such as `getJavaString` to copy input data to the CPU. + - Note: small CPU objects for metadata or temporary storage are acceptable. + +If a GPU API's behavior is unclear, inspect the implementation or docs for the SQL/cuDF/libcudf/thrust APIs invoked by the UDF. Clone the matching source if needed to understand subtle null, type, boundary, or semantic behavior under the hood. + +## Output + +Start with a clear verdict: +- `PASS`: no blocking issues found +- `FAIL`: one or more blocking issues found + +### Verdict Examples + +`PASS`: The unit test covers normal inputs plus meaningful edge cases, coverage gaps are explained, the comparison test runs the same cases through CPU and GPU/SQL paths, the implementation is general, and there are no hidden CPU fallbacks or test-derived literals. + +`PASS with non-blocking risks`: One malformed-input assertion is commented out because the CPU throws a row-level exception while the GPU path returns null for that row, and comments explain the attempted fixes and why the behavior is outside the supported GPU contract. The normal input domain and core UDF logic are still fully tested. + +`FAIL`: A test for the primary transformation is commented out, most assertions only check row counts or non-null output, or the comparison test leaves extra CPU-only unit tests unmatched. These failures weaken confidence even if comments are present. + +`FAIL`: The implementation contains test-specific literals, dispatches on exact test rows, calls the CPU UDF from the GPU/SQL path, or copies column data to the host to perform normal business logic. + +For failures, concisely list specific findings with: +- file/path +- issue +- why it matters +- suggested fix + +Also include any non-blocking risks or test gaps separately. diff --git a/skills/udf-optimize-cudf/SKILL.md b/skills/udf-optimize-cudf/SKILL.md new file mode 100644 index 00000000000..63c27b2ed3c --- /dev/null +++ b/skills/udf-optimize-cudf/SKILL.md @@ -0,0 +1,149 @@ +--- +name: udf-optimize-cudf +description: Iteratively optimizes a cuDF RapidsUDF implementation for GPU performance. Use after testing and benchmarking with udf-benchmark. Runs a loop of profiling, optimizing, testing, and benchmarking until performance converges or the iteration budget is exhausted. +license: CC-BY-4.0 AND Apache-2.0 +metadata: + spdx-file-copyright-text: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +model: inherit +--- + +# Optimize cuDF RapidsUDF + +## Workflow + +- [ ] Step 0: Create backup and establish baseline +- [ ] Steps 1-4: Iterative optimization loop (repeat up to N iterations) + - [ ] Step 1: Profile with nsys + - [ ] Step 2: Implement one targeted change + - [ ] Step 3: Run unit tests (fail → discard, retry) + - [ ] Step 4: Run microbenchmarks (no improvement → discard, retry) +- [ ] Final Step 1: Run judge subagent if requested +- [ ] Final Step 2: Review optimized implementation and report results + +## Prerequisites + +- Project directory with passing unit tests and cuDF comparison test +- MicroBenchRunner implemented and working (from the **udf-benchmark** skill) +- Benchmark data generated (reuse from the benchmark step) + +Derive `` and `` from the UDF class name. + +> **Note:** Commands require access to `/tmp` (Spark temp storage) and `/dev` (GPU device). If commands fail due to sandbox restrictions, re-run them unsandboxed. + +## Step 0: Create Backup and Establish Baseline + +1. Create a backup of the current RapidsUDF implementation: +```bash +cp src/main//com/udf/RapidsUDF. \ + src/main//com/udf/RapidsUDF..bak +``` + +2. If no `.orig.bak` exists yet, save the original unoptimized implementation: +```bash +cp src/main//com/udf/RapidsUDF. \ + src/main//com/udf/RapidsUDF..orig.bak +``` +This file is never overwritten; it preserves the pre-optimization baseline. + +3. If no baseline microbenchmark results exist, run the baseline now: +```bash +./run_micro_benchmark.sh --mode all --data-path data/bench_data__rows.parquet --rows +``` + +Record the baseline GPU time and speedup. This is the number to beat. + +## Iterative Optimization Loop + +Repeat the following steps up to **N iterations** (default: 10). Also stop early if no improvement is found after **3 consecutive failed attempts**. + +Maintain an **optimization log** throughout the loop: for each iteration, record what change was attempted and whether it improved, regressed, or had no effect. This prevents repeating failed approaches and feeds the final report. + +### Step 1: Profile with nsys + +Profile the current implementation to identify bottlenecks: +```bash +./run_micro_benchmark.sh --mode gpu --data-path data/bench_data__rows.parquet --rows --profile +``` + +Summarize libcudf kernel stats: +```bash +nsys stats --report nvtx_sum --format csv -o rapidsudf results/.nsys-rep +``` + +Consult **references/OPTIMIZATION_PATTERNS.md** for interpreting profiler output and identifying optimization opportunities. + +> **Tip:** Profiling frequently is strongly recommended. Without profiler data, optimization changes are guesses. Try using other `nsys stats` commands as needed. + +### Step 2: Implement One Targeted Change + +Based on profiling insights (or optimization patterns from the reference), make **one targeted change** to the RapidsUDF implementation. Isolating changes one at a time makes it possible to attribute performance impact. + +### Step 3: Run Unit Tests + +```bash +# Java +mvn test -Dtest=CudfComparisonTest + +# Scala +mvn test -Dsuites=com.udf.CudfComparisonTest +``` + +- **Tests pass** → proceed to Step 4 +- **Tests fail** → analyze the failure. If it is an ordinary implementation bug, fix it and rerun the test. If the targeted optimization introduces a CPU/GPU semantic mismatch that cannot be resolved, discard changes by restoring from backup: + ```bash + cp src/main//com/udf/RapidsUDF..bak \ + src/main//com/udf/RapidsUDF. + ``` + Log the failure reason, then return to Step 1. + +Sometimes, matching a certain edge case is impossible without a major performance tradeoff. If so, document the attempted fix, the benchmark evidence, and the exact behavior difference, then ask the user whether the performance-vs-correctness tradeoff is acceptable. +Do not comment out tests or accept a correctness difference during optimization unless the user explicitly approves that tradeoff. + +### Step 4: Run Microbenchmarks + +```bash +./run_micro_benchmark.sh --mode all --data-path data/bench_data__rows.parquet --rows +``` + +Compare the GPU time against the current best (from the last checkpoint). + +- **Performance improved** → create a new checkpoint: + ```bash + cp src/main//com/udf/RapidsUDF. \ + src/main//com/udf/RapidsUDF..bak + ``` + Record the new best GPU time. Reset the consecutive-failure counter. Return to Step 1. + +- **Performance did NOT improve** → discard changes: + ```bash + cp src/main//com/udf/RapidsUDF..bak \ + src/main//com/udf/RapidsUDF. + ``` + Increment the consecutive-failure counter. Return to Step 1. + +## Final Step 1: Run Judge Subagent If Requested + +If the user explicitly asked for the judge, a judge subagent, or a review agent, treat that as an explicit request for delegation: you **MUST** launch a separate subagent with `model: inherit` and instruct it to use the **udf-judge-conversion** skill. Ask it to review the `UnitTest`, `CudfComparisonTest`, optimized RapidsUDF implementation, and optimization log as a cuDF conversion. + +If the user did not request a judge/review agent, mark this step as skipped and continue to Final Step 2. If a required judge subagent is blocked by tool policy, stop and tell the user that explicit permission/instruction is needed. + +If you run the judge, include the judge verdict in the final report. If there are any blocking issues, fix them or report the last known-good checkpoint. + +## Final Step 2: Review Optimized Implementation and Report Results + +After completing all iterations (or early-stopping), review your own work to ensure the optimization did not weaken correctness, introduce hardcoded test behavior, hide CPU fallback logic, or comment out core test coverage. + +After completing all iterations (or early-stopping), report: +1. **Baseline**: starting GPU time and speedup +2. **Final**: best GPU time and speedup (from the last checkpoint) +3. **Successful optimizations**: what changes improved performance and by how much +4. **Failed optimizations**: what was attempted but did not help +5. **Review result**: self-review summary, or judge PASS/failures if the judge was requested + +## Output + +Upon successful completion: +- Optimized RapidsUDF: `src/main//com/udf/RapidsUDF.` +- Backup of best version: `src/main//com/udf/RapidsUDF..bak` +- Original unoptimized version: `src/main//com/udf/RapidsUDF..orig.bak` +- Benchmark results: `results/` diff --git a/skills/udf-optimize-cudf/references/OPTIMIZATION_PATTERNS.md b/skills/udf-optimize-cudf/references/OPTIMIZATION_PATTERNS.md new file mode 100644 index 00000000000..2e51fac8012 --- /dev/null +++ b/skills/udf-optimize-cudf/references/OPTIMIZATION_PATTERNS.md @@ -0,0 +1,22 @@ + + +# cuDF Optimization Patterns + +## Guidelines + +- **Rule of thumb:** fewer cuDF API calls typically results in better performance. Look for ways to collapse multiple operations into fewer calls. +- Explore the cuDF repo `java/src//java/ai/rapids/cudf` to find alternative cuDF methods. + +## Profiling Signals + +| Signal | Where to look | What it means | Action | +|---|---|---|---| +| High invocation count | `nvtx_sum` Instances column | Loops or too many small kernels | Batch into fewer calls | +| Low GPU utilization | `kernel_time / wall_time` | Launch/memory overhead dominates | Reduce total API calls | +| Many `make_*_column` calls | `nvtx_sum` | Excessive intermediate columns | Shorten transformation chains | +| Expensive kernel | `nvtx_sum` | Look for cheaper API (e.g., regex → stringReplace, stringSplitRecord → stringSplit) | Swap to cheaper cuDF API | +| GPU slower than CPU at large scale | Speedup results | Algorithm has serial dependencies that don't parallelize well | Rethink overall algorithm to maximize columnar parallelism and reduce divergence | +| Many gather/scatter or struct unpacking ops | `nvtx_sum` | Non-contiguous memory access patterns | Use APIs that leverage contiguous access (e.g., operate on cuDF child columns directly) | From dc9b25b640a778d84af7ae7165dafb5af59c9f13 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 19:36:08 -0700 Subject: [PATCH 02/16] cleanups. --- skills/docs/dev/VERSIONS.md | 8 ++++++++ skills/udf-convert-to-cuda/templates/cuda/Dockerfile | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/skills/docs/dev/VERSIONS.md b/skills/docs/dev/VERSIONS.md index e56fffdb236..8f811e98369 100644 --- a/skills/docs/dev/VERSIONS.md +++ b/skills/docs/dev/VERSIONS.md @@ -29,6 +29,14 @@ Update these properties together: - `` - `` +### Native CUDA build image + +File: `skills/convert-to-cuda/templates/cuda/Dockerfile` + +Update this default value: + +- `CUDA_VERSION` must match the CUDA toolkit version spark-rapids is built against (the same version the native build uses on the host). + ### Native CUDA dependency extraction File: `skills/udf-convert-to-cuda/templates/cuda/native/scripts/extract-cudf-libs.sh` diff --git a/skills/udf-convert-to-cuda/templates/cuda/Dockerfile b/skills/udf-convert-to-cuda/templates/cuda/Dockerfile index 320dc2f1423..f75a5f19376 100644 --- a/skills/udf-convert-to-cuda/templates/cuda/Dockerfile +++ b/skills/udf-convert-to-cuda/templates/cuda/Dockerfile @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # Reproducible build image for native CUDA RapidsUDF code. -ARG CUDA_VERSION=12.8.0 +ARG CUDA_VERSION=12.9.1 ARG LINUX_VERSION=rockylinux8 FROM nvidia/cuda:${CUDA_VERSION}-devel-${LINUX_VERSION} From cce30115e59b965c453dae51dc29a253496dbcdf Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 20:04:05 -0700 Subject: [PATCH 03/16] update installation --- skills/README.md | 36 +++++++++++------------------------- 1 file changed, 11 insertions(+), 25 deletions(-) diff --git a/skills/README.md b/skills/README.md index 8e5694a5b3d..71f01c0c794 100644 --- a/skills/README.md +++ b/skills/README.md @@ -10,16 +10,24 @@ Aether Agent is a set of skills to convert Apache Spark User-Defined Functions (
Table of Contents +- [Installation](#installation) - [Supported Formats](#supported-formats) - [Prerequisites](#prerequisites) - [Selecting an LLM](#selecting-an-llm) - [Quick Start](#quick-start) - - [Installing Skills](#installing-skills) - [Using Skills](#using-skills) - - [Quick Start](#quick-start-1) + - [Try the Workflow](#try-the-workflow)
+## Installation + +Install via the [skills CLI](https://github.com/vercel-labs/skills). Installing all skills is recommended, as they are designed to work together. + +```bash +npx skills add NVIDIA/spark-rapids --skill '*' [--agent ] +``` + ## Supported Formats | UDF Type | cuDF RapidsUDF | CUDA RapidsUDF | Spark SQL | @@ -47,28 +55,6 @@ For best results, we recommend the latest reasoning models from OpenAI, Anthropi Skills require any IDE or LLM that supports the [agent skills spec](https://skill.md/) (e.g., Cursor, Codex, Claude Code). -### Installing Skills - -Copy the skills from this repo into your project: - -```bash -# Claude Code -mkdir -p /path/to/your/project/.claude/skills/ -cp -r skills/* /path/to/your/project/.claude/skills/ - -# Codex -mkdir -p /path/to/your/project/.agents/skills/ -cp -r skills/* /path/to/your/project/.agents/skills/ - -# Cursor -mkdir -p /path/to/your/project/.cursor/skills/ -cp -r skills/* /path/to/your/project/.cursor/skills/ - -# Kiro -mkdir -p /path/to/your/project/.kiro/skills/ -cp -r skills/* /path/to/your/project/.kiro/skills/ -``` - ### Using Skills Skills follow a multi-step workflow: @@ -97,7 +83,7 @@ You can invoke multiple steps in a single prompt: ❯ Generate a unit test for @FormatPhoneUDF.java, then convert it to cuDF, native CUDA, or SQL and benchmark ``` -### Quick Start +### Try the Workflow Once you've installed the skills, try the workflow with one of the provided example UDFs: - Java: [FormatPhoneUDF.java](examples/FormatPhoneUDF.java) From 474c97192754fadb8ca997ffa620347981b5f5d9 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 20:06:43 -0700 Subject: [PATCH 04/16] signoff Signed-off-by: Rishi Chandra From d8ae9dc2c3e63ce85c835b48fe34783c7531b827 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 20:15:15 -0700 Subject: [PATCH 05/16] ignore pom files under skills/ in scala build --- build/make-scala-version-build-files.sh | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/build/make-scala-version-build-files.sh b/build/make-scala-version-build-files.sh index 21bf4471147..88d482ec118 100755 --- a/build/make-scala-version-build-files.sh +++ b/build/make-scala-version-build-files.sh @@ -78,6 +78,11 @@ for f in $(git ls-files '**pom.xml'); do echo "Skipping $f" continue fi + # Skills package their own pom.xml templates. Ignore those. + if [[ $f == skills/* ]]; then + echo "Skipping $f" + continue + fi echo $f tof="$TO_DIR/$f" mkdir -p $(dirname $tof) From 0f1998e5de478987d362b5cd9af7be3558047d79 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 20:16:01 -0700 Subject: [PATCH 06/16] license header --- build/make-scala-version-build-files.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build/make-scala-version-build-files.sh b/build/make-scala-version-build-files.sh index 88d482ec118..295ff44de1e 100755 --- a/build/make-scala-version-build-files.sh +++ b/build/make-scala-version-build-files.sh @@ -1,6 +1,6 @@ #!/usr/bin/env bash # -# Copyright (c) 2023-2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2023-2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. From f48b5313c05747df58f6cba717833401e30e240b Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 20:33:09 -0700 Subject: [PATCH 07/16] ignore skills in rat, fix readme link --- pom.xml | 2 ++ skills/README.md | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 450211bcc4a..a004094caa2 100644 --- a/pom.xml +++ b/pom.xml @@ -1654,6 +1654,8 @@ This will force full Scala code rebuild in downstream modules. **/target/**/* **/cufile.log **/cudf_log.txt + + skills/** thirdparty/parquet-testing/** diff --git a/skills/README.md b/skills/README.md index 71f01c0c794..ca60dcfdad8 100644 --- a/skills/README.md +++ b/skills/README.md @@ -53,7 +53,7 @@ For best results, we recommend the latest reasoning models from OpenAI, Anthropi ## Quick Start -Skills require any IDE or LLM that supports the [agent skills spec](https://skill.md/) (e.g., Cursor, Codex, Claude Code). +Skills require any IDE or LLM that supports the [agent skills spec](https://agentskills.io) (e.g., Cursor, Codex, Claude Code). ### Using Skills From 446ec504094e4f0a14f6e5df21ca10f7ed7dde12 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 20:44:14 -0700 Subject: [PATCH 08/16] regenerate scala 2.13 pom --- scala2.13/pom.xml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scala2.13/pom.xml b/scala2.13/pom.xml index 6b9a9aa8d68..a1b0b122e51 100644 --- a/scala2.13/pom.xml +++ b/scala2.13/pom.xml @@ -1654,6 +1654,8 @@ This will force full Scala code rebuild in downstream modules. **/target/**/* **/cufile.log **/cudf_log.txt + + skills/** thirdparty/parquet-testing/** From 6a709e583d60b55c2c1a0792af3f4ca3ff9aadc0 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 21:28:58 -0700 Subject: [PATCH 09/16] skip skills in scala style --- pom.xml | 1 + scala2.13/pom.xml | 1 + 2 files changed, 2 insertions(+) diff --git a/pom.xml b/pom.xml index a004094caa2..cc51b5215a0 100644 --- a/pom.xml +++ b/pom.xml @@ -1706,6 +1706,7 @@ This will force full Scala code rebuild in downstream modules. + diff --git a/scala2.13/pom.xml b/scala2.13/pom.xml index a1b0b122e51..69882d86cc1 100644 --- a/scala2.13/pom.xml +++ b/scala2.13/pom.xml @@ -1706,6 +1706,7 @@ This will force full Scala code rebuild in downstream modules. + From 53e1feff96777ec6afa102a4d6250213ee310399 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Wed, 10 Jun 2026 22:12:15 -0700 Subject: [PATCH 10/16] address comments --- skills/tests/test_export/scala_fixtures.py | 3 ++- .../templates/scala/src/main/scala/com/udf/Arm.scala | 6 +++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/skills/tests/test_export/scala_fixtures.py b/skills/tests/test_export/scala_fixtures.py index 8bf96649166..e4f1a472e49 100644 --- a/skills/tests/test_export/scala_fixtures.py +++ b/skills/tests/test_export/scala_fixtures.py @@ -148,7 +148,8 @@ def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { BENCH_GPU_SQL = """\ def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { df.createOrReplaceTempView("bench_table") - val sqlContent = scala.io.Source.fromFile("src/main/resources/integer_multiply_by_2.sql").mkString + val sqlSource = scala.io.Source.fromFile("src/main/resources/integer_multiply_by_2.sql") + val sqlContent = try sqlSource.mkString finally sqlSource.close() val benchSql = sqlContent.replace("test_table", "bench_table") spark.sql(benchSql) }""" diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala index e65ed1d6a51..dd0c48d5f7c 100644 --- a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/Arm.scala @@ -57,7 +57,11 @@ object Arm { */ def closeAll[T <: AutoCloseable](resources: Array[T]): Unit = { if (resources != null) { - resources.foreach(r => if (r != null) r.close()) + resources.foreach { r => + if (r != null) { + try { r.close() } catch { case _: Exception => } + } + } } } } From 81f6cf9b29a2a62db98baef3bd4ef30fd4ae8f83 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Fri, 12 Jun 2026 06:39:20 -0700 Subject: [PATCH 11/16] fix name --- skills/docs/dev/VERSIONS.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/skills/docs/dev/VERSIONS.md b/skills/docs/dev/VERSIONS.md index 8f811e98369..83956663ca3 100644 --- a/skills/docs/dev/VERSIONS.md +++ b/skills/docs/dev/VERSIONS.md @@ -31,7 +31,7 @@ Update these properties together: ### Native CUDA build image -File: `skills/convert-to-cuda/templates/cuda/Dockerfile` +File: `skills/udf-convert-to-cuda/templates/cuda/Dockerfile` Update this default value: From d7438d5b2abaa93637d4a7a95e0d8eda5db8f1b6 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Fri, 12 Jun 2026 08:46:45 -0700 Subject: [PATCH 12/16] remove tests from initial publish --- skills/docs/dev/TESTING.md | 33 -- skills/tests/test_export/__init__.py | 6 - skills/tests/test_export/cuda_fixtures.py | 208 -------- skills/tests/test_export/java_fixtures.py | 201 -------- skills/tests/test_export/scala_fixtures.py | 194 -------- skills/tests/test_export/test_jvm.py | 523 --------------------- skills/tests/test_export/utils.py | 91 ---- skills/tests/test_skill_frontmatter.py | 40 -- 8 files changed, 1296 deletions(-) delete mode 100644 skills/docs/dev/TESTING.md delete mode 100644 skills/tests/test_export/__init__.py delete mode 100644 skills/tests/test_export/cuda_fixtures.py delete mode 100644 skills/tests/test_export/java_fixtures.py delete mode 100644 skills/tests/test_export/scala_fixtures.py delete mode 100644 skills/tests/test_export/test_jvm.py delete mode 100644 skills/tests/test_export/utils.py delete mode 100644 skills/tests/test_skill_frontmatter.py diff --git a/skills/docs/dev/TESTING.md b/skills/docs/dev/TESTING.md deleted file mode 100644 index 505a96a0258..00000000000 --- a/skills/docs/dev/TESTING.md +++ /dev/null @@ -1,33 +0,0 @@ -# Testing - -## Setup - -Set up a local dev environment: - -```bash -python -m venv .venv -source .venv/bin/activate -pip install -e ".[dev]" -``` - -## Fast Tests - -Run the fast tests: - -```bash -pytest -m "not slow" -``` - -These are generally lightweight skill validation tests, such as verifying skill frontmatter. - -## Integration Tests - -Run the integration tests: - -```bash -pytest -m slow -s tests/test_export -``` - -These tests deterministically fill in the Java/Scala template projects from `skills/udf-gen-test/templates/` with fixture implementations, then actually compile and run Spark tests and benchmark scripts locally. - -Thus they require JDK, Maven and Maven repository access, and a GPU environment for GPU paths. diff --git a/skills/tests/test_export/__init__.py b/skills/tests/test_export/__init__.py deleted file mode 100644 index 2ae3327da34..00000000000 --- a/skills/tests/test_export/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Tests for the export package. -""" diff --git a/skills/tests/test_export/cuda_fixtures.py b/skills/tests/test_export/cuda_fixtures.py deleted file mode 100644 index 51175478c85..00000000000 --- a/skills/tests/test_export/cuda_fixtures.py +++ /dev/null @@ -1,208 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 - -""" -CUDA source fixtures for JVM export integration tests. -""" - -NATIVE_RAPIDS_UDF_SOURCE = """\ -package com.udf; - -import ai.rapids.cudf.ColumnVector; -import com.nvidia.spark.RapidsUDF; -import org.apache.hadoop.hive.ql.exec.UDF; -import org.apache.spark.sql.api.java.UDF1; - -public class IntegerMultiplyBy2NativeRapidsUDF extends UDF - implements UDF1, RapidsUDF { - public Integer evaluate(Integer value) { - if (value == null) return null; - return value * 2; - } - - @Override - public Integer call(Integer value) { - return evaluate(value); - } - - @Override - public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { - if (args.length != 1) { - throw new IllegalArgumentException("Unexpected argument count: " + args.length); - } - if (numRows != args[0].getRowCount()) { - throw new IllegalArgumentException( - "Expected " + numRows + " rows, received " + args[0].getRowCount()); - } - - NativeUDFLoader.ensureLoaded(); - return new ColumnVector(integerMultiplyBy2(args[0].getNativeView())); - } - - private static native long integerMultiplyBy2(long inputView); -} -""" - -JNI_SOURCE = """\ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -// SPDX-License-Identifier: Apache-2.0 - -#include "integer_multiply_by_2.hpp" - -#include -#include -#include - -#include - -#include -#include - -namespace { - -constexpr char const* RUNTIME_ERROR_CLASS = "java/lang/RuntimeException"; -constexpr char const* ILLEGAL_ARG_CLASS = "java/lang/IllegalArgumentException"; - -void throw_java_exception(JNIEnv* env, char const* class_name, char const* message) -{ - jclass ex_class = env->FindClass(class_name); - if (ex_class != nullptr) { - env->ThrowNew(ex_class, message); - } -} - -} // namespace - -extern "C" { - -JNIEXPORT jlong JNICALL -Java_com_udf_IntegerMultiplyBy2NativeRapidsUDF_integerMultiplyBy2(JNIEnv* env, - jclass, - jlong input_view) -{ - try { - auto input = reinterpret_cast(input_view); - if (input == nullptr) { - throw_java_exception(env, ILLEGAL_ARG_CLASS, "input column view is null"); - return 0; - } - if (input->type().id() != cudf::type_id::INT32) { - throw_java_exception(env, ILLEGAL_ARG_CLASS, "input must be INT32"); - return 0; - } - - std::unique_ptr result = integer_multiply_by_2(*input); - return reinterpret_cast(result.release()); - } catch (std::bad_alloc const& e) { - auto message = std::string("Unable to allocate native memory: ") + e.what(); - throw_java_exception(env, RUNTIME_ERROR_CLASS, message.c_str()); - } catch (std::invalid_argument const& e) { - throw_java_exception(env, ILLEGAL_ARG_CLASS, e.what()); - } catch (std::exception const& e) { - throw_java_exception(env, RUNTIME_ERROR_CLASS, e.what()); - } - return 0; -} - -} -""" - -CUDA_SOURCE = """\ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -// SPDX-License-Identifier: Apache-2.0 - -#include "integer_multiply_by_2.hpp" - -#include -#include -#include -#include -#include -#include - -#include -#include - -#include -#include - -namespace { - -__global__ void multiply_by_2_kernel(int32_t const* input, int32_t* output, cudf::size_type size) -{ - auto const idx = static_cast(blockIdx.x * blockDim.x + threadIdx.x); - if (idx < size) { - output[idx] = input[idx] * 2; - } -} - -} // namespace - -std::unique_ptr integer_multiply_by_2( - cudf::column_view const& input, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - if (input.type().id() != cudf::type_id::INT32) { - throw std::invalid_argument("input must be INT32"); - } - - auto const row_count = input.size(); - auto null_mask = cudf::copy_bitmask(input, stream, mr); - auto result = cudf::make_numeric_column( - input.type(), row_count, std::move(null_mask), input.null_count(), stream, mr); - - if (row_count > 0) { - constexpr int threads_per_block = 256; - int const blocks = (row_count + threads_per_block - 1) / threads_per_block; - multiply_by_2_kernel<<>>( - input.data(), result->mutable_view().data(), row_count); - CUDF_CHECK_CUDA(stream.value()); - } - - return result; -} -""" - -HEADER_SOURCE = """\ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -// SPDX-License-Identifier: Apache-2.0 - -#pragma once - -#include -#include -#include -#include - -#include -#include - -#include - -std::unique_ptr integer_multiply_by_2( - cudf::column_view const& input, - rmm::cuda_stream_view stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); -""" - -CMAKE_SOURCE_FILES = """\ -set(SOURCE_FILES - "src/IntegerMultiplyBy2Jni.cpp" - "src/integer_multiply_by_2.cu" -) -""" - -PLACEHOLDER_FILES = ( - "src/main/java/com/udf/PlaceholderUDFNameNativeRapidsUDF.java", - "native/src/main/cpp/src/PlaceholderUDFNameJni.cpp", - "native/src/main/cpp/src/placeholder_udf_name.cu", - "native/src/main/cpp/src/placeholder_udf_name.hpp", -) - -NATIVE_SOURCE_FILES = { - "src/main/java/com/udf/IntegerMultiplyBy2NativeRapidsUDF.java": NATIVE_RAPIDS_UDF_SOURCE, - "native/src/main/cpp/src/IntegerMultiplyBy2Jni.cpp": JNI_SOURCE, - "native/src/main/cpp/src/integer_multiply_by_2.cu": CUDA_SOURCE, - "native/src/main/cpp/src/integer_multiply_by_2.hpp": HEADER_SOURCE, -} diff --git a/skills/tests/test_export/java_fixtures.py b/skills/tests/test_export/java_fixtures.py deleted file mode 100644 index 2845308800d..00000000000 --- a/skills/tests/test_export/java_fixtures.py +++ /dev/null @@ -1,201 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Java source code fixtures for integration tests. -""" - -from .utils import replace_java_todo_method - -NAME = "java" -REPLACE_TODO_FN = replace_java_todo_method -TEST_SELECTOR_FLAG = "-Dtest" - -# --------------------------------------------------------------------------- -# UDF source code -# --------------------------------------------------------------------------- - -UDF_SOURCE = """\ -package com.udf; - -import org.apache.hadoop.hive.ql.exec.UDF; - -public class IntegerMultiplyBy2UDF extends UDF { - public Integer evaluate(Integer value) { - if (value == null) return null; - return value * 2; - } -} -""" - -RAPIDS_UDF_SOURCE = """\ -package com.udf; - -import ai.rapids.cudf.ColumnVector; -import ai.rapids.cudf.Scalar; -import com.nvidia.spark.RapidsUDF; -import org.apache.hadoop.hive.ql.exec.UDF; - -public class IntegerMultiplyBy2RapidsUDF extends UDF implements RapidsUDF { - public Integer evaluate(Integer value) { - if (value == null) return null; - return value * 2; - } - - @Override - public ColumnVector evaluateColumnar(int numRows, ColumnVector... args) { - try (Scalar two = Scalar.fromInt(2)) { - return args[0].mul(two); - } - } -} -""" - -SQL_SOURCE = """\ -SELECT *, - value * 2 AS result -FROM test_table -""" - -# --------------------------------------------------------------------------- -# Unit test methods -# --------------------------------------------------------------------------- - -UNIT_TEST_METHODS = { - "createTestData": """\ - public static Dataset createTestData(SparkSession spark) { - StructType schema = new StructType(new StructField[]{ - DataTypes.createStructField("id", DataTypes.IntegerType, false), - DataTypes.createStructField("value", DataTypes.IntegerType, true) - }); - List data = Arrays.asList( - RowFactory.create(1, 123), - RowFactory.create(2, 0), - RowFactory.create(3, -5), - RowFactory.create(4, null) - ); - return spark.createDataFrame(data, schema); - }""", - "registerUDF": """\ - public static void registerUDF(SparkSession spark, String udfName) { - spark.sql("CREATE TEMPORARY FUNCTION " + udfName - + " AS 'com.udf.IntegerMultiplyBy2UDF'"); - }""", - "executeUDF": """\ - public static Dataset executeUDF(SparkSession spark, String udfName, Dataset testDF) { - testDF.createOrReplaceTempView("test_table"); - return spark.sql("SELECT *, " + udfName - + "(value) AS result FROM test_table"); - }""", - "verifyUDFResults": """\ - public static void verifyUDFResults(Dataset resultDF, Dataset testDF) { - Row[] results = (Row[]) resultDF.sort("id").collect(); - Assert.assertEquals(246, (int) results[0].getAs("result")); - Assert.assertEquals(0, (int) results[1].getAs("result")); - Assert.assertEquals(-10, (int) results[2].getAs("result")); - Assert.assertTrue(results[3].isNullAt(results[3].fieldIndex("result"))); - }""", -} - -RAPIDS_UDF_REGISTER = """\ - public static void registerRapidsUDF(SparkSession spark, String udfName) { - spark.sql("CREATE TEMPORARY FUNCTION " + udfName - + " AS 'com.udf.IntegerMultiplyBy2RapidsUDF'"); - }""" - -NATIVE_RAPIDS_UDF_REGISTER = """\ - public static void registerRapidsUDF(SparkSession spark, String udfName) { - spark.sql("CREATE TEMPORARY FUNCTION " + udfName - + " AS 'com.udf.IntegerMultiplyBy2NativeRapidsUDF'"); - }""" - -# --------------------------------------------------------------------------- -# BenchUtils methods -# --------------------------------------------------------------------------- - -BENCH_GENERATE = """\ - public static Dataset generateSyntheticData( - SparkSession spark, long numRows, int numPartitions) { - Dataset baseDF = spark.range(0, numRows, 1, numPartitions).toDF("id"); - return baseDF.select( - col("id"), - expr("CAST(rand() * 1000 AS INT)").alias("value") - ); - }""" - -BENCH_CPU = """\ - public static Dataset executeCpu(SparkSession spark, Dataset df) { - df.createOrReplaceTempView("bench_table"); - spark.sql("CREATE TEMPORARY FUNCTION integer_multiply_by_2" - + " AS 'com.udf.IntegerMultiplyBy2UDF'"); - return spark.sql("SELECT *, integer_multiply_by_2(value)" - + " AS result FROM bench_table"); - }""" - -BENCH_GPU_CUDF = """\ - public static Dataset executeGpu(SparkSession spark, Dataset df) { - df.createOrReplaceTempView("bench_table"); - spark.sql("CREATE TEMPORARY FUNCTION integer_multiply_by_2_rapids" - + " AS 'com.udf.IntegerMultiplyBy2RapidsUDF'"); - return spark.sql("SELECT *, integer_multiply_by_2_rapids(value)" - + " AS result FROM bench_table"); - }""" - -BENCH_GPU_CUDA = """\ - public static Dataset executeGpu(SparkSession spark, Dataset df) { - df.createOrReplaceTempView("bench_table"); - spark.sql("CREATE TEMPORARY FUNCTION integer_multiply_by_2_native" - + " AS 'com.udf.IntegerMultiplyBy2NativeRapidsUDF'"); - return spark.sql("SELECT *, integer_multiply_by_2_native(value)" - + " AS result FROM bench_table"); - }""" - -BENCH_GPU_SQL = """\ - public static Dataset executeGpu(SparkSession spark, Dataset df) { - df.createOrReplaceTempView("bench_table"); - try { - String sqlContent = new String( - java.nio.file.Files.readAllBytes( - java.nio.file.Paths.get("src/main/resources/integer_multiply_by_2.sql"))); - String benchSql = sqlContent.replace("test_table", "bench_table"); - return spark.sql(benchSql); - } catch (java.io.IOException e) { - throw new RuntimeException(e); - } - }""" - -# --------------------------------------------------------------------------- -# MicroBenchRunner methods -# --------------------------------------------------------------------------- - -MICRO_PREPARE_CPU = """\ - public static Object[] prepareCpuData(HostColumnVector[] hostColumns, int numRows) { - Integer[] values = new Integer[numRows]; - for (int i = 0; i < numRows; i++) { - values[i] = hostColumns[1].isNull(i) - ? null : hostColumns[1].getInt(i); - } - return new Object[] { values }; - }""" - -MICRO_EXECUTE_CPU = """\ - public static void executeCpu(Object[] data, int numRows) { - Integer[] values = (Integer[]) data[0]; - com.udf.IntegerMultiplyBy2UDF udf = new com.udf.IntegerMultiplyBy2UDF(); - for (int i = 0; i < numRows; i++) { - udf.evaluate(values[i]); - } - }""" - -MICRO_EXECUTE_GPU = """\ - public static ColumnVector executeGpu(Table table, int numRows) { - com.udf.IntegerMultiplyBy2RapidsUDF udf = new com.udf.IntegerMultiplyBy2RapidsUDF(); - return udf.evaluateColumnar(numRows, table.getColumn(1)); - }""" - -MICRO_EXECUTE_GPU_CUDA = """\ - public static ColumnVector executeGpu(Table table, int numRows) { - com.udf.IntegerMultiplyBy2NativeRapidsUDF udf = - new com.udf.IntegerMultiplyBy2NativeRapidsUDF(); - return udf.evaluateColumnar(numRows, table.getColumn(1)); - }""" diff --git a/skills/tests/test_export/scala_fixtures.py b/skills/tests/test_export/scala_fixtures.py deleted file mode 100644 index e4f1a472e49..00000000000 --- a/skills/tests/test_export/scala_fixtures.py +++ /dev/null @@ -1,194 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Scala source code fixtures for integration tests. -""" - -from .utils import replace_scala_todo_method - -NAME = "scala" -REPLACE_TODO_FN = replace_scala_todo_method -TEST_SELECTOR_FLAG = "-Dsuites" - -# --------------------------------------------------------------------------- -# UDF source code -# --------------------------------------------------------------------------- - -UDF_SOURCE = """\ -package com.udf - -class IntegerMultiplyBy2UDF extends Function1[Integer, Integer] with Serializable { - override def apply(value: Integer): Integer = { - if (value == null) null else value * 2 - } -} -""" - -RAPIDS_UDF_SOURCE = """\ -package com.udf - -import ai.rapids.cudf._ -import com.nvidia.spark.RapidsUDF -import Arm.withResource - -class IntegerMultiplyBy2RapidsUDF extends Function1[Integer, Integer] with Serializable with RapidsUDF { - override def apply(value: Integer): Integer = { - if (value == null) null else value * 2 - } - - override def evaluateColumnar(numRows: Int, args: ColumnVector*): ColumnVector = { - withResource(Scalar.fromInt(2)) { two => - args.head.mul(two) - } - } -} -""" - -SQL_SOURCE = """\ -SELECT *, - value * 2 AS result -FROM test_table -""" - -# --------------------------------------------------------------------------- -# Unit test methods -# --------------------------------------------------------------------------- - -UNIT_TEST_METHODS = { - "createTestData": """\ - def createTestData(spark: SparkSession): DataFrame = { - val schema = StructType(Seq( - StructField("id", IntegerType, nullable = false), - StructField("value", IntegerType, nullable = true) - )) - val testData = Seq( - Row(1, 123), - Row(2, 0), - Row(3, -5), - Row(4, null) - ) - spark.createDataFrame(spark.sparkContext.parallelize(testData), schema) - }""", - "registerUDF": """\ - def registerUDF(spark: SparkSession, udfName: String): Unit = { - spark.udf.register(udfName, new IntegerMultiplyBy2UDF()) - }""", - "executeUDF": """\ - def executeUDF(spark: SparkSession, udfName: String, testDF: DataFrame): DataFrame = { - testDF.createOrReplaceTempView("test_table") - spark.sql(s"SELECT *, $udfName(value) AS result FROM test_table") - }""", - "verifyUDFResults": """\ - def verifyUDFResults(resultDF: DataFrame, testDF: DataFrame): Unit = { - val results = resultDF.collect().sortBy(_.getAs[Int]("id")) - assert(results(0).getAs[Int]("result") === 246) - assert(results(1).getAs[Int]("result") === 0) - assert(results(2).getAs[Int]("result") === -10) - assert(results(3).isNullAt(results(3).fieldIndex("result"))) - }""", -} - -RAPIDS_UDF_REGISTER = """\ - def registerRapidsUDF(spark: SparkSession, udfName: String): Unit = { - spark.udf.register(udfName, new IntegerMultiplyBy2RapidsUDF()) - }""" - -NATIVE_RAPIDS_UDF_REGISTER = """\ - def registerRapidsUDF(spark: SparkSession, udfName: String): Unit = { - spark.udf.register( - udfName, - new IntegerMultiplyBy2NativeRapidsUDF(), - org.apache.spark.sql.types.IntegerType) - }""" - -# --------------------------------------------------------------------------- -# BenchUtils methods -# --------------------------------------------------------------------------- - -BENCH_GENERATE = """\ - def generateSyntheticData( - spark: SparkSession, - numRows: Long, - numPartitions: Int - ): DataFrame = { - val baseDF = spark.range(0, numRows, 1, numPartitions) - baseDF.select( - col("id"), - (rand() * 1000).cast(IntegerType).alias("value") - ) - }""" - -BENCH_CPU = """\ - def executeCpu(spark: SparkSession, df: DataFrame): DataFrame = { - import com.udf.IntegerMultiplyBy2UDF - df.createOrReplaceTempView("bench_table") - spark.udf.register("integer_multiply_by_2", new IntegerMultiplyBy2UDF()) - spark.sql("SELECT *, integer_multiply_by_2(value) AS result FROM bench_table") - }""" - -BENCH_GPU_CUDF = """\ - def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { - import com.udf.IntegerMultiplyBy2RapidsUDF - df.createOrReplaceTempView("bench_table") - spark.udf.register("integer_multiply_by_2_rapids", new IntegerMultiplyBy2RapidsUDF()) - spark.sql("SELECT *, integer_multiply_by_2_rapids(value) AS result FROM bench_table") - }""" - -BENCH_GPU_CUDA = """\ - def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { - df.createOrReplaceTempView("bench_table") - spark.udf.register( - "integer_multiply_by_2_native", - new com.udf.IntegerMultiplyBy2NativeRapidsUDF(), - org.apache.spark.sql.types.IntegerType) - spark.sql("SELECT *, integer_multiply_by_2_native(value) AS result FROM bench_table") - }""" - -BENCH_GPU_SQL = """\ - def executeGpu(spark: SparkSession, df: DataFrame): DataFrame = { - df.createOrReplaceTempView("bench_table") - val sqlSource = scala.io.Source.fromFile("src/main/resources/integer_multiply_by_2.sql") - val sqlContent = try sqlSource.mkString finally sqlSource.close() - val benchSql = sqlContent.replace("test_table", "bench_table") - spark.sql(benchSql) - }""" - -# --------------------------------------------------------------------------- -# MicroBenchRunner methods -# --------------------------------------------------------------------------- - -MICRO_PREPARE_CPU = """\ - def prepareCpuData( - hostColumns: Array[HostColumnVector], - numRows: Int - ): Array[AnyRef] = { - val values = Array.tabulate(numRows) { i => - if (hostColumns(1).isNull(i)) null - else Int.box(hostColumns(1).getInt(i)) - } - Array[AnyRef](values) - }""" - -MICRO_EXECUTE_CPU = """\ - def executeCpu(data: Array[AnyRef], numRows: Int): Unit = { - val values = data(0).asInstanceOf[Array[Integer]] - val udf = new com.udf.IntegerMultiplyBy2UDF() - var i = 0 - while (i < numRows) { - udf.apply(values(i)) - i += 1 - } - }""" - -MICRO_EXECUTE_GPU = """\ - def executeGpu(table: Table, numRows: Int): ColumnVector = { - val udf = new com.udf.IntegerMultiplyBy2RapidsUDF() - udf.evaluateColumnar(numRows, table.getColumn(1)) - }""" - -MICRO_EXECUTE_GPU_CUDA = """\ - def executeGpu(table: Table, numRows: Int): ColumnVector = { - val udf = new com.udf.IntegerMultiplyBy2NativeRapidsUDF() - udf.evaluateColumnar(numRows, table.getColumn(1)) - }""" diff --git a/skills/tests/test_export/test_jvm.py b/skills/tests/test_export/test_jvm.py deleted file mode 100644 index 44cade95b59..00000000000 --- a/skills/tests/test_export/test_jvm.py +++ /dev/null @@ -1,523 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Integration tests for the JVM export directories. -""" - -import os -import shutil -import stat -import tempfile -from pathlib import Path - -import pytest - -from . import cuda_fixtures, java_fixtures, scala_fixtures -from .utils import run_mvn, run_script - -pytestmark = pytest.mark.slow - -SKILLS_DIR = Path(__file__).resolve().parents[2] -TEMPLATES_DIR = SKILLS_DIR / "udf-gen-test" / "templates" -CUDA_TEMPLATES_DIR = ( - SKILLS_DIR - / "udf-convert-to-cuda" - / "templates" - / "cuda" -) - -LANG_CONFIGS = [java_fixtures, scala_fixtures] -TARGETS = ["cudf", "sql", "cuda"] -LANG_TARGET_PARAMS = [(cfg, target) for cfg in LANG_CONFIGS for target in TARGETS] - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _comparison_test_methods(cfg, target): - """Return the comparison test method stubs for the given target.""" - if target == "cudf": - return {"registerRapidsUDF": cfg.RAPIDS_UDF_REGISTER} - if target == "cuda": - return {"registerRapidsUDF": cfg.NATIVE_RAPIDS_UDF_REGISTER} - # target == "sql" - no additional methods - return {} - - -def _bench_utils_methods(cfg, target): - """Return the BenchUtils method stubs for the given target.""" - methods = { - "generateSyntheticData": cfg.BENCH_GENERATE, - "executeCpu": cfg.BENCH_CPU, - } - if target == "cudf": - methods["executeGpu"] = cfg.BENCH_GPU_CUDF - elif target == "cuda": - methods["executeGpu"] = cfg.BENCH_GPU_CUDA - else: - methods["executeGpu"] = cfg.BENCH_GPU_SQL - return methods - - -def _micro_bench_methods(cfg, target): - """Return the MicroBenchRunner method stubs.""" - methods = { - "prepareCpuData": cfg.MICRO_PREPARE_CPU, - "executeCpu": cfg.MICRO_EXECUTE_CPU, - } - if target == "cuda": - methods["executeGpu"] = cfg.MICRO_EXECUTE_GPU_CUDA - else: - methods["executeGpu"] = cfg.MICRO_EXECUTE_GPU - return methods - - -def _build_project_dir(cfg): - """Copy export directory to a temp directory and resolve pom.xml.""" - export_dir = TEMPLATES_DIR / cfg.NAME - tmp_dir = tempfile.mkdtemp(prefix=f"test_{cfg.NAME}_") - project_dir = os.path.join(tmp_dir, cfg.NAME) - shutil.copytree(str(export_dir), project_dir) - - return tmp_dir, project_dir - - -def _copy_cuda_templates(project_dir): - """Copy CUDA add-on templates and replace placeholders with fixture sources.""" - shutil.copytree(str(CUDA_TEMPLATES_DIR), project_dir, dirs_exist_ok=True) - - extract_script = Path(project_dir) / "native" / "scripts" / "extract-cudf-libs.sh" - extract_script.chmod( - extract_script.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH - ) - - project_path = Path(project_dir) - for rel_path in cuda_fixtures.PLACEHOLDER_FILES: - (project_path / rel_path).unlink(missing_ok=True) - - for rel_path, source in cuda_fixtures.NATIVE_SOURCE_FILES.items(): - path = project_path / rel_path - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(source) - - cmake_path = project_path / "native" / "src" / "main" / "cpp" / "CMakeLists.txt" - cmake = cmake_path.read_text() - cmake = cmake.replace( - """\ -set(SOURCE_FILES - "src/PlaceholderUDFNameJni.cpp" - "src/placeholder_udf_name.cu" -) -""", - cuda_fixtures.CMAKE_SOURCE_FILES, - ) - cmake_path.write_text(cmake) - - -def _fill_stubs(cfg, project_dir, target): - """Write UDF sources and fill in all TODO stubs in the project.""" - ext = f".{cfg.NAME}" - - def _replace_stubs(path, methods): - with open(path, "r") as f: - source = f.read() - for method_name, impl in methods.items(): - source = cfg.REPLACE_TODO_FN(source, method_name, impl) - with open(path, "w") as f: - f.write(source) - - src_dir = os.path.join(project_dir, "src", "main", cfg.NAME, "com", "udf") - test_dir = os.path.join(project_dir, "src", "test", cfg.NAME, "com", "udf") - - # Write CPU UDF source. - with open(os.path.join(src_dir, f"IntegerMultiplyBy2UDF{ext}"), "w") as f: - f.write(cfg.UDF_SOURCE) - - # Fill in unit test stubs. - _replace_stubs(os.path.join(test_dir, f"UnitTest{ext}"), cfg.UNIT_TEST_METHODS) - - if target == "cudf": - # Write RapidsUDF source. - with open(os.path.join(src_dir, f"IntegerMultiplyBy2RapidsUDF{ext}"), "w") as f: - f.write(cfg.RAPIDS_UDF_SOURCE) - - # Fill comparison test stubs. - _replace_stubs( - os.path.join(test_dir, f"CudfComparisonTest{ext}"), - _comparison_test_methods(cfg, "cudf"), - ) - - # Fill MicroBenchRunner stubs. - _replace_stubs( - os.path.join(src_dir, "bench", f"MicroBenchRunner{ext}"), - _micro_bench_methods(cfg, target), - ) - elif target == "cuda": - _copy_cuda_templates(project_dir) - - # Fill comparison test stubs. - _replace_stubs( - os.path.join(test_dir, f"CudfComparisonTest{ext}"), - _comparison_test_methods(cfg, "cuda"), - ) - - # Fill MicroBenchRunner stubs. - _replace_stubs( - os.path.join(src_dir, "bench", f"MicroBenchRunner{ext}"), - _micro_bench_methods(cfg, target), - ) - else: - # Write SQL file. - resources_dir = os.path.join(project_dir, "src", "main", "resources") - os.makedirs(resources_dir, exist_ok=True) - with open(os.path.join(resources_dir, "integer_multiply_by_2.sql"), "w") as f: - f.write(cfg.SQL_SOURCE) - - # Replace SQL file path placeholder in the comparison test. - sql_test_path = os.path.join(test_dir, f"SqlComparisonTest{ext}") - with open(sql_test_path, "r") as f: - content = f.read() - content = content.replace("placeholder_udf_name", "integer_multiply_by_2") - with open(sql_test_path, "w") as f: - f.write(content) - - # Fill comparison test stubs. - _replace_stubs(sql_test_path, _comparison_test_methods(cfg, "sql")) - - # Fill bench utils stubs. - _replace_stubs( - os.path.join(src_dir, "bench", f"BenchUtils{ext}"), - _bench_utils_methods(cfg, target), - ) - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -@pytest.fixture(scope="module", params=LANG_CONFIGS, ids=lambda c: c.NAME) -def project_dir(request): - """Clean copy of the export template with resolved pom.xml (stubs not filled).""" - cfg = request.param - tmp_dir, proj = _build_project_dir(cfg) - yield (proj, cfg) - shutil.rmtree(tmp_dir, ignore_errors=True) - - -@pytest.fixture( - scope="module", - params=LANG_TARGET_PARAMS, - ids=lambda p: f"{p[0].NAME}-{p[1]}", # "language-target" -) -def project_with_fixtures(request): - """Clean copy with resolved pom.xml and all stubs filled in.""" - cfg, target = request.param - tmp_dir, proj = _build_project_dir(cfg) - _fill_stubs(cfg, proj, target) - yield (proj, cfg, target) - shutil.rmtree(tmp_dir, ignore_errors=True) - - -@pytest.fixture( - scope="class", - params=LANG_TARGET_PARAMS, - ids=lambda p: f"{p[0].NAME}-{p[1]}", # "language-target" -) -def project_with_broken_gpu(request): - """Project with deliberately broken GPU implementation.""" - cfg, target = request.param - tmp_dir, proj = _build_project_dir(cfg) - _fill_stubs(cfg, proj, target) - - def _break_gpu_source(source: str) -> str: - # Change multiplier from 2 to 3 - source = source.replace("fromInt(2)", "fromInt(3)") - source = source.replace("* 2", "* 3") - return source - - def _insert_memory_leak(source: str) -> str: - # Inject an unclosed Scalar. - idx = source.index("evaluateColumnar") - brace = source.index("{", idx) - return ( - source[: brace + 1] + '\nScalar.fromString("LEAKED");' + source[brace + 1 :] - ) - - if target == "cudf": - path = os.path.join( - proj, - "src", - "main", - cfg.NAME, - "com", - "udf", - f"IntegerMultiplyBy2RapidsUDF.{cfg.NAME}", - ) - elif target == "cuda": - path = os.path.join( - proj, - "native", - "src", - "main", - "cpp", - "src", - "integer_multiply_by_2.cu", - ) - else: - path = os.path.join( - proj, "src", "main", "resources", "integer_multiply_by_2.sql" - ) - - # Read and overwrite with the broken source. - with open(path, "r") as f: - content = f.read() - - broken = _break_gpu_source(content) - if target == "cudf": - broken = _insert_memory_leak(broken) - - with open(path, "w") as f: - f.write(broken) - - yield (proj, cfg, target) - shutil.rmtree(tmp_dir, ignore_errors=True) - - -@pytest.fixture(scope="class", params=LANG_CONFIGS, ids=lambda c: c.NAME) -def project_with_broken_schema(request): - """Project with wrong column name in generateSyntheticData.""" - cfg = request.param - tmp_dir, proj = _build_project_dir(cfg) - _fill_stubs(cfg, proj, "cudf") - - def _break_bench_source(source: str) -> str: - # Cause a schema error due to unresolved column. - return source.replace('.alias("value")', '.alias("wrong_column")') - - ext = f".{cfg.NAME}" - bench_path = os.path.join( - proj, - "src", - "main", - cfg.NAME, - "com", - "udf", - "bench", - f"BenchUtils{ext}", - ) - - # Read and overwrite with the broken source. - with open(bench_path, "r") as f: - source = f.read() - with open(bench_path, "w") as f: - f.write(_break_bench_source(source)) - - yield (proj, cfg) - shutil.rmtree(tmp_dir, ignore_errors=True) - - -# --------------------------------------------------------------------------- -# Tests -# --------------------------------------------------------------------------- - - -class TestCompilation: - """Verify the export directory compiles.""" - - def test_compile_smoke(self, project_dir): - """ - Smoke test: compile the project as-is with TODO stubs, without - completing any of the methods, so we can catch any simple compile errors. - """ - proj, cfg = project_dir - result = run_mvn(proj, "clean", "compile") - assert result.returncode == 0, f"{cfg.NAME} smoke compile failed" - - def test_compile_with_fixtures(self, project_with_fixtures): - """ - Compile after writing UDF sources and filling in TODO stubs, to make - sure our fixtures are valid source code. - """ - proj, cfg, _target = project_with_fixtures - # test-compile compiles both main and test sources. - result = run_mvn(proj, "clean", "test-compile") - assert result.returncode == 0, f"{cfg.NAME} compile with fixtures failed" - - -class TestComparisonTest: - """Run the comparison test suite.""" - - def test_run_comparison_test(self, project_with_fixtures): - """Execute the comparison test (CudfComparisonTest or SqlComparisonTest).""" - proj, cfg, target = project_with_fixtures - if target == "sql": - suite = "com.udf.SqlComparisonTest" - else: - suite = "com.udf.CudfComparisonTest" - - result = run_mvn( - proj, - "test", - extra_args=[ - *(["-Pcuda-native-udf"] if target == "cuda" else []), - f"{cfg.TEST_SELECTOR_FLAG}={suite}", - ], - ) - assert result.returncode == 0, f"{suite} failed" - - -class TestBench: - """Test the benchmark pipeline (GenData + BenchRunner).""" - - def test_validate(self, project_with_fixtures): - """GenData validation: generate a small dataset and run validation.""" - proj, _, target = project_with_fixtures - result = run_mvn( - proj, - "compile", - "exec:java", - extra_args=[ - *(["-Pcuda-native-udf"] if target == "cuda" else []), - "-Dexec.mainClass=com.udf.bench.GenData", - "-Dexec.classpathScope=compile", - "-Dexec.args=--rows 1000 --validate --spark-conf spark.master=local[*]", - ], - ) - assert result.returncode == 0, "GenData validate failed" - - def test_spark_e2e(self, project_with_fixtures): - """End-to-end: GenData generates data, SparkBenchRunner benchmarks CPU/GPU.""" - proj, _, target = project_with_fixtures - - data_dir = os.path.join(proj, "data", "bench_input") - result_path = os.path.join(proj, "results", "bench_result.json") - mvn_args = ["--mvn-arg", "-Pcuda-native-udf"] if target == "cuda" else [] - try: - # GenData: generate parquet - gen_result = run_script( - os.path.join(proj, "run_gen_data.sh"), - args=["--rows", "1000", "--output-path", data_dir, *mvn_args], - ) - assert gen_result.returncode == 0, "run_gen_data.sh failed" - - # BenchRunner: run both benchmarks - for mode in ["cpu", "gpu"]: - bench_result = run_script( - os.path.join(proj, "run_spark_benchmark.sh"), - args=[ - "--mode", - mode, - "--data-path", - data_dir, - "--result-path", - result_path, - *mvn_args, - ], - ) - assert ( - bench_result.returncode == 0 - ), f"run_spark_benchmark.sh --mode {mode} failed" - assert os.path.isfile( - result_path - ), f"Result file not created: {result_path}" - finally: - shutil.rmtree(data_dir, ignore_errors=True) - if os.path.isfile(result_path): - os.remove(result_path) - - def test_micro_e2e(self, project_with_fixtures): - """End-to-end: GenData generates data, MicroBenchRunner benchmarks CPU/GPU.""" - proj, _, target = project_with_fixtures - if target not in {"cudf", "cuda"}: - pytest.skip("MicroBenchRunner only applies to RapidsUDF targets") - - data_dir = os.path.join(proj, "data", "micro_input") - mvn_args = ["--mvn-arg", "-Pcuda-native-udf"] if target == "cuda" else [] - try: - # GenData: generate parquet - gen_result = run_script( - os.path.join(proj, "run_gen_data.sh"), - args=["--rows", "1000", "--output-path", data_dir, *mvn_args], - ) - assert gen_result.returncode == 0, "run_gen_data.sh failed" - - # MicroBenchRunner: run both benchmarks - bench_result = run_script( - os.path.join(proj, "run_micro_benchmark.sh"), - args=["--mode", "all", "--data-path", data_dir, *mvn_args], - ) - assert bench_result.returncode == 0, ( - "run_micro_benchmark.sh failed:\n" - + bench_result.stdout - + bench_result.stderr - ) - finally: - shutil.rmtree(data_dir, ignore_errors=True) - - -class TestErrors: - """Verify that errors are caught by the test harness.""" - - def test_comparison_catches_gpu_error(self, project_with_broken_gpu): - """Comparison test should fail when GPU implementation produces wrong results.""" - proj, cfg, target = project_with_broken_gpu - if target == "sql": - suite = "com.udf.SqlComparisonTest" - else: - suite = "com.udf.CudfComparisonTest" - - result = run_mvn( - proj, - "test", - extra_args=[ - *(["-Pcuda-native-udf"] if target == "cuda" else []), - f"{cfg.TEST_SELECTOR_FLAG}={suite}", - ], - ) - assert ( - result.returncode != 0 - ), f"{suite} should have failed with broken GPU implementation" - combined = result.stdout + result.stderr - assert ( # 123 * 2 vs. 123 * 3, since we swapped multiplier - "246" in combined and "369" in combined - ), "Expected to see mismatch in test output" - - if target == "cudf": - # Re-run with debug.memory.leaks=true to verify memory leak is detected - leak_result = run_mvn( - proj, - "test", - extra_args=[ - f"{cfg.TEST_SELECTOR_FLAG}={suite}", - "-Ddebug.memory.leaks=true", - ], - ) - assert ( - leak_result.returncode != 0 - ), f"{suite} should have failed with broken GPU implementation" - leak_output = leak_result.stdout + leak_result.stderr - assert "A SCALAR WAS LEAKED" in leak_output, "Expected memory leak" - - def test_bench_validate_catches_error(self, project_with_broken_schema): - """GenData --validate should fail when synthetic data has wrong schema.""" - proj, cfg = project_with_broken_schema - result = run_mvn( - proj, - "compile", - "exec:java", - extra_args=[ - "-Dexec.mainClass=com.udf.bench.GenData", - "-Dexec.classpathScope=compile", - "-Dexec.args=--rows 1000 --validate --spark-conf spark.master=local[*]", - ], - ) - assert ( - result.returncode != 0 - ), "GenData --validate should have failed with wrong schema" - assert ( - "org.apache.spark.sql.AnalysisException" in result.stderr - ), "Expected to see schema error message" diff --git a/skills/tests/test_export/utils.py b/skills/tests/test_export/utils.py deleted file mode 100644 index d80f0e38b54..00000000000 --- a/skills/tests/test_export/utils.py +++ /dev/null @@ -1,91 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Shared test utilities for JVM skill integration tests. -""" - -import re -import subprocess -import sys -from typing import Optional - - -def run_mvn( - work_dir: str, - *goals: str, - extra_args: Optional[list[str]] = None, - timeout: int = 300, -) -> subprocess.CompletedProcess: - """Run Maven in work_dir with the given goals.""" - cmd = ["mvn", *goals, "-q"] - if extra_args: - cmd.extend(extra_args) - result = subprocess.run( - cmd, - cwd=work_dir, - capture_output=True, - text=True, - timeout=timeout, - ) - - # Write output to stdout/stderr so it is visible via pytest -s (and on errors) - if result.stdout: - sys.stdout.write(result.stdout) - if result.stderr: - sys.stderr.write(result.stderr) - return result - - -def run_script( - script_path: str, - args: Optional[list[str]] = None, - timeout: int = 300, -) -> subprocess.CompletedProcess: - """Run a bash script with the given arguments.""" - cmd = ["bash", script_path] - if args: - cmd.extend(args) - result = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=timeout, - ) - - # Write output to stdout/stderr so it is visible via pytest -s (and on errors) - if result.stdout: - sys.stdout.write(result.stdout) - if result.stderr: - sys.stderr.write(result.stderr) - return result - - -def replace_java_todo_method(source: str, method_name: str, new_body: str) -> str: - """ - Replace a TODO method in a Java source file with a real implementation. - """ - pattern = re.compile(r" public static \S+ " + re.escape(method_name) + r"\b") - match = pattern.search(source) - if not match: - raise ValueError(f"Could not find TODO method '{method_name}' in source") - - start = match.start() - brace_pos = source.index("{", start) - end_pos = source.index("}", brace_pos + 1) + 1 - return source[:start] + new_body + source[end_pos:] - - -def replace_scala_todo_method(source: str, method_name: str, new_body: str) -> str: - """ - Replace a TODO method stub in a Scala source file with a real implementation. - Assumes stubs look like "def foo(...) = ???" - """ - pattern = re.compile( - r" def " + re.escape(method_name) + r"\b.*?\?\?\?", - re.DOTALL, - ) - result = pattern.sub(new_body, source, count=1) - if result == source: - raise ValueError(f"Could not find TODO method '{method_name}' in source") - return result diff --git a/skills/tests/test_skill_frontmatter.py b/skills/tests/test_skill_frontmatter.py deleted file mode 100644 index cd2fe7c4a62..00000000000 --- a/skills/tests/test_skill_frontmatter.py +++ /dev/null @@ -1,40 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -""" -Tests for checked-in skill metadata. -""" - -from pathlib import Path - -import pytest -import yaml - -SKILLS_DIR = Path(__file__).resolve().parents[1] -SKILL_FILES = sorted(SKILLS_DIR.glob("*/SKILL.md")) - - -def _read_frontmatter(path: Path) -> str: - """Return the YAML frontmatter body from a SKILL.md file.""" - lines = path.read_text(encoding="utf-8").splitlines() - if not lines or lines[0] != "---": - raise ValueError(f"{path} must start with YAML frontmatter") - - try: - end = lines.index("---", 1) - except ValueError as e: - raise ValueError(f"{path} must close YAML frontmatter with ---") from e - - return "\n".join(lines[1:end]) - - -@pytest.mark.parametrize("skill_file", SKILL_FILES, ids=lambda p: p.parent.name) -def test_skill_frontmatter_loads(skill_file: Path) -> None: - frontmatter = _read_frontmatter(skill_file) - parsed = yaml.safe_load(frontmatter) - - assert isinstance(parsed, dict), f"{skill_file} frontmatter must parse to a map" - assert isinstance(parsed.get("name"), str), f"{skill_file} must define name" - assert isinstance( - parsed.get("description"), str - ), f"{skill_file} must define description" From 99c0795624d862c65eb1d20d6dd24bcc37f91ac6 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Fri, 12 Jun 2026 08:48:34 -0700 Subject: [PATCH 13/16] defer pyproject to follow-up --- skills/pyproject.toml | 40 ---------------------------------------- 1 file changed, 40 deletions(-) delete mode 100644 skills/pyproject.toml diff --git a/skills/pyproject.toml b/skills/pyproject.toml deleted file mode 100644 index e977287231d..00000000000 --- a/skills/pyproject.toml +++ /dev/null @@ -1,40 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -[build-system] -requires = ["setuptools>=61.0"] -build-backend = "setuptools.build_meta" - -[project] -name = "aether-agent" -version = "0.1.0" -description = "Convert Spark UDFs into GPU implementations" -authors = [ - {name = "Rishi Chandra", email = "rishic@nvidia.com"} -] -readme = "README.md" -requires-python = ">=3.10" -classifiers = [ - "Programming Language :: Python :: 3", - "Operating System :: OS Independent", -] - -[project.optional-dependencies] -dev = [ - "pytest==8.4.1", - "PyYAML==6.0.3", - "isort==6.0.1", - "black==25.1.0", - "ruff==0.12.8", -] - -[tool.setuptools] -packages = [] - -[tool.pyright] -typeCheckingMode = "standard" - -[tool.pytest.ini_options] -markers = [ - "slow: integration tests", -] From 0e7cbf84ca40932895ea0585e87b615af77c9979 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Mon, 15 Jun 2026 13:18:48 -0700 Subject: [PATCH 14/16] address comments --- .../references/NATIVE_BUILD_ENV.md | 2 +- .../templates/java/.mvn/jvm.config | 2 +- .../templates/java/run_gen_data.sh | 3 +-- .../templates/java/run_spark_benchmark.sh | 3 +-- .../java/com/udf/bench/MicroBenchRunner.java | 23 +++++++++++------- .../test/java/com/udf/CudfComparisonTest.java | 2 +- .../test/java/com/udf/SqlComparisonTest.java | 2 +- .../java/src/test/java/com/udf/UnitTest.java | 2 +- .../templates/scala/.mvn/jvm.config | 2 +- .../templates/scala/run_gen_data.sh | 3 +-- .../templates/scala/run_spark_benchmark.sh | 3 +-- .../com/udf/bench/MicroBenchRunner.scala | 24 +++++++++++++++---- .../scala/com/udf/CudfComparisonTest.scala | 2 +- .../scala/com/udf/SqlComparisonTest.scala | 2 +- .../src/test/scala/com/udf/UnitTest.scala | 2 +- 15 files changed, 46 insertions(+), 31 deletions(-) diff --git a/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md b/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md index a45d912d9a5..56af56fd34d 100644 --- a/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md +++ b/skills/udf-convert-to-cuda/references/NATIVE_BUILD_ENV.md @@ -31,7 +31,7 @@ The native build compiles against the prebuilt libcudf in the spark-rapids jar, ```bash curl -fsSL https://nvidia.github.io/spark-rapids/docs/download.html \ - | perl -0777 -ne 'while (/built against CUDA\s+(\d+\.\d+)(?:\s+or\s+CUDA\s+(\d+\.\d+))?/g) { print "$1\n"; print "$2\n" if defined $2 }' + | grep -Eo '[^<>]*built against CUDA[^<>]*' ``` 2. Check the active toolkit (`nvcc --version`). CMake uses `$CUDACXX`, else `nvcc` on `PATH`, else `$CUDAToolkit_ROOT/bin/nvcc` — the default `PATH` `nvcc` may not be the one you want. diff --git a/skills/udf-gen-test/templates/java/.mvn/jvm.config b/skills/udf-gen-test/templates/java/.mvn/jvm.config index 9d7bf1a15b2..0ae13fa9a86 100644 --- a/skills/udf-gen-test/templates/java/.mvn/jvm.config +++ b/skills/udf-gen-test/templates/java/.mvn/jvm.config @@ -1,4 +1,4 @@ --Xmx5g +-Xmx16g -ea --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.lang.invoke=ALL-UNNAMED diff --git a/skills/udf-gen-test/templates/java/run_gen_data.sh b/skills/udf-gen-test/templates/java/run_gen_data.sh index 44c802c1864..1a7b4b2adbc 100644 --- a/skills/udf-gen-test/templates/java/run_gen_data.sh +++ b/skills/udf-gen-test/templates/java/run_gen_data.sh @@ -39,8 +39,7 @@ if [ -z "$ROWS" ]; then fi SPARK_CONFS=( - --spark-conf spark.master="local[*]" - --spark-conf spark.driver.memory="16g" + --spark-conf spark.master="local[8]" --spark-conf spark.rapids.sql.enabled="true" --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" --spark-conf spark.locality.wait="0s" diff --git a/skills/udf-gen-test/templates/java/run_spark_benchmark.sh b/skills/udf-gen-test/templates/java/run_spark_benchmark.sh index 5ce799d5413..d8b2b1d1b70 100644 --- a/skills/udf-gen-test/templates/java/run_spark_benchmark.sh +++ b/skills/udf-gen-test/templates/java/run_spark_benchmark.sh @@ -45,8 +45,7 @@ if [ -z "$RESULT_PATH" ]; then fi SPARK_CONFS=( - --spark-conf spark.master="local[*]" - --spark-conf spark.driver.memory="16g" + --spark-conf spark.master="local[8]" --spark-conf spark.rapids.sql.enabled="true" --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" --spark-conf spark.locality.wait="0s" diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java index d258f94334d..041337077b1 100644 --- a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java @@ -237,14 +237,14 @@ private static Table readParquetData(String dataPath, int maxRows) { totalRows += tables[i].getRowCount(); if (maxRows > 0 && totalRows >= maxRows) break; } - Table combined = (count == 1) ? tables[0] : Table.concatenate(Arrays.copyOf(tables, count)); - try (Table src = combined) { - return limitTable(src, maxRows); + if (count == 1) { + return limitTable(tables[0], maxRows); } - } finally { - if (count > 1) { - closeAll(tables); + try (Table combined = Table.concatenate(Arrays.copyOf(tables, count))) { + return limitTable(combined, maxRows); } + } finally { + closeAll(tables); } } @@ -276,10 +276,15 @@ private static double getTableSizeMB(Table table) { /** Copy all device columns to host memory. */ private static HostColumnVector[] copyAllToHost(Table table) { HostColumnVector[] hostCols = new HostColumnVector[table.getNumberOfColumns()]; - for (int i = 0; i < hostCols.length; i++) { - hostCols[i] = table.getColumn(i).copyToHost(); + try { + for (int i = 0; i < hostCols.length; i++) { + hostCols[i] = table.getColumn(i).copyToHost(); + } + return hostCols; + } catch (Exception e) { + closeAll(hostCols); + throw e; } - return hostCols; } /** Close all resources in an array. */ diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java index 8217963425d..e8109a25cd2 100644 --- a/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/CudfComparisonTest.java @@ -22,7 +22,7 @@ public static void setUp() { origContextClassLoader = TestUtils.installMutableClassLoader(); spark = SparkSession.builder() .appName("UDF vs. RapidsUDF Comparison Test") - .master("local[*]") + .master("local[4]") .config("spark.plugins", "com.nvidia.spark.SQLPlugin") .config("spark.rapids.memory.gpu.pool", "NONE") .config("spark.rapids.sql.explain", "NONE") diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java index d335432d920..26d08fc5f44 100644 --- a/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/SqlComparisonTest.java @@ -26,7 +26,7 @@ public static void setUp() { origContextClassLoader = TestUtils.installMutableClassLoader(); spark = SparkSession.builder() .appName("UDF vs. SQL Comparison Test") - .master("local[*]") + .master("local[4]") .config("spark.plugins", "com.nvidia.spark.SQLPlugin") .config("spark.rapids.skipGpuArchitectureCheck", "true") .config("spark.rapids.sql.mode", "explainOnly") diff --git a/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java b/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java index c6fb7a26c02..17d6b177fe5 100644 --- a/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java +++ b/skills/udf-gen-test/templates/java/src/test/java/com/udf/UnitTest.java @@ -30,7 +30,7 @@ public static void setUp() { origContextClassLoader = TestUtils.installMutableClassLoader(); spark = SparkSession.builder() .appName("UDF Unit Test") - .master("local[*]") + .master("local[4]") .config("spark.plugins", "com.nvidia.spark.SQLPlugin") .config("spark.rapids.skipGpuArchitectureCheck", "true") .config("spark.rapids.sql.mode", "explainOnly") diff --git a/skills/udf-gen-test/templates/scala/.mvn/jvm.config b/skills/udf-gen-test/templates/scala/.mvn/jvm.config index 4d641a99f65..f8f3f2490b0 100644 --- a/skills/udf-gen-test/templates/scala/.mvn/jvm.config +++ b/skills/udf-gen-test/templates/scala/.mvn/jvm.config @@ -1,4 +1,4 @@ --Xmx5g +-Xmx16g --add-opens=java.base/java.lang=ALL-UNNAMED --add-opens=java.base/java.lang.invoke=ALL-UNNAMED --add-opens=java.base/java.lang.reflect=ALL-UNNAMED diff --git a/skills/udf-gen-test/templates/scala/run_gen_data.sh b/skills/udf-gen-test/templates/scala/run_gen_data.sh index 44c802c1864..1a7b4b2adbc 100644 --- a/skills/udf-gen-test/templates/scala/run_gen_data.sh +++ b/skills/udf-gen-test/templates/scala/run_gen_data.sh @@ -39,8 +39,7 @@ if [ -z "$ROWS" ]; then fi SPARK_CONFS=( - --spark-conf spark.master="local[*]" - --spark-conf spark.driver.memory="16g" + --spark-conf spark.master="local[8]" --spark-conf spark.rapids.sql.enabled="true" --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" --spark-conf spark.locality.wait="0s" diff --git a/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh b/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh index 5ce799d5413..d8b2b1d1b70 100644 --- a/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh +++ b/skills/udf-gen-test/templates/scala/run_spark_benchmark.sh @@ -45,8 +45,7 @@ if [ -z "$RESULT_PATH" ]; then fi SPARK_CONFS=( - --spark-conf spark.master="local[*]" - --spark-conf spark.driver.memory="16g" + --spark-conf spark.master="local[8]" --spark-conf spark.rapids.sql.enabled="true" --spark-conf spark.plugins="com.nvidia.spark.SQLPlugin" --spark-conf spark.locality.wait="0s" diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala index 3cf59d07fe1..9944f726743 100644 --- a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala @@ -224,11 +224,15 @@ object MicroBenchRunner { tables += t totalRows += t.getRowCount } - val combined = if (tables.length == 1) tables(0) - else Table.concatenate(tables.toArray: _*) - withResource(combined) { src => limitTable(src, maxRows) } + if (tables.length == 1) { + limitTable(tables(0), maxRows) + } else { + withResource(Table.concatenate(tables.toArray: _*)) { combined => + limitTable(combined, maxRows) + } + } } finally { - if (tables.length > 1) closeAll(tables.toArray) + closeAll(tables.toArray) } } @@ -256,7 +260,17 @@ object MicroBenchRunner { /** Copy all device columns to host memory. */ private def copyAllToHost(table: Table): Array[HostColumnVector] = { - Array.tabulate(table.getNumberOfColumns)(i => table.getColumn(i).copyToHost()) + val hostCols = new Array[HostColumnVector](table.getNumberOfColumns) + try { + for (i <- hostCols.indices) { + hostCols(i) = table.getColumn(i).copyToHost() + } + hostCols + } catch { + case e: Throwable => + closeAll(hostCols) + throw e + } } /** Parse CLI arguments. */ diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala index ddc481c28b4..32361dc95c0 100644 --- a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/CudfComparisonTest.scala @@ -16,7 +16,7 @@ class CudfComparisonTest extends AnyFunSuite with BeforeAndAfterAll { override def beforeAll(): Unit = { spark = SparkSession.builder() .appName("UDF vs. RapidsUDF Comparison Test") - .master("local[*]") + .master("local[4]") .config("spark.plugins", "com.nvidia.spark.SQLPlugin") .config("spark.rapids.memory.gpu.pool", "NONE") .config("spark.rapids.sql.explain", "NONE") diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala index bb83e405665..8aa167dbc22 100644 --- a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/SqlComparisonTest.scala @@ -16,7 +16,7 @@ class SqlComparisonTest extends AnyFunSuite with BeforeAndAfterAll { override def beforeAll(): Unit = { spark = SparkSession.builder() .appName("UDF vs. SQL Comparison Test") - .master("local[*]") + .master("local[4]") .config("spark.plugins", "com.nvidia.spark.SQLPlugin") .config("spark.rapids.skipGpuArchitectureCheck", "true") .config("spark.rapids.sql.mode", "explainOnly") diff --git a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala index eceffabca6e..71298f472c3 100644 --- a/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala +++ b/skills/udf-gen-test/templates/scala/src/test/scala/com/udf/UnitTest.scala @@ -73,7 +73,7 @@ class UnitTest extends AnyFunSuite with BeforeAndAfterAll { override def beforeAll(): Unit = { spark = SparkSession.builder() .appName("UDF Unit Test") - .master("local[*]") + .master("local[4]") .config("spark.plugins", "com.nvidia.spark.SQLPlugin") .config("spark.rapids.skipGpuArchitectureCheck", "true") .config("spark.rapids.sql.mode", "explainOnly") From f07e94390f44f0dd640e36f7b71d64d1b4c3f685 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Mon, 15 Jun 2026 15:32:31 -0700 Subject: [PATCH 15/16] add a note on heap size --- skills/udf-benchmark/SKILL.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/skills/udf-benchmark/SKILL.md b/skills/udf-benchmark/SKILL.md index 66e5005a748..d1efc8c8d24 100644 --- a/skills/udf-benchmark/SKILL.md +++ b/skills/udf-benchmark/SKILL.md @@ -49,6 +49,8 @@ If validation fails, analyze the error and fix the BenchUtils implementation. ## Step 3: Generate Data and Run Benchmarks +The scripts set the default heap size to 16g in `.mvn/jvm.config`; adjust depending on data size. + ### Generate benchmark data (10M rows): ```bash ./run_gen_data.sh --rows 10000000 From 466dd349cdb033406cbd364e6007f9eb4729e0c9 Mon Sep 17 00:00:00 2001 From: Rishi Chandra Date: Tue, 16 Jun 2026 07:48:12 -0700 Subject: [PATCH 16/16] exit(1) on errors in microbenchmarks --- .../java/src/main/java/com/udf/bench/MicroBenchRunner.java | 2 ++ .../scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala | 2 ++ 2 files changed, 4 insertions(+) diff --git a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java index 041337077b1..877bfe214c4 100644 --- a/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java +++ b/skills/udf-gen-test/templates/java/src/main/java/com/udf/bench/MicroBenchRunner.java @@ -164,6 +164,7 @@ public static void main(String[] args) { } catch (Exception e) { System.err.printf("CPU benchmark failed: %s%n", e.getMessage()); e.printStackTrace(System.err); + System.exit(1); } finally { closeAll(hostColumns); } @@ -182,6 +183,7 @@ public static void main(String[] args) { } catch (Exception e) { System.err.printf("GPU benchmark failed: %s%n", e.getMessage()); e.printStackTrace(System.err); + System.exit(1); } } diff --git a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala index 9944f726743..f1f22cc4469 100644 --- a/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala +++ b/skills/udf-gen-test/templates/scala/src/main/scala/com/udf/bench/MicroBenchRunner.scala @@ -155,6 +155,7 @@ object MicroBenchRunner { case e: Exception => System.err.println(s"CPU benchmark failed: ${e.getMessage}") e.printStackTrace(System.err) + System.exit(1) } finally { closeAll(hostColumns) } @@ -175,6 +176,7 @@ object MicroBenchRunner { case e: Exception => System.err.println(s"GPU benchmark failed: ${e.getMessage}") e.printStackTrace(System.err) + System.exit(1) } }