Integrating TensorFlow for Java with Spark‑Scala for Distributed Machine Learning Prediction

In Qunar's intelligent risk‑control scenario, the team needed to predict user risk from massive hourly feature data (≈3 million records) within one hour, requiring a high‑performance distributed solution.

Framework selection : After evaluating pure machine‑learning frameworks (TensorFlow, PyTorch) and big‑data platforms (Spark, Flink, Hadoop), the team chose a hybrid approach—using Spark for distributed data processing and TensorFlow for model inference—because Spark’s in‑memory batch processing offers superior performance for offline prediction workloads.

Why TensorFlow for Java & Spark‑Scala? The Java/Scala stack avoids the extra Python‑to‑Java conversion layer present in PySpark, reducing overhead. Both TensorFlow for Java and Spark‑Scala run natively on the JVM, providing better latency for per‑task inference.

Performance comparison : Benchmarks showed that Spark‑Scala outperforms PySpark due to (1) elimination of the Python‑to‑Java bridge on the driver and (2) avoiding per‑task Python process creation on executors. TensorFlow for Java and TensorFlow for Python have similar core performance, but Java incurs less preprocessing overhead.

Implementation details :

Model training is performed in Python using TensorFlow + Keras (MNIST CNN example). The trained model is saved in protobuf format for cross‑platform loading.

Model loading in Java: val bundle = tf.SavedModelBundle.load(modelPath, modelTag)

Inference code creates a Tensor from feature arrays and runs the session: val y = sess.runner().feed("serving_default_hmc_input:0", x).fetch("StatefulPartitionedCall:0").run().get(0)

Integration with Spark: the prediction method is registered as a UDF and applied to a DataFrame, or batch‑processed via mapPartitions to reduce object creation.

Optimization & pitfalls :

Batch prediction with mapPartitions cuts execution time from ~20 min to ~9 min for 3 million records.

Model hot‑update is achieved by storing the SavedModel on HDFS and loading it at runtime, avoiding service redeployment.

Native library errors (missing libtensorflow_jni.so ) were resolved by running Spark in local mode to ensure consistent C library versions across nodes.

Deployment : The assembled JAR (including Spark‑Scala and TensorFlow dependencies) is submitted to a YARN cluster with appropriate driver and executor resources.

Conclusion : Combining TensorFlow for Java with Spark‑Scala provides an effective, high‑throughput solution for large‑scale offline model prediction, demonstrating the importance of framework selection, performance tuning, and operational considerations in big‑data AI pipelines.