1 /***********************************************************************
2  * Copyright (c) 2013-2024 Commonwealth Computer Research, Inc.
3  * All rights reserved. This program and the accompanying materials
4  * are made available under the terms of the Apache License, Version 2.0
5  * which accompanies this distribution and is available at
6  * http://www.opensource.org/licenses/apache2.0.php.
7  ***********************************************************************/
8 
9 package org.apache.spark.geomesa.api.python
10 
11 import com.typesafe.scalalogging.LazyLogging
12 import net.razorvine.pickle.{IObjectPickler, Pickler}
13 import org.apache.spark.api.java.JavaRDD
14 import org.apache.spark.api.python.SerDeUtil
15 import org.apache.spark.rdd.RDD
16 import org.locationtech.geomesa.utils.text.WKBUtils
17 import org.locationtech.jts.geom.Geometry
18 
19 import java.io.OutputStream
20 
21 object GeoMesaSeDerUtil extends LazyLogging {
22 
23   implicit def toPickledRDD(rdd: RDD[_]): JavaRDD[Array[Byte]] =
24     rdd.mapPartitions{ i =>  AutoBatchedPickler(i) }.toJavaRDD
25 
26   implicit def toPickledRDD(jrdd: JavaRDD[_]): JavaRDD[Array[Byte]] = toPickledRDD(jrdd.rdd)
27 
28 }
29 
30 object AutoBatchedPickler extends LazyLogging {
31 
32   val module = "geomesa_pyspark.types"
33   val function = "_deserialize_from_wkb"
34 
35   val pickler = new IObjectPickler {
36     def pickle(o: Object, os: OutputStream, p: Pickler): Unit = {
37       import net.razorvine.pickle.Opcodes
38       os.write(Opcodes.GLOBAL)
39       os.write(s"$module\n$function\n".getBytes)
40       p.save(WKBUtils.write(o.asInstanceOf[Geometry]))
41       os.write(Opcodes.TUPLE1)
42       os.write(Opcodes.REDUCE)
43     }
44   }
45   Pickler.registerCustomPickler(classOf[Geometry], pickler)
46   logger.info("registered JTS Geometry to WKB pickler")
47 
48   def apply(i: Iterator[Any]): AutoBatchedPickler =  new AutoBatchedPickler(i)
49 }
50 
51 class AutoBatchedPickler(i: Iterator[Any]) extends SerDeUtil.AutoBatchedPickler(i)