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.locationtech.geomesa.fs.tools.compact
10 
11 import com.beust.jcommander.{Parameter, ParameterException, Parameters}
12 import com.typesafe.scalalogging.LazyLogging
13 import org.apache.hadoop.fs.Path
14 import org.locationtech.geomesa.fs.data.FileSystemDataStore
15 import org.locationtech.geomesa.fs.storage.orc.OrcFileSystemStorage
16 import org.locationtech.geomesa.fs.storage.parquet.ParquetFileSystemStorage
17 import org.locationtech.geomesa.fs.tools.FsDataStoreCommand
18 import org.locationtech.geomesa.fs.tools.FsDataStoreCommand.{FsDistributedCommand, FsParams, PartitionParam}
19 import org.locationtech.geomesa.fs.tools.compact.FileSystemCompactionJob.{OrcCompactionJob, ParquetCompactionJob}
20 import org.locationtech.geomesa.fs.tools.compact.FsCompactCommand.CompactCommand
21 import org.locationtech.geomesa.jobs.JobResult.{JobFailure, JobSuccess}
22 import org.locationtech.geomesa.tools.Command.CommandException
23 import org.locationtech.geomesa.tools.DistributedRunParam.RunModes
24 import org.locationtech.geomesa.tools._
25 import org.locationtech.geomesa.tools.ingest.IngestCommand
26 import org.locationtech.geomesa.tools.utils.ParameterConverters.BytesConverter
27 import org.locationtech.geomesa.tools.utils.TerminalCallback.PrintProgress
28 import org.locationtech.geomesa.utils.io.PathUtils
29 import org.locationtech.geomesa.utils.text.TextTools
30 
31 import java.util.Locale
32 import java.util.concurrent.{CountDownLatch, Executors}
33 import scala.util.control.NonFatal
34 
35 // need to mixin FsDistributedCommand to pick up base libjars file
36 class FsCompactCommand extends CompactCommand with FsDistributedCommand
37 
38 object FsCompactCommand {
39 
40   trait CompactCommand extends FsDataStoreCommand with DistributedCommand with LazyLogging {
41 
42     import scala.collection.JavaConverters._
43 
44     override val name: String = "compact"
45     override val params = new CompactParams
46 
47     override def libjarsFiles: Seq[String] = Seq("org/locationtech/geomesa/tools/ingest-libjars.list")
48 
49     override def execute(): Unit = withDataStore(compact)
50 
51     def compact(ds: FileSystemDataStore): Unit = {
52       Command.user.info("Beginning compaction process...")
53 
54       val storage = ds.storage(params.featureName)
55 
56       val toCompact = if (params.partitions.isEmpty) { storage.getPartitions } else {
57         val filtered = params.partitions.asScala.flatMap(storage.metadata.getPartition)
58         if (filtered.lengthCompare(params.partitions.size()) != 0) {
59           val unmatched = params.partitions.asScala.filterNot(name => filtered.exists(_.name == name))
60           throw new ParameterException(s"Partition(s) ${unmatched.mkString(", ")} cannot be found in metadata")
61         }
62         filtered
63       }
64 
65       val mode = params.mode.getOrElse {
66         if (PathUtils.isRemote(storage.context.root.toString)) { RunModes.Distributed } else { RunModes.Local }
67       }
68       val fileSize = Option(params.targetFileSize).map(_.longValue)
69 
70       Command.user.info(s"Compacting ${toCompact.size} partitions in ${mode.toString.toLowerCase(Locale.US)} mode")
71 
72       val start = System.currentTimeMillis()
73       val status = new PrintProgress(System.err, TextTools.buildString(' ', 60), '\u003d', '\u003e', '\u003e')
74 
75       mode match {
76         case RunModes.Local =>
77           val total = toCompact.length
78           val latch = new CountDownLatch(total)
79           val executor = Executors.newFixedThreadPool(math.max(1, math.min(params.threads, total)))
80 
81           try {
82             toCompact.foreach { p =>
83               executor.submit(
84                 new Runnable() {
85                   override def run(): Unit = {
86                     try {
87                       logger.info(s"Compacting ${p.name}")
88                       storage.compact(Some(p.name), fileSize)
89                     } catch {
90                       case NonFatal(e) => logger.error(s"Error processing partition '${p.name}':", e)
91                     } finally {
92                       latch.countDown()
93                     }
94                   }
95                 }
96               )
97             }
98           } finally {
99             executor.shutdown()
100           }
101 
102           while (latch.getCount > 0) {
103             Thread.sleep(1000)
104             status("", 1f - latch.getCount.toFloat / total, Seq.empty, done = false)
105           }
106           status("", 1f, Seq.empty, done = true)
107           Command.user.info("Compacting metadata")
108           storage.metadata.compact(None, None, math.max(1, params.threads))
109           Command.user.info(s"Local compaction complete in ${TextTools.getTime(start)}")
110 
111         case RunModes.Distributed =>
112           val encoding = storage.metadata.encoding
113           val job = if (ParquetFileSystemStorage.Encoding.equalsIgnoreCase(encoding)) {
114             new ParquetCompactionJob()
115           } else if (OrcFileSystemStorage.Encoding.equalsIgnoreCase(encoding)) {
116             new OrcCompactionJob()
117           } else {
118             throw new ParameterException(s"Compaction is not supported for encoding '$encoding'")
119           }
120           val tempDir = Option(params.tempPath).map(t => new Path(t))
121           job.run(storage, toCompact.toSeq, fileSize, tempDir, libjarsFiles, libjarsPaths, status) match {
122             case JobSuccess(message, counts) =>
123               Command.user.info(s"Distributed compaction complete in ${TextTools.getTime(start)}")
124               val success = counts(FileSystemCompactionJob.MappedCounter)
125               val failed = counts(FileSystemCompactionJob.FailedCounter)
126               Command.user.info(IngestCommand.getStatInfo(success, failed, "Compacted", message))
127 
128             case JobFailure(message) =>
129               Command.user.error(s"Distributed compaction failed in ${TextTools.getTime(start)}")
130               throw new CommandException(message)
131           }
132       }
133     }
134   }
135 
136   @Parameters(commandDescription = "Compact partitions")
137   class CompactParams extends FsParams
138       with RequiredTypeNameParam with TempPathParam with PartitionParam with DistributedRunParam {
139 
140     @Parameter(names = Array("-t", "--threads"), description = "Number of threads if using local mode")
141     var threads: Integer = 4
142 
143     @Parameter(
144       names = Array("--target-file-size"),
145       description = "Target size for data files",
146       converter = classOf[BytesConverter])
147     var targetFileSize: java.lang.Long = _
148   }
149 }