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.fs.storage.orc.jobs
10 
11 import org.apache.hadoop.conf.Configuration
12 import org.apache.hadoop.fs.Path
13 import org.apache.hadoop.io.NullWritable
14 import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat
15 import org.apache.hadoop.mapreduce.{RecordWriter, TaskAttemptContext}
16 import org.apache.orc.mapred.OrcStruct
17 import org.apache.orc.mapreduce.{OrcMapreduceRecordWriter, OrcOutputFormat}
18 import org.apache.orc.{OrcConf, OrcFile, TypeDescription}
19 import org.geotools.api.feature.simple.{SimpleFeature, SimpleFeatureType}
20 import org.locationtech.geomesa.fs.storage.common.jobs.StorageConfiguration
21 import org.locationtech.geomesa.fs.storage.orc.jobs.OrcSimpleFeatureOutputFormat.OrcRecordWriter
22 import org.locationtech.geomesa.fs.storage.orc.utils.OrcOutputFormatWriter
23 
24 class OrcSimpleFeatureOutputFormat extends FileOutputFormat[Void, SimpleFeature] {
25 
26   private val delegate = new OrcOutputFormat[OrcStruct]
27 
28   def getRecordWriter(context: TaskAttemptContext, file: Path): RecordWriter[Void, SimpleFeature] = {
29     val options = org.apache.orc.mapred.OrcOutputFormat.buildOptions(context.getConfiguration)
30     val writer = new OrcMapreduceRecordWriter[OrcStruct](OrcFile.createWriter(file, options))
31     getRecordWriter(context, writer)
32   }
33 
34   override def getRecordWriter(context: TaskAttemptContext): RecordWriter[Void, SimpleFeature] =
35     getRecordWriter(context, delegate.getRecordWriter(context))
36 
37   override def getDefaultWorkFile(context: TaskAttemptContext, extension: String): Path =
38     delegate.getDefaultWorkFile(context, extension)
39 
40   private def getRecordWriter(context: TaskAttemptContext,
41                               orcWriter: RecordWriter[NullWritable, OrcStruct]): RecordWriter[Void, SimpleFeature] = {
42     val sft = StorageConfiguration.getSft(context.getConfiguration)
43     val description = OrcSimpleFeatureOutputFormat.getDescription(context.getConfiguration)
44     new OrcRecordWriter(sft, description, orcWriter)
45   }
46 }
47 
48 object OrcSimpleFeatureOutputFormat {
49 
50   def setDescription(conf: Configuration, description: TypeDescription): Unit =
51     conf.set(OrcConf.MAPRED_OUTPUT_SCHEMA.getAttribute, description.toString)
52 
53   def getDescription(conf: Configuration): TypeDescription =
54     TypeDescription.fromString(OrcConf.MAPRED_OUTPUT_SCHEMA.getString(conf))
55 
56   class OrcRecordWriter(sft: SimpleFeatureType,
57                         description: TypeDescription,
58                         delegate: RecordWriter[NullWritable, OrcStruct])
59       extends RecordWriter[Void, SimpleFeature] {
60 
61     private val writer = OrcOutputFormatWriter(sft, description)
62     private val key = NullWritable.get()
63     private val struct = new OrcStruct(description)
64 
65     override def write(key: Void, value: SimpleFeature): Unit = {
66       writer.apply(value, struct)
67       delegate.write(this.key, struct)
68     }
69 
70     override def close(context: TaskAttemptContext): Unit = delegate.close(context)
71   }
72 }