1 /*
2  * Copyright (C) 2020 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.timezone.location.data_pipeline.steps;
18 
19 import com.android.timezone.location.common.LicenseSupport;
20 import com.android.timezone.location.common.LicenseSupport.License;
21 import com.android.timezone.location.data_pipeline.steps.Types.Pair;
22 import com.android.timezone.location.data_pipeline.steps.Types.ProtoStorageFormat;
23 import com.android.timezone.location.data_pipeline.steps.Types.TzS2CellUnion;
24 import com.android.timezone.location.data_pipeline.steps.Types.TzS2Polygons;
25 import com.android.timezone.location.data_pipeline.util.NamedFuture;
26 
27 import com.beust.jcommander.JCommander;
28 import com.beust.jcommander.Parameter;
29 import com.beust.jcommander.converters.FileConverter;
30 import com.google.common.base.Stopwatch;
31 import com.google.common.geometry.S2CellId;
32 import com.google.common.geometry.S2CellUnion;
33 import com.google.common.geometry.S2Polygon;
34 import com.google.common.geometry.S2RegionCoverer;
35 
36 import java.io.File;
37 import java.io.IOException;
38 import java.util.ArrayList;
39 import java.util.HashSet;
40 import java.util.List;
41 import java.util.Objects;
42 import java.util.Set;
43 import java.util.concurrent.ExecutionException;
44 import java.util.concurrent.ExecutorService;
45 import java.util.concurrent.Executors;
46 import java.util.concurrent.Future;
47 import java.util.concurrent.TimeUnit;
48 
49 /**
50  * An executable class that takes a set of {@link TzS2Polygons} files and produces
51  * {@link TzS2CellUnion} files, one per input file. See {@link #main(String[])} for usage.
52  */
53 public final class TzS2PolygonsToTzS2CellUnions {
54 
55     private static final int MAX_CELL_UNION_CELLS = 1000000;
56 
57     private final File mInputDir;
58 
59     private final ExecutorService mExecutorService;
60 
61     private final File mOutputDir;
62 
63     private final int mMaxS2Level;
64 
65     private final ProtoStorageFormat mProtoStorageFormat;
66 
TzS2PolygonsToTzS2CellUnions( File inputDir, ExecutorService executorService, File outputDir, int maxS2Level, ProtoStorageFormat protoStorageFormat)67     private TzS2PolygonsToTzS2CellUnions(
68             File inputDir, ExecutorService executorService, File outputDir, int maxS2Level,
69             ProtoStorageFormat protoStorageFormat) {
70         this.mInputDir = Objects.requireNonNull(inputDir);
71         this.mExecutorService = Objects.requireNonNull(executorService);
72         this.mOutputDir = Objects.requireNonNull(outputDir);
73         if (maxS2Level < 1 || maxS2Level > S2CellId.MAX_LEVEL) {
74             throw new IllegalArgumentException("Bad S2 level=" + maxS2Level);
75         }
76         this.mMaxS2Level = maxS2Level;
77         this.mProtoStorageFormat = Objects.requireNonNull(protoStorageFormat);
78     }
79 
80     private static class Arguments {
81         @Parameter(names = "--input",
82                 description = "The input directory containing the TzS2Polygons files",
83                 required = true,
84                 converter = FileConverter.class)
85         File inputDir;
86 
87         @Parameter(names = "--num-threads",
88                 description = "The number of threads to use",
89                 required = true)
90         int numThreads;
91 
92         @Parameter(names = "--output",
93                 description = "The output directory to store the TzS2CellUnion files",
94                 required = true,
95                 converter = FileConverter.class)
96         File outputDir;
97 
98         @Parameter(names = "--max-s2-level",
99                 description = "The maximum S2 level of the cells to include in the S2 cell unions",
100                 required = true)
101         int maxS2Level;
102 
103     }
104 
105     /**
106      * See {@link TzS2PolygonsToTzS2CellUnions} for the purpose of this class.
107      */
main(String[] args)108     public static void main(String[] args) throws Exception {
109         Arguments arguments = new Arguments();
110         JCommander.newBuilder()
111                 .addObject(arguments)
112                 .build()
113                 .parse(args);
114 
115         File inputDir = arguments.inputDir;
116         int threads = arguments.numThreads;
117         File outputDir = arguments.outputDir;
118         int maxS2Level = arguments.maxS2Level;
119         ProtoStorageFormat protoStorageFormat = Types.DEFAULT_PROTO_STORAGE_FORMAT;
120 
121         outputDir.mkdirs();
122 
123         ExecutorService executorService = Executors.newFixedThreadPool(threads);
124         TzS2PolygonsToTzS2CellUnions converter =
125                 new TzS2PolygonsToTzS2CellUnions(inputDir, executorService, outputDir, maxS2Level,
126                         protoStorageFormat);
127 
128         try {
129             converter.execute();
130         } finally {
131             System.out.println("Waiting for shutdown");
132             executorService.shutdown();
133             executorService.awaitTermination(5, TimeUnit.SECONDS);
134         }
135     }
136 
execute()137     private void execute() throws Exception {
138         LicenseSupport.copyLicenseFile(mInputDir, mOutputDir);
139 
140         List<File> inputFiles = TzS2Polygons.listFiles(mInputDir, mProtoStorageFormat);
141         List<NamedFuture<Pair<String, File>>> futures = new ArrayList<>();
142         for (File file : inputFiles) {
143             Future<Pair<String, File>> future =
144                     mExecutorService.submit(() -> processFile(file, mMaxS2Level, mOutputDir));
145             NamedFuture<Pair<String, File>> namedFuture = new NamedFuture<>(file.getName(), future);
146             futures.add(namedFuture);
147         }
148 
149         Set<String> knownTzIds = new HashSet<>();
150         for (NamedFuture<Pair<String, File>> future : futures) {
151             try {
152                 if (!future.isDone()) {
153                     System.out.println("Waiting for " + future.getName());
154                 }
155                 Pair<String, File> output = future.get();
156                 String tzId = output.a;
157                 if (knownTzIds.contains(tzId)) {
158                     throw new IllegalStateException("Multiple entries found for: " + tzId);
159                 }
160                 knownTzIds.add(tzId);
161             } catch (InterruptedException | ExecutionException e) {
162                 throw new RuntimeException(e);
163             }
164         }
165     }
166 
processFile( File tzS2PolygonFile, int maxS2Level, File outputDir)167     private Pair<String, File> processFile(
168             File tzS2PolygonFile, int maxS2Level, File outputDir) throws IOException {
169         TzS2Polygons tzS2Polygons = TzS2Polygons.load(tzS2PolygonFile, mProtoStorageFormat);
170         TzS2CellUnion tzS2CellUnion = createTzS2CellUnion(tzS2Polygons, maxS2Level);
171         File outputFile = TzIds.createFile(
172                 outputDir, tzS2CellUnion.tzId, TzS2CellUnion.getFileSuffix(mProtoStorageFormat));
173         TzS2CellUnion.store(tzS2CellUnion, outputFile, mProtoStorageFormat, License.ODBL);
174         return new Pair<>(tzS2Polygons.tzId, outputFile);
175     }
176 
createTzS2CellUnion(TzS2Polygons tzPolygons, int maxS2Level)177     private static TzS2CellUnion createTzS2CellUnion(TzS2Polygons tzPolygons, int maxS2Level) {
178         Stopwatch stopwatch = Stopwatch.createStarted();
179 
180         S2RegionCoverer s2RegionCovererQuad = new S2RegionCoverer();
181         s2RegionCovererQuad.setMinLevel(1);
182         s2RegionCovererQuad.setMaxLevel(maxS2Level);
183         s2RegionCovererQuad.setMaxCells(MAX_CELL_UNION_CELLS);
184 
185         List<S2Polygon> s2Polygons = tzPolygons.s2PolygonList;
186         String tzId = tzPolygons.tzId;
187         ArrayList<S2CellId> cellIds = new ArrayList<>();
188         for (S2Polygon s2Polygon : s2Polygons) {
189             S2CellUnion covering = s2RegionCovererQuad.getCovering(s2Polygon);
190             cellIds.addAll(covering.cellIds());
191         }
192         if (cellIds.size() >= (MAX_CELL_UNION_CELLS * 95) / 100) {
193             System.err.println(
194                     "Possible overflow. for " + tzPolygons.tzId + "size=" + cellIds.size() + "...");
195             throw new IllegalStateException();
196         }
197         S2CellUnion combinedCellUnion = new S2CellUnion();
198         combinedCellUnion.initFromCellIds(cellIds);
199         System.out.printf("Created S2CellUnion for %s containing %s cells at level %S in %s...\n",
200                 tzId, cellIds.size(), maxS2Level, stopwatch.elapsed());
201         return new TzS2CellUnion(tzId, combinedCellUnion);
202     }
203 }
204