from_protobuf: Re-enable protobuf integration tests on OSS Spark and Databricks [databricks] [fast-ut] [reduced-it]

integration_tests/README.md

@@ -86,6 +86,8 @@ For manual installation, you need to setup your environment:
 - pytest
   : A framework that makes it easy to write small, readable tests, and can scale to support complex
   functional testing for applications and libraries (requires  Python 3.6+).
+- protobuf
+  : Provides Protocol Buffers APIs for protobuf integration-test fixtures.
 - sre_yield
   : Provides a set of APIs to generate string data from a regular expression.
 - pandas
@@ -299,6 +301,18 @@ You do need to have access to a compatible GPU with the needed CUDA drivers. The
 `--runtime_env` is used to specify the environment you are running the tests in. Valid values are `databricks`,`emr`,`dataproc`,`dataproc_serverless` and `apache`. This is generally used
 when certain environments have different behavior, and the tests don't have a good way to auto-detect the environment yet.
 
+#### Protobuf tests on Databricks
+
+On Databricks, `INCLUDE_SPARK_PROTOBUF_JAR` controls only external `spark-protobuf` jar injection; it
+does not control protobuf test eligibility. Apache Spark runs require a matching external jar and
+skip the protobuf tests when this variable is set to `false`. Databricks runs use the runtime-bundled
+protobuf implementation instead, so `run_pyspark_from_build.sh --runtime_env=databricks` does not
+inject a matching jar from either the build dependencies or `LOCAL_JAR_PATH`, even if the variable
+is explicitly set to `true`.
+
+The smoke tests detect the bundled runtime independently and use a static descriptor set, so they do
+not depend on Spark's private, runtime-specific shaded protobuf classes.
+
 ### timezone
 
 The RAPIDS plugin currently only supports the UTC time zone. Spark uses the default system time zone unless explicitly set otherwise.

integration_tests/pom.xml

@@ -132,6 +132,7 @@
                                     <includes>
                                         <include>parquet-hadoop*.jar</include>
                                         <include>spark-avro*.jar</include>
+                                        <include>spark-protobuf*.jar</include>
                                     </includes>
                                 </filesets>
                             </filesets>
@@ -166,6 +167,24 @@
                             </artifactItems>
                         </configuration>
                     </execution>
+                    <execution>
+                        <id>copy-spark-protobuf</id>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>copy</goal>
+                        </goals>
+                        <configuration>
+                            <skip>${spark.protobuf.copy.skip}</skip>
+                            <useBaseVersion>true</useBaseVersion>
+                            <artifactItems>
+                                <artifactItem>
+                                    <groupId>org.apache.spark</groupId>
+                                    <artifactId>spark-protobuf_${scala.binary.version}</artifactId>
+                                    <version>${spark.version}</version>
+                                </artifactItem>
+                            </artifactItems>
+                        </configuration>
+                    </execution>
                 </executions>
             </plugin>
             <plugin>

integration_tests/requirements.txt

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 pytest
+protobuf
 sre_yield
 pandas
 pyarrow == 17.0.0 ; python_version == '3.8'

integration_tests/run_pyspark_from_build.sh

@@ -29,6 +29,8 @@
 #   - SPARK_HOME: Path to your Apache Spark installation.
 #   - SKIP_TESTS: If set to true, skips running the Python integration tests.
 #   - INCLUDE_SPARK_AVRO_JAR: If set to true, includes Avro tests.
+#   - INCLUDE_SPARK_PROTOBUF_JAR: Controls external spark-protobuf jar injection; setting it to
+#                                false also disables protobuf tests on Apache Spark.
 #   - TEST: Specifies a specific test to run.
 #   - TEST_TAGS: Allows filtering tests based on tags.
 #   - TEST_TYPE: Specifies the type of tests to run.
@@ -48,6 +50,9 @@
 #   To run all tests, including Avro tests:
 #     INCLUDE_SPARK_AVRO_JAR=true ./run_pyspark_from_build.sh
 #
+#   To run without injecting an external spark-protobuf jar:
+#     INCLUDE_SPARK_PROTOBUF_JAR=false ./run_pyspark_from_build.sh
+#
 #   To run a specific test:
 #     TEST=my_test ./run_pyspark_from_build.sh
 #
@@ -60,6 +65,28 @@
 
 set -ex
 
