load all kryo versions

Author: skoppu22

cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/KryoRegister.java

@@ -22,7 +22,10 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.LinkedHashMap;
+import java.util.LinkedHashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -31,7 +34,6 @@
 
 import com.esotericsoftware.kryo.Kryo;
 import com.esotericsoftware.kryo.Serializer;
-import org.apache.cassandra.bridge.BaseCassandraBridgeFactory;
 import org.apache.cassandra.bridge.BigNumberConfigImpl;
 import org.apache.cassandra.bridge.CassandraBridgeFactory;
 import org.apache.cassandra.bridge.CassandraVersion;
@@ -142,26 +144,33 @@ public static void setup(@NotNull SparkConf configuration)
         LOGGER.info("Setting up Kryo");
         configuration.set(SPARK_SERIALIZER, "org.apache.spark.serializer.KryoSerializer");
 
-        // Add KryoRegister to SparkConf serialization if not already there
+        // Preserve any pre-existing (e.g. user-supplied) registrators and keep them first.
+        // LinkedHashSet gives a stable, predictable registration order on the driver; the same
+        // resulting spark.kryo.registrator string is then propagated to all executors via SparkConf.
         Set<String> registratorsSet = Arrays.stream(configuration.get(SPARK_REGISTRATORS, "").split(","))
                                             .filter(string -> string != null && !string.isEmpty())
-                                            .collect(Collectors.toSet());
-
-        // TODO: Find a better way to initialize Kryo serializer, instead of relaying
-        //  on Cassandra version specified as parameter of Spark job. Can we get Cassandra version from Sidecar?
-        CassandraVersion cassandraVersion = BaseCassandraBridgeFactory.getCassandraVersion(configuration.get(CASSANDRA_VERSION, "4.0.0"));
-        Class<?> registratorClass = KRYO_REGISTRATORS.get(cassandraVersion);
-        if (registratorClass == null)
+                                            .collect(Collectors.toCollection(LinkedHashSet::new));
+
+        // SSTable based bridge selection feature selects the bridge version, which may differ
+        // from cassandra.version; registering every loadable bridge's registrator ensures Spark
+        // can serialize objects for whichever bridge is chosen. Only implemented (bundled) versions
+        // are used, so we never attempt to load a bridge JAR that is not available.
+        List<Class<?>> registratorClasses = Arrays.stream(CassandraVersion.implementedVersions())
+                                                  .map(KRYO_REGISTRATORS::get)
+                                                  .filter(Objects::nonNull)
+                                                  .collect(Collectors.toList());
+        if (registratorClasses.isEmpty())
         {
-            throw new IllegalArgumentException("Kryo registrator not configured for Cassandra version: " + cassandraVersion);
+            throw new IllegalStateException("No Kryo registrators configured for implemented Cassandra versions: "
+                                            + Arrays.toString(CassandraVersion.implementedVersions()));
         }
 
-        registratorsSet.add(registratorClass.getName());
+        registratorClasses.forEach(registratorClass -> registratorsSet.add(registratorClass.getName()));
         String registratorsString = String.join(",", registratorsSet);
         LOGGER.info("Setting kryo registrators: " + registratorsString);
         configuration.set(SPARK_REGISTRATORS, registratorsString);
 
-        configuration.registerKryoClasses(new Class<?>[]{registratorClass});
+        configuration.registerKryoClasses(registratorClasses.toArray(new Class<?>[0]));
     }
 
     public static class V40 extends KryoRegister

cassandra-analytics-core/src/test/java/org/apache/cassandra/spark/KryoRegisterTest.java

@@ -19,9 +19,14 @@
 
 package org.apache.cassandra.spark;
 
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+
 import org.junit.jupiter.api.Test;
 
 import org.apache.cassandra.bridge.CassandraVersion;
+import org.apache.spark.SparkConf;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatNoException;
@@ -92,4 +97,48 @@ void testValidateWithNullCassandraVersionString()
         .describedAs("Validation should work with null cassandraVersion string")
         .isThrownBy(() -> KryoRegister.validateKryoRegistratorExists(CassandraVersion.FOURZERO, null));
     }
+
+    @Test
+    void testSetupRegistersAllImplementedVersions()
+    {
+        SparkConf conf = new SparkConf();
+        KryoRegister.setup(conf);
+
+        List<String> expected = Arrays.stream(CassandraVersion.implementedVersions())
+                                      .map(KryoRegister.KRYO_REGISTRATORS::get)
+                                      .map(Class::getName)
+                                      .collect(Collectors.toList());
+
+        // setup() must register a registrator for every implemented (bundled) bridge version,
+        // independent of spark.cassandra_analytics.cassandra.version, so serialization works for
+        // whichever bridge the SSTable-version analyzer selects at runtime.
+        assertThat(expected).isNotEmpty();
+        List<String> registrators = Arrays.asList(conf.get("spark.kryo.registrator").split(","));
+        assertThat(registrators).containsAll(expected);
+    }
+
+    @Test
+    void testSetupDoesNotDependOnCassandraVersionConfig()
+    {
+        // Even with a cassandra.version that differs from the bridge that may be selected,
+        // setup() registers all implemented versions and never throws based on that config.
+        SparkConf conf = new SparkConf()
+                         .set("spark.cassandra_analytics.cassandra.version", "5.0.0");
+        assertThatNoException().isThrownBy(() -> KryoRegister.setup(conf));
+        assertThat(conf.get("spark.serializer")).isEqualTo("org.apache.spark.serializer.KryoSerializer");
+    }
+
+    @Test
+    void testSetupPreservesExistingRegistrators()
+    {
+        SparkConf conf = new SparkConf()
+                         .set("spark.kryo.registrator", "com.example.CustomRegistrator");
+        KryoRegister.setup(conf);
+
+        List<String> registrators = Arrays.asList(conf.get("spark.kryo.registrator").split(","));
+        assertThat(registrators).contains("com.example.CustomRegistrator");
+        assertThat(registrators.get(0))
+        .describedAs("pre-existing registrators should be kept first (deterministic order)")
+        .isEqualTo("com.example.CustomRegistrator");
+    }
 }