1 /***********************************************************************
2  * Copyright (c) 2013-2025 General Atomics Integrated Intelligence, 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  * https://www.apache.org/licenses/LICENSE-2.0
7  ***********************************************************************/
8 
9 package org.locationtech.geomesa.features.exporters
10 
11 import org.geotools.api.data.Transaction
12 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
13 import org.geotools.data.shapefile.files.ShpFiles
14 import org.geotools.data.shapefile.{ShapefileDataStore, ShapefileDataStoreFactory}
15 import org.geotools.util.URLs
16 import org.locationtech.geomesa.utils.io.WithClose
17 import org.locationtech.jts.geom.Geometry
18 
19 import java.io.File
20 import java.net.URL
21 
22 /**
23  * Exports to a shapefile, which is a collection of files
24  *
25  * @param file pointer to a .shp file (not a directory)
26  */
27 class ShapefileExporter(file: File) extends FeatureExporter {
28 
29   import scala.collection.JavaConverters._
30 
31   private val url = URLs.fileToUrl(file)
32   private var ds: ShapefileDataStore = _
33   private var mappings: Map[Int, Int] = _
34 
35   override def start(sft: SimpleFeatureType): Unit = {
36     // ensure that the parent directory exists, otherwise the data store will error out
37     Option(file.getParentFile).filterNot(_.exists()).foreach(_.mkdirs())
38     ds = new ShapefileDataStoreFactory().createDataStore(url).asInstanceOf[ShapefileDataStore]
39     ds.createSchema(sft)
40 
41     var i = -1
42     var j = 0
43     // map attributes according to the shapefile data store
44     // default geometry goes to attribute 0, other geometries are dropped
45     // byte arrays are dropped, but everything else is kept or transformed
46     // see: `org.geotools.data.shapefile.ShapefileDataStore.createDbaseHeader()`
47     val attributes = sft.getAttributeDescriptors.asScala.flatMap { d =>
48       i += 1
49       val binding = d.getType.getBinding
50       if (classOf[Geometry].isAssignableFrom(binding) || binding == classOf[Array[Byte]] ) { None } else {
51         j += 1
52         Some(i -> j)
53       }
54     }
55     mappings = attributes.toMap + (sft.indexOf(sft.getGeometryDescriptor.getLocalName) -> 0)
56   }
57 
58   override def export(features: Iterator[SimpleFeature]): Option[Long] = {
59     var count = 0L
60 
61     WithClose(ds.getFeatureWriterAppend(Transaction.AUTO_COMMIT)) { writer =>
62       features.foreach { feature =>
63         val toWrite = writer.next()
64         mappings.foreach { case (from, to) => toWrite.setAttribute(to, feature.getAttribute(from)) }
65         // copy over the user data
66         // note: shapefile doesn't support provided fid
67         toWrite.getUserData.putAll(feature.getUserData)
68         writer.write()
69         count += 1L
70       }
71     }
72 
73     Some(count)
74   }
75 
76   /**
77    * Counts the total bytes written, across the various files in this shapefile
78    *
79    * @return bytes written
80    */
81   def bytes: Long = {
82     val files = new ShpFiles(url)
83     try {
84       var sum = 0L
85       files.getFileNames.asScala.values.foreach { file =>
86         sum += URLs.urlToFile(new URL(file)).length()
87       }
88       sum
89     } finally {
90       files.dispose()
91     }
92   }
93 
94   override def close(): Unit = Option(ds).foreach(_.dispose)
95 }