+is_databricks_runtime_arg() {
+    local expect_runtime_env_value=false
+    local runtime_env=""
+    local arg
+    for arg in "$@"; do
+        if [[ "$expect_runtime_env_value" == "true" ]]; then
+            runtime_env="$arg"
+            expect_runtime_env_value=false
+            continue
+        fi
+        case "$arg" in
+            --runtime_env)
+                expect_runtime_env_value=true
+                ;;
+            --runtime_env=*)
+                runtime_env="${arg#*=}"
+                ;;
+        esac
+    done
+    [[ "${runtime_env,,}" == "databricks" ]]
+}
+
 SCRIPTPATH="$( cd "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 cd "$SCRIPTPATH"
 
@@ -102,6 +129,7 @@ else
     # support alternate local jars NOT building from the source code
     if [ -d "$LOCAL_JAR_PATH" ]; then
         AVRO_JARS=$(echo "$LOCAL_JAR_PATH"/spark-avro*.jar)
+        PROTOBUF_JARS=$(echo "$LOCAL_JAR_PATH"/spark-protobuf*.jar)
         PLUGIN_JAR=$(echo "$LOCAL_JAR_PATH"/rapids-4-spark_*.jar)
         if [ -f $(echo $LOCAL_JAR_PATH/parquet-hadoop*.jar) ]; then
             export INCLUDE_PARQUET_HADOOP_TEST_JAR=true
@@ -118,6 +146,7 @@ else
     else
         [[ "$SCALA_VERSION" != "2.12"  ]] && TARGET_DIR=${TARGET_DIR/integration_tests/scala$SCALA_VERSION\/integration_tests}
         AVRO_JARS=$(echo "$TARGET_DIR"/dependency/spark-avro*.jar)
+        PROTOBUF_JARS=$(echo "$TARGET_DIR"/dependency/spark-protobuf*.jar)
         PARQUET_HADOOP_TESTS=$(echo "$TARGET_DIR"/dependency/parquet-hadoop*.jar)
         # remove the log4j.properties file so it doesn't conflict with ours, ignore errors
         # if it isn't present or already removed
@@ -143,9 +172,29 @@ else
         AVRO_JARS=""
     fi
 
-    # ALL_JARS includes dist.jar integration-test.jar avro.jar parquet.jar if they exist
+    INCLUDE_SPARK_PROTOBUF_JAR_REQUESTED=$(echo "${INCLUDE_SPARK_PROTOBUF_JAR}" | tr '[:upper:]' '[:lower:]')
+    PROTOBUF_JAR_COUNT=$(readlink -e $PROTOBUF_JARS 2>/dev/null | wc -l)
+    if is_databricks_runtime_arg "$@"; then
+        export INCLUDE_SPARK_PROTOBUF_JAR=false
+        PROTOBUF_JARS=""
+    elif [[ "$INCLUDE_SPARK_PROTOBUF_JAR_REQUESTED" != "false" \
+          && "$PROTOBUF_JAR_COUNT" -eq 1 ]];
+    then
+        export INCLUDE_SPARK_PROTOBUF_JAR=true
+    else
+        if [[ "$INCLUDE_SPARK_PROTOBUF_JAR_REQUESTED" != "false" \
+              && "$PROTOBUF_JAR_COUNT" -gt 1 ]]; then
+            >&2 echo "WARNING: Multiple spark-protobuf jars were found (matched: $PROTOBUF_JARS); not injecting spark-protobuf."
+        elif [[ "$INCLUDE_SPARK_PROTOBUF_JAR_REQUESTED" == "true" ]]; then
+            >&2 echo "WARNING: INCLUDE_SPARK_PROTOBUF_JAR=true was requested but a spark-protobuf jar was not found (searched: $PROTOBUF_JARS)."
+        fi
+        export INCLUDE_SPARK_PROTOBUF_JAR=false
+        PROTOBUF_JARS=""
+    fi
+
+    # ALL_JARS includes dist.jar integration-test.jar avro.jar parquet.jar protobuf.jar if they exist
     # Remove non-existing paths and canonicalize the paths including get rid of links and `..`
-    ALL_JARS=$(readlink -e $PLUGIN_JAR $TEST_JARS $AVRO_JARS $PARQUET_HADOOP_TESTS || true)
+    ALL_JARS=$(readlink -e $PLUGIN_JAR $TEST_JARS $AVRO_JARS $PARQUET_HADOOP_TESTS $PROTOBUF_JARS || true)
     # `:` separated jars
     ALL_JARS="${ALL_JARS//$'\n'/:}"
 

integration_tests/src/main/python/conftest.py

@@ -16,6 +16,8 @@
 import os
 import pytest
 import random
+import shutil
+import tempfile
 import warnings
 
 # TODO redo _spark stuff using fixtures
@@ -715,6 +717,16 @@ def spark_tmp_path(request):
     if not debug:
         fs.delete(path)
 
+# Driver-local counterpart to spark_tmp_path; spark_tmp_path lives in the
+# default Hadoop FS, which is not local on distributed setups.
+@pytest.fixture
+def local_tmp_path(request):
+    debug = request.config.getoption('debug_tmp_path')
+    ret = tempfile.mkdtemp(prefix='pyspark_tests_')
+    yield ret
+    if not debug:
+        shutil.rmtree(ret, ignore_errors=True)
+
 class TmpTableFactory:
   def __init__(self, base_id):
       self.base_id = base_id

integration_tests/src/main/python/protobuf_test.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+
+import pytest
+
+from asserts import assert_gpu_fallback_collect
+from marks import allow_non_gpu
+from spark_session import is_spark_protobuf_available
+import pyspark.sql.functions as f
+
+pytestmark = pytest.mark.skipif(
+    not is_spark_protobuf_available(), reason="from_protobuf is unavailable")
+
+
+@pytest.fixture(scope="module")
+def from_protobuf_fn():
+    from pyspark.sql.protobuf.functions import from_protobuf
+    return from_protobuf
+
+
+def _encode_varint(value):
+    out = bytearray()
+    value &= 0xFFFFFFFFFFFFFFFF
+    while True:
+        bits = value & 0x7F
+        value >>= 7
+        if value:
+            out.append(bits | 0x80)
+        else:
+            out.append(bits)
+            return bytes(out)
+
+
+def _encode_simple_message(i32_value, s_value):
+    buf = bytearray()
+    buf += _encode_varint((1 << 3) | 0)  # field 1, VARINT
+    buf += _encode_varint(i32_value)
+    s_bytes = s_value.encode("utf-8")
+    buf += _encode_varint((2 << 3) | 2)  # field 2, LENGTH-DELIMITED
+    buf += _encode_varint(len(s_bytes))
+    buf += s_bytes
+    return bytes(buf)
+
+
+# Avoid depending on whichever unshaded protobuf runtime the Spark driver provides.
+_simple_desc_bytes = bytes.fromhex(
+    "0a360a0c73696d706c652e70726f746f12047465737422200a0653696d706c65"
+    "120b0a0369333218012001280512090a0173180220012809")
+
+
+@pytest.fixture
+def simple_desc(local_tmp_path):
+    # from_protobuf reads descFilePath via java.io.File on the driver.
+    desc_path = local_tmp_path + "/simple.desc"
+    with open(desc_path, "wb") as fp:
+        fp.write(_simple_desc_bytes)
+    return desc_path, _simple_desc_bytes
+
+
+_smoke_rows = [(1, "a"), (-2, "bb"), (0, ""), (12345, "hello")]
+
+
+def _make_smoke_df(spark):
+    encoded = [(_encode_simple_message(i, s),) for (i, s) in _smoke_rows]
+    return spark.createDataFrame(encoded, ["bin"])
+
+
+@allow_non_gpu("ProjectExec", "ProtobufDataToCatalyst")
+def test_from_protobuf_smoke_path_api(simple_desc, from_protobuf_fn):
+    desc_path, _ = simple_desc
+
+    def run(spark):
+        return _make_smoke_df(spark).select(
+            from_protobuf_fn(f.col("bin"), "test.Simple", desc_path).alias("d"))
+
+    assert_gpu_fallback_collect(run, "ProtobufDataToCatalyst")
+
+
+@allow_non_gpu("ProjectExec", "ProtobufDataToCatalyst")
+def test_from_protobuf_smoke_binary_descriptor_api(simple_desc, from_protobuf_fn):
+    if "binaryDescriptorSet" not in inspect.signature(from_protobuf_fn).parameters:
+        pytest.skip("binaryDescriptorSet kwarg is Spark 3.5+ only")
+    _, desc_bytes = simple_desc
+
+    def run(spark):
+        return _make_smoke_df(spark).select(
+            from_protobuf_fn(f.col("bin"), "test.Simple",
+                             binaryDescriptorSet=bytearray(desc_bytes)).alias("d"))
+
+    assert_gpu_fallback_collect(run, "ProtobufDataToCatalyst")

integration_tests/src/main/python/spark_session.py

@@ -337,6 +337,13 @@ def is_databricks143():
 def is_databricks173_or_later():
     return is_databricks_version_or_later(17, 3)
 
+def is_spark_protobuf_available():
+    if is_databricks_runtime():
+        return is_databricks122_or_later()
+    return (is_spark_340_or_later() and
+            os.environ.get('INCLUDE_SPARK_PROTOBUF_JAR', 'true').lower() != 'false')
+
+
 def supports_delta_lake_deletion_vectors():
     """Whether the current Delta Lake runtime provides the deletion-vector feature."""
     if is_databricks_runtime():

pom.xml

@@ -98,6 +98,7 @@
                 <rapids.delta.artifactId1>rapids-4-spark-delta-21x</rapids.delta.artifactId1>
                 <rapids.delta.artifactId2>rapids-4-spark-delta-22x</rapids.delta.artifactId2>
                 <rapids.delta.artifactId3>rapids-4-spark-delta-23x</rapids.delta.artifactId3>
+                <spark.protobuf.copy.skip>true</spark.protobuf.copy.skip>
             </properties>
             <dependencies>
                 <dependency>
@@ -130,6 +131,7 @@
                 <rapids.delta.artifactId1>rapids-4-spark-delta-21x</rapids.delta.artifactId1>
                 <rapids.delta.artifactId2>rapids-4-spark-delta-22x</rapids.delta.artifactId2>
                 <rapids.delta.artifactId3>rapids-4-spark-delta-23x</rapids.delta.artifactId3>
+                <spark.protobuf.copy.skip>true</spark.protobuf.copy.skip>
             </properties>
             <modules>
                 <module>delta-lake/delta-21x</module>
@@ -154,6 +156,7 @@
                 <rapids.delta.artifactId1>rapids-4-spark-delta-21x</rapids.delta.artifactId1>
                 <rapids.delta.artifactId2>rapids-4-spark-delta-22x</rapids.delta.artifactId2>
                 <rapids.delta.artifactId3>rapids-4-spark-delta-23x</rapids.delta.artifactId3>
+                <spark.protobuf.copy.skip>true</spark.protobuf.copy.skip>
             </properties>
             <modules>
                 <module>delta-lake/delta-21x</module>
@@ -178,6 +181,7 @@
                 <rapids.delta.artifactId1>rapids-4-spark-delta-21x</rapids.delta.artifactId1>
                 <rapids.delta.artifactId2>rapids-4-spark-delta-22x</rapids.delta.artifactId2>
                 <rapids.delta.artifactId3>rapids-4-spark-delta-23x</rapids.delta.artifactId3>
+                <spark.protobuf.copy.skip>true</spark.protobuf.copy.skip>
             </properties>
             <modules>
                 <module>delta-lake/delta-21x</module>
@@ -202,6 +206,7 @@
                 <rapids.delta.artifactId1>rapids-4-spark-delta-21x</rapids.delta.artifactId1>
                 <rapids.delta.artifactId2>rapids-4-spark-delta-22x</rapids.delta.artifactId2>
                 <rapids.delta.artifactId3>rapids-4-spark-delta-23x</rapids.delta.artifactId3>
+                <spark.protobuf.copy.skip>true</spark.protobuf.copy.skip>
             </properties>
             <modules>
                 <module>delta-lake/delta-21x</module>
@@ -1050,6 +1055,8 @@
         <rapids.shimplify.skip>false</rapids.shimplify.skip>
         <rapids.build.info.skip>false</rapids.build.info.skip>
         <buildver>330</buildver>
+        <!-- spark-protobuf is a Spark 3.4.0+ module; release33x profiles override to `true`. -->
+        <spark.protobuf.copy.skip>false</spark.protobuf.copy.skip>
         <maven.compiler.source>1.8</maven.compiler.source>
         <java.major.version>8</java.major.version>
         <scala.compiler.release>${java.major.version}</scala.compiler.release>

scala2.13/integration_tests/pom.xml

@@ -132,6 +132,7 @@
                                     <includes>
                                         <include>parquet-hadoop*.jar</include>
                                         <include>spark-avro*.jar</include>
+                                        <include>spark-protobuf*.jar</include>
                                     </includes>
                                 </filesets>
                             </filesets>
@@ -166,6 +167,24 @@
                             </artifactItems>
                         </configuration>
                     </execution>
+                    <execution>
+                        <id>copy-spark-protobuf</id>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>copy</goal>
+                        </goals>
+                        <configuration>
+                            <skip>${spark.protobuf.copy.skip}</skip>
+                            <useBaseVersion>true</useBaseVersion>
+                            <artifactItems>
+                                <artifactItem>
+                                    <groupId>org.apache.spark</groupId>
+                                    <artifactId>spark-protobuf_${scala.binary.version}</artifactId>
+                                    <version>${spark.version}</version>
+                                </artifactItem>
+                            </artifactItems>
+                        </configuration>
+                    </execution>
                 </executions>
             </plugin>
             <plugin>