Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: GDAL
4 : * Purpose: "partition" step of "vector pipeline"
5 : * Author: Even Rouault <even dot rouault at spatialys.com>
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2025, Even Rouault <even dot rouault at spatialys.com>
9 : *
10 : * SPDX-License-Identifier: MIT
11 : ****************************************************************************/
12 :
13 : #include "gdalalg_vector_partition.h"
14 :
15 : #include "cpl_vsi.h"
16 : #include "cpl_mem_cache.h"
17 : #include "ogr_p.h"
18 :
19 : #include <algorithm>
20 : #include <set>
21 : #include <string_view>
22 :
23 : #ifndef _
24 : #define _(x) (x)
25 : #endif
26 :
27 : //! @cond Doxygen_Suppress
28 :
29 : constexpr int DIRECTORY_CREATION_MODE = 0755;
30 :
31 : constexpr const char *NULL_MARKER = "__HIVE_DEFAULT_PARTITION__";
32 :
33 : constexpr const char *DEFAULT_PATTERN_HIVE = "part_%010d";
34 : constexpr const char *DEFAULT_PATTERN_FLAT_NO_FIELD = "{LAYER_NAME}_%010d";
35 : constexpr const char *DEFAULT_PATTERN_FLAT = "{LAYER_NAME}_{FIELD_VALUE}_%010d";
36 :
37 : constexpr char DIGIT_ZERO = '0';
38 :
39 : /************************************************************************/
40 : /* GetConstructorOptions() */
41 : /************************************************************************/
42 :
43 : /* static */
44 : GDALVectorPartitionAlgorithm::ConstructorOptions
45 131 : GDALVectorPartitionAlgorithm::GetConstructorOptions(bool standaloneStep)
46 : {
47 131 : GDALVectorPartitionAlgorithm::ConstructorOptions options;
48 131 : options.SetStandaloneStep(standaloneStep);
49 131 : options.SetAddInputLayerNameArgument(false);
50 131 : options.SetAddDefaultArguments(false);
51 131 : return options;
52 : }
53 :
54 : /************************************************************************/
55 : /* GDALVectorPartitionAlgorithm::GDALVectorPartitionAlgorithm() */
56 : /************************************************************************/
57 :
58 131 : GDALVectorPartitionAlgorithm::GDALVectorPartitionAlgorithm(bool standaloneStep)
59 : : GDALVectorPipelineStepAlgorithm(NAME, DESCRIPTION, HELP_URL,
60 131 : GetConstructorOptions(standaloneStep))
61 : {
62 131 : if (standaloneStep)
63 : {
64 93 : AddVectorInputArgs(false);
65 93 : AddProgressArg();
66 : }
67 : else
68 : {
69 38 : AddVectorHiddenInputDatasetArg();
70 : }
71 :
72 262 : AddArg(GDAL_ARG_NAME_OUTPUT, 'o', _("Output directory"), &m_output)
73 131 : .SetRequired()
74 131 : .SetIsInput()
75 131 : .SetMinCharCount(1)
76 131 : .SetPositional();
77 :
78 131 : constexpr const char *OVERWRITE_APPEND_EXCLUSION_GROUP = "overwrite-append";
79 131 : AddOverwriteArg(&m_overwrite)
80 131 : .SetMutualExclusionGroup(OVERWRITE_APPEND_EXCLUSION_GROUP);
81 131 : AddAppendLayerArg(&m_appendLayer)
82 131 : .SetMutualExclusionGroup(OVERWRITE_APPEND_EXCLUSION_GROUP);
83 131 : AddUpdateArg(&m_update).SetHidden();
84 :
85 : AddOutputFormatArg(&m_format, /* bStreamAllowed = */ false,
86 131 : /* bGDALGAllowed = */ false)
87 : .AddMetadataItem(GAAMDI_REQUIRED_CAPABILITIES,
88 393 : {GDAL_DCAP_VECTOR, GDAL_DCAP_CREATE});
89 131 : AddCreationOptionsArg(&m_creationOptions);
90 131 : AddLayerCreationOptionsArg(&m_layerCreationOptions);
91 :
92 : auto &fieldNameArg = AddArg(
93 : "field", 0, _("Attribute or geometry field(s) on which to partition"),
94 131 : &m_fields);
95 262 : SetAutoCompleteFunctionForFieldName(fieldNameArg, nullptr,
96 : /* attributeFields = */ true,
97 : /* geometryFields = */ true,
98 131 : m_inputDataset);
99 262 : AddArg("scheme", 0, _("Partitioning scheme"), &m_scheme)
100 131 : .SetChoices(SCHEME_HIVE, SCHEME_FLAT)
101 131 : .SetDefault(m_scheme);
102 : AddArg("pattern", 0,
103 : _("Filename pattern ('part_%010d' for scheme=hive, "
104 : "'{LAYER_NAME}_{FIELD_VALUE}_%010d' for scheme=flat)"),
105 262 : &m_pattern)
106 131 : .SetMinCharCount(1)
107 : .AddValidationAction(
108 72 : [this]()
109 : {
110 8 : if (!m_pattern.empty())
111 : {
112 8 : const auto nPercentPos = m_pattern.find('%');
113 8 : if (nPercentPos == std::string::npos)
114 : {
115 1 : ReportError(CE_Failure, CPLE_IllegalArg, "%s",
116 : "Missing '%' character in pattern");
117 1 : return false;
118 : }
119 13 : if (nPercentPos + 1 < m_pattern.size() &&
120 6 : m_pattern.find('%', nPercentPos + 1) !=
121 : std::string::npos)
122 : {
123 1 : ReportError(
124 : CE_Failure, CPLE_IllegalArg, "%s",
125 : "A single '%' character is expected in pattern");
126 1 : return false;
127 : }
128 6 : bool percentFound = false;
129 9 : for (size_t i = nPercentPos + 1; i < m_pattern.size(); ++i)
130 : {
131 7 : if (m_pattern[i] >= DIGIT_ZERO && m_pattern[i] <= '9')
132 : {
133 : // ok
134 : }
135 4 : else if (m_pattern[i] == 'd')
136 : {
137 3 : percentFound = true;
138 3 : break;
139 : }
140 : else
141 : {
142 1 : break;
143 : }
144 : }
145 6 : if (!percentFound)
146 : {
147 3 : ReportError(
148 : CE_Failure, CPLE_IllegalArg, "%s",
149 : "pattern value must include a single "
150 : "'%[0]?[1-9]?[0]?d' part number specification");
151 3 : return false;
152 : }
153 3 : m_partDigitCount =
154 3 : atoi(m_pattern.c_str() + nPercentPos + 1);
155 3 : if (m_partDigitCount > 10)
156 : {
157 1 : ReportError(CE_Failure, CPLE_IllegalArg,
158 : "Number of digits in part number "
159 : "specification should be in [1,10] range");
160 1 : return false;
161 : }
162 2 : m_partDigitLeadingZeroes =
163 2 : m_pattern[nPercentPos + 1] == DIGIT_ZERO;
164 : }
165 2 : return true;
166 131 : });
167 : AddArg("feature-limit", 0, _("Maximum number of features per file"),
168 262 : &m_featureLimit)
169 131 : .SetMinValueExcluded(0);
170 : AddArg("max-file-size", 0,
171 : _("Maximum file size (MB or GB suffix can be used)"),
172 262 : &m_maxFileSizeStr)
173 : .AddValidationAction(
174 24 : [this]()
175 : {
176 : bool ok;
177 : {
178 6 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
179 12 : ok = CPLParseMemorySize(m_maxFileSizeStr.c_str(),
180 : &m_maxFileSize,
181 11 : nullptr) == CE_None &&
182 5 : m_maxFileSize > 0;
183 : }
184 6 : if (!ok)
185 : {
186 1 : ReportError(CE_Failure, CPLE_IllegalArg,
187 : "Invalid value for max-file-size");
188 1 : return false;
189 : }
190 5 : else if (m_maxFileSize < 1024 * 1024)
191 : {
192 1 : ReportError(CE_Failure, CPLE_IllegalArg,
193 : "max-file-size should be at least one MB");
194 1 : return false;
195 : }
196 4 : return true;
197 131 : });
198 : AddArg("omit-partitioned-field", 0,
199 : _("Whether to omit partitioned fields from target layer definition"),
200 131 : &m_omitPartitionedFields);
201 : AddArg("skip-errors", 0, _("Skip errors when writing features"),
202 131 : &m_skipErrors);
203 :
204 : // Hidden for now
205 :
206 : AddArg("max-cache-size", 0,
207 : _("Maximum number of datasets simultaneously opened"),
208 262 : &m_maxCacheSize)
209 131 : .SetMinValueIncluded(0) // 0 = unlimited
210 131 : .SetDefault(m_maxCacheSize)
211 131 : .SetHidden();
212 :
213 : AddArg("transaction-size", 0,
214 262 : _("Maximum number of features per transaction"), &m_transactionSize)
215 131 : .SetMinValueIncluded(1)
216 131 : .SetDefault(m_transactionSize)
217 131 : .SetHidden();
218 :
219 131 : AddValidationAction(
220 53 : [this]()
221 : {
222 48 : if (m_fields.empty() && m_featureLimit == 0 && m_maxFileSize == 0)
223 : {
224 1 : ReportError(
225 : CE_Failure, CPLE_IllegalArg,
226 : "When 'fields' argument is not specified, "
227 : "'feature-limit' and/or 'max-file-size' must be specified");
228 1 : return false;
229 : }
230 47 : return true;
231 : });
232 131 : }
233 :
234 : /************************************************************************/
235 : /* PercentEncode() */
236 : /************************************************************************/
237 :
238 20305 : static void PercentEncode(std::string &out, const std::string_view &s)
239 : {
240 92085 : for (unsigned char c : s)
241 : {
242 71780 : if (c > 32 && c <= 127 && c != ':' && c != '/' && c != '\\' &&
243 71733 : c != '>' && c != '%' && c != '=')
244 : {
245 71733 : out += c;
246 : }
247 : else
248 : {
249 47 : out += CPLSPrintf("%%%02X", c);
250 : }
251 : }
252 20305 : }
253 :
254 10237 : static std::string PercentEncode(const std::string_view &s)
255 : {
256 10237 : std::string out;
257 10237 : PercentEncode(out, s);
258 10237 : return out;
259 : }
260 :
261 : /************************************************************************/
262 : /* GetEstimatedFeatureSize() */
263 : /************************************************************************/
264 :
265 10000 : static size_t GetEstimatedFeatureSize(
266 : const OGRFeature *poFeature, const std::vector<bool> &abPartitionedFields,
267 : const bool omitPartitionedFields,
268 : const std::vector<OGRFieldType> &aeSrcFieldTypes, bool bIsBinary)
269 : {
270 10000 : size_t nSize = 16;
271 10000 : const int nFieldCount = poFeature->GetFieldCount();
272 10000 : nSize += 4 * nFieldCount;
273 110000 : for (int i = 0; i < nFieldCount; ++i)
274 : {
275 100000 : if (!(omitPartitionedFields && abPartitionedFields[i]))
276 : {
277 100000 : switch (aeSrcFieldTypes[i])
278 : {
279 10000 : case OFTInteger:
280 10000 : nSize += bIsBinary ? sizeof(int) : 11;
281 10000 : break;
282 10000 : case OFTInteger64:
283 10000 : nSize += bIsBinary ? sizeof(int64_t) : 21;
284 10000 : break;
285 10000 : case OFTReal:
286 : // Decimal representation
287 10000 : nSize += bIsBinary ? sizeof(double) : 15;
288 10000 : break;
289 10000 : case OFTString:
290 10000 : nSize += 4 + strlen(poFeature->GetFieldAsStringUnsafe(i));
291 10000 : break;
292 10000 : case OFTBinary:
293 : {
294 10000 : int nCount = 0;
295 10000 : CPL_IGNORE_RET_VAL(poFeature->GetFieldAsBinary(i, &nCount));
296 10000 : nSize += 4 + nCount;
297 10000 : break;
298 : }
299 5000 : case OFTIntegerList:
300 : {
301 5000 : int nCount = 0;
302 5000 : CPL_IGNORE_RET_VAL(
303 5000 : poFeature->GetFieldAsIntegerList(i, &nCount));
304 5000 : nSize += 4 + (bIsBinary ? sizeof(int) : 11) * nCount;
305 5000 : break;
306 : }
307 5000 : case OFTInteger64List:
308 : {
309 5000 : int nCount = 0;
310 5000 : CPL_IGNORE_RET_VAL(
311 5000 : poFeature->GetFieldAsInteger64List(i, &nCount));
312 5000 : nSize += 4 + (bIsBinary ? sizeof(int64_t) : 21) * nCount;
313 5000 : break;
314 : }
315 5000 : case OFTRealList:
316 : {
317 5000 : int nCount = 0;
318 5000 : CPL_IGNORE_RET_VAL(
319 5000 : poFeature->GetFieldAsDoubleList(i, &nCount));
320 5000 : nSize += 4 + (bIsBinary ? sizeof(double) : 15) * nCount;
321 5000 : break;
322 : }
323 5000 : case OFTStringList:
324 : {
325 5000 : CSLConstList papszIter = poFeature->GetFieldAsStringList(i);
326 5000 : nSize += 4;
327 15000 : for (; papszIter && *papszIter; ++papszIter)
328 10000 : nSize += 4 + strlen(*papszIter);
329 5000 : break;
330 : }
331 10000 : case OFTTime:
332 : // Decimal representation
333 10000 : nSize += 4 + sizeof("HH:MM:SS.sss");
334 10000 : break;
335 10000 : case OFTDate:
336 : // Decimal representation
337 10000 : nSize += 4 + sizeof("YYYY-MM-DD");
338 10000 : break;
339 10000 : case OFTDateTime:
340 : // Decimal representation
341 10000 : nSize += 4 + sizeof("YYYY-MM-DDTHH:MM:SS.sss+HH:MM");
342 10000 : break;
343 0 : case OFTWideString:
344 : case OFTWideStringList:
345 0 : break;
346 : }
347 : }
348 : }
349 :
350 10000 : const int nGeomFieldCount = poFeature->GetGeomFieldCount();
351 10000 : nSize += 4 * nGeomFieldCount;
352 20000 : for (int i = 0; i < nGeomFieldCount; ++i)
353 : {
354 10000 : const auto poGeom = poFeature->GetGeomFieldRef(i);
355 10000 : if (poGeom)
356 10000 : nSize += poGeom->WkbSize();
357 : }
358 :
359 10000 : return nSize;
360 : }
361 :
362 : /************************************************************************/
363 : /* GetCurrentOutputLayer() */
364 : /************************************************************************/
365 :
366 : constexpr int MIN_FILE_SIZE = 65536;
367 :
368 : namespace
369 : {
370 : struct Layer
371 : {
372 : bool bUseTransactions = false;
373 : std::unique_ptr<GDALDataset> poDS{};
374 : OGRLayer *poLayer = nullptr;
375 : GIntBig nFeatureCount = 0;
376 : int nFileCounter = 1;
377 : GIntBig nFileSize = MIN_FILE_SIZE;
378 :
379 151 : ~Layer()
380 151 : {
381 151 : if (poDS)
382 : {
383 87 : CPL_IGNORE_RET_VAL(poDS->CommitTransaction());
384 : }
385 151 : }
386 : };
387 : } // namespace
388 :
389 10115 : static bool GetCurrentOutputLayer(
390 : GDALAlgorithm *const alg, const OGRFeatureDefn *const poSrcFeatureDefn,
391 : OGRLayer *const poSrcLayer, const std::string &osKey,
392 : const std::vector<OGRwkbGeometryType> &aeGeomTypes,
393 : const std::string &osLayerDir, const std::string &osScheme,
394 : const std::string &osPatternIn, bool partDigitLeadingZeroes,
395 : size_t partDigitCount, const int featureLimit, const GIntBig maxFileSize,
396 : const bool omitPartitionedFields,
397 : const std::vector<bool> &abPartitionedFields,
398 : const std::vector<bool> &abPartitionedGeomFields, const char *pszExtension,
399 : GDALDriver *const poOutDriver, const CPLStringList &datasetCreationOptions,
400 : const CPLStringList &layerCreationOptions,
401 : const OGRFeatureDefn *const poFeatureDefnWithoutPartitionedFields,
402 : const int nSpatialIndexPerFeatureConstant,
403 : const int nSpatialIndexPerLog2FeatureCountConstant, bool bUseTransactions,
404 : lru11::Cache<std::string, std::shared_ptr<Layer>> &oCacheOutputLayer,
405 : std::shared_ptr<Layer> &outputLayer)
406 : {
407 : const std::string osPattern =
408 10115 : !osPatternIn.empty() ? osPatternIn
409 10111 : : osScheme == GDALVectorPartitionAlgorithm::SCHEME_HIVE
410 : ? DEFAULT_PATTERN_HIVE
411 8 : : osKey.empty() ? DEFAULT_PATTERN_FLAT_NO_FIELD
412 20238 : : DEFAULT_PATTERN_FLAT;
413 :
414 10115 : bool bLimitReached = false;
415 10115 : bool bOpenOrCreateNewFile = true;
416 10115 : if (oCacheOutputLayer.tryGet(osKey, outputLayer))
417 : {
418 10025 : if (featureLimit > 0 && outputLayer->nFeatureCount >= featureLimit)
419 : {
420 4 : bLimitReached = true;
421 : }
422 20019 : else if (maxFileSize > 0 &&
423 19996 : outputLayer->nFileSize +
424 : (nSpatialIndexPerFeatureConstant > 0
425 9998 : ? (outputLayer->nFeatureCount *
426 9998 : nSpatialIndexPerFeatureConstant +
427 4999 : static_cast<int>(std::ceil(
428 4999 : log2(outputLayer->nFeatureCount)))) *
429 4999 : nSpatialIndexPerLog2FeatureCountConstant
430 : : 0) >=
431 : maxFileSize)
432 : {
433 2 : bLimitReached = true;
434 : }
435 : else
436 : {
437 10019 : bOpenOrCreateNewFile = false;
438 : }
439 : }
440 : else
441 : {
442 90 : outputLayer = std::make_unique<Layer>();
443 90 : outputLayer->bUseTransactions = bUseTransactions;
444 : }
445 :
446 20246 : const auto SubstituteVariables = [&osKey, poSrcLayer](const std::string &s)
447 : {
448 10119 : CPLString ret(s);
449 : ret.replaceAll("{LAYER_NAME}",
450 10119 : PercentEncode(poSrcLayer->GetDescription()));
451 :
452 10119 : if (ret.find("{FIELD_VALUE}") != std::string::npos)
453 : {
454 16 : std::string fieldValue;
455 : const CPLStringList aosTokens(
456 8 : CSLTokenizeString2(osKey.c_str(), "/", 0));
457 16 : for (int i = 0; i < aosTokens.size(); ++i)
458 : {
459 : const CPLStringList aosFieldNameValue(
460 8 : CSLTokenizeString2(aosTokens[i], "=", 0));
461 8 : if (!fieldValue.empty())
462 0 : fieldValue += '_';
463 : fieldValue +=
464 8 : aosFieldNameValue.size() == 2
465 16 : ? (strcmp(aosFieldNameValue[1], NULL_MARKER) == 0
466 : ? std::string("__NULL__")
467 : : aosFieldNameValue[1])
468 8 : : std::string("__EMPTY__");
469 : }
470 8 : ret.replaceAll("{FIELD_VALUE}", fieldValue);
471 : }
472 10119 : return ret;
473 10115 : };
474 :
475 10115 : const auto nPercentPos = osPattern.find('%');
476 10115 : CPLAssert(nPercentPos !=
477 : std::string::npos); // checked by validation action
478 : const std::string osPatternPrefix =
479 30345 : SubstituteVariables(osPattern.substr(0, nPercentPos));
480 10115 : const auto nAfterDPos = osPattern.find('d', nPercentPos + 1) + 1;
481 : const std::string osPatternSuffix =
482 10115 : nAfterDPos < osPattern.size()
483 10123 : ? SubstituteVariables(osPattern.substr(nAfterDPos))
484 20234 : : std::string();
485 :
486 99 : const auto GetBasenameFromCounter = [partDigitCount, partDigitLeadingZeroes,
487 : &osPatternPrefix,
488 489 : &osPatternSuffix](int nCounter)
489 : {
490 198 : const std::string sCounter(CPLSPrintf("%d", nCounter));
491 99 : std::string s(osPatternPrefix);
492 99 : if (sCounter.size() < partDigitCount)
493 : {
494 192 : s += std::string(partDigitCount - sCounter.size(),
495 96 : partDigitLeadingZeroes ? DIGIT_ZERO : ' ');
496 : }
497 99 : s += sCounter;
498 99 : s += osPatternSuffix;
499 198 : return s;
500 10115 : };
501 :
502 10115 : if (bOpenOrCreateNewFile)
503 : {
504 : std::string osDatasetDir =
505 96 : osScheme == GDALVectorPartitionAlgorithm::SCHEME_HIVE
506 : ? CPLFormFilenameSafe(osLayerDir.c_str(), osKey.c_str(),
507 : nullptr)
508 96 : : osLayerDir;
509 96 : outputLayer->nFeatureCount = 0;
510 :
511 96 : bool bCreateNewFile = true;
512 96 : if (bLimitReached)
513 : {
514 6 : ++outputLayer->nFileCounter;
515 : }
516 : else
517 : {
518 90 : outputLayer->nFileCounter = 1;
519 :
520 : VSIStatBufL sStat;
521 90 : if (VSIStatL(osDatasetDir.c_str(), &sStat) != 0)
522 : {
523 71 : if (VSIMkdirRecursive(osDatasetDir.c_str(),
524 71 : DIRECTORY_CREATION_MODE) != 0)
525 : {
526 0 : alg->ReportError(CE_Failure, CPLE_AppDefined,
527 : "Cannot create directory '%s'",
528 : osDatasetDir.c_str());
529 3 : return false;
530 : }
531 : }
532 :
533 90 : int nMaxCounter = 0;
534 : std::unique_ptr<VSIDIR, decltype(&VSICloseDir)> psDir(
535 90 : VSIOpenDir(osDatasetDir.c_str(), 0, nullptr), VSICloseDir);
536 90 : if (psDir)
537 : {
538 110 : while (const auto *psEntry = VSIGetNextDirEntry(psDir.get()))
539 : {
540 : const std::string osName(
541 40 : CPLGetBasenameSafe(psEntry->pszName));
542 32 : if (cpl::starts_with(osName, osPatternPrefix) &&
543 12 : cpl::ends_with(osName, osPatternSuffix))
544 : {
545 10 : nMaxCounter = std::max(
546 : nMaxCounter,
547 10 : atoi(osName
548 20 : .substr(osPatternPrefix.size(),
549 10 : osName.size() -
550 10 : osPatternPrefix.size() -
551 10 : osPatternSuffix.size())
552 10 : .c_str()));
553 : }
554 20 : }
555 : }
556 :
557 90 : if (nMaxCounter > 0)
558 : {
559 9 : outputLayer->nFileCounter = nMaxCounter;
560 :
561 : const std::string osFilename = CPLFormFilenameSafe(
562 : osDatasetDir.c_str(),
563 9 : GetBasenameFromCounter(nMaxCounter).c_str(), pszExtension);
564 : auto poDS = std::unique_ptr<GDALDataset>(GDALDataset::Open(
565 : osFilename.c_str(),
566 9 : GDAL_OF_VECTOR | GDAL_OF_UPDATE | GDAL_OF_VERBOSE_ERROR));
567 9 : if (!poDS)
568 1 : return false;
569 8 : auto poDstLayer = poDS->GetLayer(0);
570 8 : if (!poDstLayer)
571 : {
572 1 : alg->ReportError(CE_Failure, CPLE_AppDefined,
573 : "No layer in %s", osFilename.c_str());
574 1 : return false;
575 : }
576 :
577 : // Check if the existing output layer has the expected layer
578 : // definition
579 7 : const auto poRefFeatureDefn =
580 : poFeatureDefnWithoutPartitionedFields
581 : ? poFeatureDefnWithoutPartitionedFields
582 : : poSrcFeatureDefn;
583 7 : const auto poDstFeatureDefn = poDstLayer->GetLayerDefn();
584 7 : bool bSameDefinition = (poDstFeatureDefn->GetFieldCount() ==
585 7 : poRefFeatureDefn->GetFieldCount());
586 7 : for (int i = 0;
587 31 : bSameDefinition && i < poRefFeatureDefn->GetFieldCount();
588 : ++i)
589 : {
590 : const auto poRefFieldDefn =
591 24 : poRefFeatureDefn->GetFieldDefn(i);
592 : const auto poDstFieldDefn =
593 24 : poDstFeatureDefn->GetFieldDefn(i);
594 24 : bSameDefinition =
595 24 : EQUAL(poRefFieldDefn->GetNameRef(),
596 48 : poDstFieldDefn->GetNameRef()) &&
597 24 : poRefFieldDefn->GetType() == poDstFieldDefn->GetType();
598 : }
599 7 : bSameDefinition =
600 13 : bSameDefinition && (poDstFeatureDefn->GetGeomFieldCount() ==
601 6 : poRefFeatureDefn->GetGeomFieldCount());
602 21 : for (int i = 0; bSameDefinition &&
603 10 : i < poRefFeatureDefn->GetGeomFieldCount();
604 : ++i)
605 : {
606 : const auto poRefFieldDefn =
607 4 : poRefFeatureDefn->GetGeomFieldDefn(i);
608 : const auto poDstFieldDefn =
609 4 : poDstFeatureDefn->GetGeomFieldDefn(i);
610 4 : bSameDefinition =
611 4 : (poRefFeatureDefn->GetGeomFieldCount() == 1 ||
612 0 : EQUAL(poRefFieldDefn->GetNameRef(),
613 : poDstFieldDefn->GetNameRef()));
614 : }
615 :
616 7 : if (!bSameDefinition)
617 : {
618 1 : alg->ReportError(CE_Failure, CPLE_AppDefined,
619 : "%s does not have the same feature "
620 : "definition as the source layer",
621 : osFilename.c_str());
622 1 : return false;
623 : }
624 :
625 6 : if (VSIStatL(osFilename.c_str(), &sStat) == 0)
626 : {
627 6 : outputLayer->nFileSize = sStat.st_size;
628 : }
629 :
630 6 : GIntBig nFeatureCount = 0;
631 9 : if (((featureLimit == 0 ||
632 3 : (nFeatureCount = poDstLayer->GetFeatureCount(true)) <
633 9 : featureLimit)) &&
634 0 : (maxFileSize == 0 || outputLayer->nFileSize < maxFileSize))
635 : {
636 3 : bCreateNewFile = false;
637 3 : outputLayer->poDS = std::move(poDS);
638 3 : outputLayer->poLayer = poDstLayer;
639 3 : outputLayer->nFeatureCount = nFeatureCount;
640 :
641 3 : if (bUseTransactions)
642 : {
643 3 : if (outputLayer->poDS->StartTransaction() !=
644 : OGRERR_NONE)
645 : {
646 0 : return false;
647 : }
648 : }
649 : }
650 : else
651 : {
652 3 : ++outputLayer->nFileCounter;
653 : }
654 : }
655 : }
656 :
657 93 : if (bCreateNewFile)
658 : {
659 90 : outputLayer->nFileSize = MIN_FILE_SIZE;
660 :
661 96 : if (bUseTransactions && outputLayer->poDS &&
662 6 : outputLayer->poDS->CommitTransaction() != OGRERR_NONE)
663 : {
664 3 : return false;
665 : }
666 :
667 : const std::string osFilename = CPLFormFilenameSafe(
668 : osDatasetDir.c_str(),
669 90 : GetBasenameFromCounter(outputLayer->nFileCounter).c_str(),
670 90 : pszExtension);
671 90 : outputLayer->poDS.reset(
672 : poOutDriver->Create(osFilename.c_str(), 0, 0, 0, GDT_Unknown,
673 : datasetCreationOptions.List()));
674 90 : if (!outputLayer->poDS)
675 : {
676 0 : alg->ReportError(CE_Failure, CPLE_AppDefined,
677 : "Cannot create dataset '%s'",
678 : osFilename.c_str());
679 0 : return false;
680 : }
681 :
682 90 : CPLStringList modLayerCreationOptions(layerCreationOptions);
683 90 : const char *pszSrcFIDColumn = poSrcLayer->GetFIDColumn();
684 90 : if (pszSrcFIDColumn[0])
685 : {
686 114 : const char *pszLCO = poOutDriver->GetMetadataItem(
687 57 : GDAL_DS_LAYER_CREATIONOPTIONLIST);
688 104 : if (pszLCO && strstr(pszLCO, "'FID'") &&
689 47 : layerCreationOptions.FetchNameValue("FID") == nullptr)
690 : modLayerCreationOptions.SetNameValue("FID",
691 46 : pszSrcFIDColumn);
692 : }
693 :
694 0 : std::unique_ptr<OGRGeomFieldDefn> poFirstGeomFieldDefn;
695 90 : if (poSrcFeatureDefn->GetGeomFieldCount())
696 : {
697 68 : poFirstGeomFieldDefn = std::make_unique<OGRGeomFieldDefn>(
698 68 : *poSrcFeatureDefn->GetGeomFieldDefn(0));
699 68 : if (abPartitionedGeomFields[0])
700 : {
701 6 : if (aeGeomTypes[0] == wkbNone)
702 1 : poFirstGeomFieldDefn.reset();
703 : else
704 10 : whileUnsealing(poFirstGeomFieldDefn.get())
705 5 : ->SetType(aeGeomTypes[0]);
706 : }
707 : }
708 180 : auto poLayer = outputLayer->poDS->CreateLayer(
709 90 : poSrcLayer->GetDescription(), poFirstGeomFieldDefn.get(),
710 90 : modLayerCreationOptions.List());
711 90 : if (!poLayer)
712 : {
713 1 : return false;
714 : }
715 89 : outputLayer->poLayer = poLayer;
716 89 : int iField = -1;
717 432 : for (const auto *poFieldDefn : poSrcFeatureDefn->GetFields())
718 : {
719 344 : ++iField;
720 344 : if (omitPartitionedFields && abPartitionedFields[iField])
721 23 : continue;
722 321 : if (poLayer->CreateField(poFieldDefn) != OGRERR_NONE)
723 : {
724 1 : alg->ReportError(CE_Failure, CPLE_AppDefined,
725 : "Cannot create field '%s'",
726 : poFieldDefn->GetNameRef());
727 1 : return false;
728 : }
729 : }
730 88 : int iGeomField = -1;
731 68 : for (const auto *poGeomFieldDefn :
732 224 : poSrcFeatureDefn->GetGeomFields())
733 : {
734 69 : ++iGeomField;
735 69 : if (iGeomField > 0)
736 : {
737 3 : OGRGeomFieldDefn oClone(poGeomFieldDefn);
738 3 : if (abPartitionedGeomFields[iGeomField])
739 : {
740 2 : if (aeGeomTypes[iGeomField] == wkbNone)
741 0 : continue;
742 4 : whileUnsealing(&oClone)->SetType(
743 2 : aeGeomTypes[iGeomField]);
744 : }
745 3 : if (poLayer->CreateGeomField(&oClone) != OGRERR_NONE)
746 : {
747 1 : alg->ReportError(CE_Failure, CPLE_AppDefined,
748 : "Cannot create geometry field '%s'",
749 : poGeomFieldDefn->GetNameRef());
750 1 : return false;
751 : }
752 : }
753 : }
754 :
755 87 : if (bUseTransactions)
756 : {
757 70 : if (outputLayer->poDS->StartTransaction() != OGRERR_NONE)
758 0 : return false;
759 : }
760 : }
761 :
762 90 : const auto nCounter = CPLGetErrorCounter();
763 90 : oCacheOutputLayer.insert(osKey, outputLayer);
764 : // In case insertion caused an eviction and old dataset
765 : // flushing caused an error
766 90 : if (CPLGetErrorCounter() != nCounter)
767 0 : return false;
768 : }
769 :
770 10109 : return true;
771 : }
772 :
773 : /************************************************************************/
774 : /* GDALVectorPartitionAlgorithm::RunStep() */
775 : /************************************************************************/
776 :
777 46 : bool GDALVectorPartitionAlgorithm::RunStep(GDALPipelineStepRunContext &ctxt)
778 : {
779 46 : auto poSrcDS = m_inputDataset[0].GetDatasetRef();
780 46 : CPLAssert(poSrcDS);
781 :
782 46 : auto poOutDriver = poSrcDS->GetDriver();
783 : const char *pszExtensions =
784 46 : poOutDriver ? poOutDriver->GetMetadataItem(GDAL_DMD_EXTENSIONS)
785 46 : : nullptr;
786 46 : if (m_format.empty())
787 : {
788 1 : if (!pszExtensions)
789 : {
790 1 : ReportError(CE_Failure, CPLE_AppDefined,
791 : "Cannot infer output format. Please specify "
792 : "'output-format' argument");
793 1 : return false;
794 : }
795 : }
796 : else
797 : {
798 45 : poOutDriver = GetGDALDriverManager()->GetDriverByName(m_format.c_str());
799 90 : if (!(poOutDriver && (pszExtensions = poOutDriver->GetMetadataItem(
800 45 : GDAL_DMD_EXTENSIONS)) != nullptr))
801 : {
802 1 : ReportError(CE_Failure, CPLE_AppDefined,
803 : "Output driver has no known file extension");
804 1 : return false;
805 : }
806 : }
807 44 : CPLAssert(poOutDriver);
808 :
809 : const bool bFormatSupportsAppend =
810 46 : poOutDriver->GetMetadataItem(GDAL_DCAP_UPDATE) ||
811 2 : poOutDriver->GetMetadataItem(GDAL_DCAP_APPEND);
812 44 : if (m_appendLayer && !bFormatSupportsAppend)
813 : {
814 1 : ReportError(CE_Failure, CPLE_AppDefined,
815 : "Driver '%s' does not support update",
816 1 : poOutDriver->GetDescription());
817 1 : return false;
818 : }
819 :
820 43 : const bool bParquetOutput = EQUAL(poOutDriver->GetDescription(), "PARQUET");
821 43 : if (bParquetOutput && m_scheme == SCHEME_HIVE)
822 : {
823 : // Required for Parquet Hive partitioning
824 4 : m_omitPartitionedFields = true;
825 : }
826 :
827 86 : const CPLStringList aosExtensions(CSLTokenizeString(pszExtensions));
828 43 : const char *pszExtension = aosExtensions[0];
829 :
830 86 : const CPLStringList datasetCreationOptions(m_creationOptions);
831 86 : const CPLStringList layerCreationOptions(m_layerCreationOptions);
832 :
833 : // We don't have driver metadata for that (and that would be a bit
834 : // tricky because some formats are half-text/half-binary), so...
835 : const bool bOutputFormatIsBinary =
836 39 : bParquetOutput || EQUAL(poOutDriver->GetDescription(), "GPKG") ||
837 83 : EQUAL(poOutDriver->GetDescription(), "SQLite") ||
838 1 : EQUAL(poOutDriver->GetDescription(), "FlatGeoBuf");
839 :
840 : // Below values have been experimentally determined and are not based
841 : // on rocket science...
842 43 : int nSpatialIndexPerFeatureConstant = 0;
843 43 : int nSpatialIndexPerLog2FeatureCountConstant = 0;
844 43 : if (CPLTestBool(
845 : layerCreationOptions.FetchNameValueDef("SPATIAL_INDEX", "YES")))
846 : {
847 42 : if (EQUAL(poOutDriver->GetDescription(), "GPKG"))
848 : {
849 36 : nSpatialIndexPerFeatureConstant =
850 : static_cast<int>(sizeof(double) * 4 + sizeof(uint32_t));
851 36 : nSpatialIndexPerLog2FeatureCountConstant = 1;
852 : }
853 6 : else if (EQUAL(poOutDriver->GetDescription(), "FlatGeoBuf"))
854 : {
855 0 : nSpatialIndexPerFeatureConstant = 1;
856 0 : nSpatialIndexPerLog2FeatureCountConstant =
857 : static_cast<int>(sizeof(double) * 4 + sizeof(uint64_t));
858 : }
859 : }
860 :
861 : const bool bUseTransactions =
862 43 : (EQUAL(poOutDriver->GetDescription(), "GPKG") ||
863 81 : EQUAL(poOutDriver->GetDescription(), "SQLite")) &&
864 38 : !m_skipErrors;
865 :
866 : VSIStatBufL sStat;
867 43 : if (VSIStatL(m_output.c_str(), &sStat) == 0)
868 : {
869 11 : if (m_overwrite)
870 : {
871 5 : bool emptyDir = true;
872 5 : bool hasDirLevel1WithEqual = false;
873 :
874 : // Do a sanity check to verify that this looks like a directory
875 : // generated by partition
876 :
877 5 : if (m_scheme == SCHEME_HIVE)
878 : {
879 : std::unique_ptr<VSIDIR, decltype(&VSICloseDir)> psDir(
880 3 : VSIOpenDir(m_output.c_str(), -1, nullptr), VSICloseDir);
881 3 : if (psDir)
882 : {
883 : while (const auto *psEntry =
884 6 : VSIGetNextDirEntry(psDir.get()))
885 : {
886 5 : emptyDir = false;
887 5 : if (VSI_ISDIR(psEntry->nMode))
888 : {
889 5 : std::string_view v(psEntry->pszName);
890 5 : if (std::count_if(
891 129 : v.begin(), v.end(), [](char c)
892 134 : { return c == '/' || c == '\\'; }) == 1)
893 : {
894 2 : const auto nPosDirSep = v.find_first_of("/\\");
895 2 : const auto nPosEqual = v.find('=', nPosDirSep);
896 2 : if (nPosEqual != std::string::npos)
897 : {
898 2 : hasDirLevel1WithEqual = true;
899 2 : break;
900 : }
901 : }
902 : }
903 3 : }
904 : }
905 :
906 3 : if (!hasDirLevel1WithEqual && !emptyDir)
907 : {
908 1 : ReportError(
909 : CE_Failure, CPLE_AppDefined,
910 : "Rejecting removing '%s' as it does not look like "
911 : "a directory generated by this utility. If you are "
912 : "sure, remove it manually and re-run",
913 : m_output.c_str());
914 1 : return false;
915 : }
916 : }
917 : else
918 : {
919 2 : bool hasSubDir = false;
920 : std::unique_ptr<VSIDIR, decltype(&VSICloseDir)> psDir(
921 2 : VSIOpenDir(m_output.c_str(), 0, nullptr), VSICloseDir);
922 2 : if (psDir)
923 : {
924 : while (const auto *psEntry =
925 6 : VSIGetNextDirEntry(psDir.get()))
926 : {
927 5 : if (VSI_ISDIR(psEntry->nMode))
928 : {
929 1 : hasSubDir = true;
930 1 : break;
931 : }
932 4 : }
933 : }
934 :
935 2 : if (hasSubDir)
936 : {
937 1 : ReportError(
938 : CE_Failure, CPLE_AppDefined,
939 : "Rejecting removing '%s' as it does not look like "
940 : "a directory generated by this utility. If you are "
941 : "sure, remove it manually and re-run",
942 : m_output.c_str());
943 1 : return false;
944 : }
945 : }
946 :
947 3 : if (VSIRmdirRecursive(m_output.c_str()) != 0)
948 : {
949 0 : ReportError(CE_Failure, CPLE_AppDefined, "Cannot remove '%s'",
950 : m_output.c_str());
951 0 : return false;
952 : }
953 : }
954 6 : else if (!m_appendLayer)
955 : {
956 1 : ReportError(CE_Failure, CPLE_AppDefined,
957 : "'%s' already exists. Specify --overwrite or --append",
958 : m_output.c_str());
959 1 : return false;
960 : }
961 : }
962 40 : if (VSIStatL(m_output.c_str(), &sStat) != 0)
963 : {
964 35 : if (VSIMkdir(m_output.c_str(), DIRECTORY_CREATION_MODE) != 0)
965 : {
966 1 : ReportError(CE_Failure, CPLE_AppDefined,
967 : "Cannot create directory '%s'", m_output.c_str());
968 1 : return false;
969 : }
970 : }
971 :
972 91 : for (OGRLayer *poSrcLayer : poSrcDS->GetLayers())
973 : {
974 : const std::string osLayerDir =
975 63 : m_scheme == SCHEME_HIVE
976 : ? CPLFormFilenameSafe(
977 : m_output.c_str(),
978 177 : PercentEncode(poSrcLayer->GetDescription()).c_str(),
979 : nullptr)
980 120 : : m_output;
981 120 : if (m_scheme == SCHEME_HIVE &&
982 57 : VSIStatL(osLayerDir.c_str(), &sStat) != 0)
983 : {
984 50 : if (VSIMkdir(osLayerDir.c_str(), DIRECTORY_CREATION_MODE) != 0)
985 : {
986 0 : ReportError(CE_Failure, CPLE_AppDefined,
987 : "Cannot create directory '%s'", osLayerDir.c_str());
988 0 : return false;
989 : }
990 : }
991 :
992 63 : const auto poSrcFeatureDefn = poSrcLayer->GetLayerDefn();
993 :
994 : struct Field
995 : {
996 : int nIdx{};
997 : bool bIsGeom = false;
998 : std::string encodedFieldName{};
999 : OGRFieldType eType{};
1000 : };
1001 :
1002 63 : std::vector<Field> asFields;
1003 126 : std::vector<bool> abPartitionedFields(poSrcFeatureDefn->GetFieldCount(),
1004 63 : false);
1005 : std::vector<bool> abPartitionedGeomFields(
1006 63 : poSrcFeatureDefn->GetGeomFieldCount(), false);
1007 124 : for (const std::string &fieldName : m_fields)
1008 : {
1009 63 : int nIdx = poSrcFeatureDefn->GetFieldIndex(fieldName.c_str());
1010 63 : if (nIdx < 0)
1011 : {
1012 3 : if (EQUAL(fieldName.c_str(), "OGR_GEOMETRY") &&
1013 0 : poSrcFeatureDefn->GetGeomFieldCount() > 0)
1014 : {
1015 0 : CPLError(CE_Warning, CPLE_AppDefined,
1016 : "'%s' is deprecated. Please use '%s' instead",
1017 : "OGR_GEOMETRY",
1018 : OGR_GEOMETRY_DEFAULT_NON_EMPTY_NAME);
1019 0 : nIdx = 0;
1020 : }
1021 3 : else if (EQUAL(fieldName.c_str(),
1022 4 : OGR_GEOMETRY_DEFAULT_NON_EMPTY_NAME) &&
1023 1 : poSrcFeatureDefn->GetGeomFieldCount() > 0)
1024 : {
1025 1 : nIdx = 0;
1026 : }
1027 : else
1028 : nIdx =
1029 2 : poSrcFeatureDefn->GetGeomFieldIndex(fieldName.c_str());
1030 3 : if (nIdx < 0)
1031 : {
1032 1 : ReportError(CE_Failure, CPLE_AppDefined,
1033 : "Cannot find field '%s' in layer '%s'",
1034 : fieldName.c_str(),
1035 1 : poSrcLayer->GetDescription());
1036 2 : return false;
1037 : }
1038 : else
1039 : {
1040 2 : abPartitionedGeomFields[nIdx] = true;
1041 4 : Field f;
1042 2 : f.nIdx = nIdx;
1043 2 : f.bIsGeom = true;
1044 2 : if (fieldName.empty())
1045 : f.encodedFieldName =
1046 0 : OGR_GEOMETRY_DEFAULT_NON_EMPTY_NAME;
1047 : else
1048 2 : f.encodedFieldName = PercentEncode(fieldName);
1049 2 : asFields.push_back(std::move(f));
1050 : }
1051 : }
1052 : else
1053 : {
1054 : const auto eType =
1055 60 : poSrcFeatureDefn->GetFieldDefn(nIdx)->GetType();
1056 60 : if (eType != OFTString && eType != OFTInteger &&
1057 : eType != OFTInteger64)
1058 : {
1059 1 : ReportError(
1060 : CE_Failure, CPLE_NotSupported,
1061 : "Field '%s' not valid for partitioning. Only fields of "
1062 : "type String, Integer or Integer64, or geometry fields,"
1063 : " are accepted",
1064 : fieldName.c_str());
1065 1 : return false;
1066 : }
1067 59 : abPartitionedFields[nIdx] = true;
1068 118 : Field f;
1069 59 : f.nIdx = nIdx;
1070 59 : f.bIsGeom = false;
1071 59 : f.encodedFieldName = PercentEncode(fieldName);
1072 59 : f.eType = eType;
1073 59 : asFields.push_back(std::move(f));
1074 : }
1075 : }
1076 :
1077 61 : std::vector<OGRFieldType> aeSrcFieldTypes;
1078 305 : for (const auto *poFieldDefn : poSrcFeatureDefn->GetFields())
1079 : {
1080 244 : aeSrcFieldTypes.push_back(poFieldDefn->GetType());
1081 : }
1082 :
1083 : std::unique_ptr<OGRFeatureDefn> poFeatureDefnWithoutPartitionedFields(
1084 61 : poSrcFeatureDefn->Clone());
1085 61 : std::vector<int> anMapForSetFrom;
1086 61 : if (m_omitPartitionedFields)
1087 : {
1088 : // Sort fields by descending index (so we can delete them easily)
1089 10 : std::vector<Field> sortedFields(asFields);
1090 10 : std::sort(sortedFields.begin(), sortedFields.end(),
1091 4 : [](const Field &a, const Field &b)
1092 4 : { return a.nIdx > b.nIdx; });
1093 24 : for (const auto &field : sortedFields)
1094 : {
1095 14 : if (!field.bIsGeom)
1096 14 : poFeatureDefnWithoutPartitionedFields->DeleteFieldDefn(
1097 14 : field.nIdx);
1098 : }
1099 : anMapForSetFrom =
1100 20 : poFeatureDefnWithoutPartitionedFields->ComputeMapForSetFrom(
1101 10 : poSrcFeatureDefn);
1102 : }
1103 :
1104 : lru11::Cache<std::string, std::shared_ptr<Layer>> oCacheOutputLayer(
1105 61 : m_maxCacheSize, 0);
1106 61 : std::shared_ptr<Layer> outputLayer = std::make_unique<Layer>();
1107 61 : outputLayer->bUseTransactions = bUseTransactions;
1108 :
1109 61 : GIntBig nTotalFeatures = 1;
1110 61 : GIntBig nFeatureIter = 0;
1111 61 : if (ctxt.m_pfnProgress)
1112 5 : nTotalFeatures = poSrcLayer->GetFeatureCount(true);
1113 : const double dfInvTotalFeatures =
1114 61 : 1.0 / static_cast<double>(std::max<GIntBig>(1, nTotalFeatures));
1115 :
1116 61 : std::string osAttrQueryString;
1117 61 : if (const char *pszAttrQueryString = poSrcLayer->GetAttrQueryString())
1118 0 : osAttrQueryString = pszAttrQueryString;
1119 :
1120 61 : std::string osKeyTmp;
1121 61 : std::vector<OGRwkbGeometryType> aeGeomTypesTmp;
1122 : const auto BuildKey =
1123 10118 : [&osKeyTmp, &aeGeomTypesTmp](const std::vector<Field> &fields,
1124 : const OGRFeature *poFeature)
1125 70847 : -> std::pair<const std::string &,
1126 : const std::vector<OGRwkbGeometryType> &>
1127 : {
1128 10118 : osKeyTmp.clear();
1129 10118 : aeGeomTypesTmp.resize(poFeature->GetDefnRef()->GetGeomFieldCount());
1130 20236 : for (const auto &field : fields)
1131 : {
1132 10118 : if (!osKeyTmp.empty())
1133 8 : osKeyTmp += '/';
1134 10118 : osKeyTmp += field.encodedFieldName;
1135 10118 : osKeyTmp += '=';
1136 10118 : if (field.bIsGeom)
1137 : {
1138 9 : const auto poGeom = poFeature->GetGeomFieldRef(field.nIdx);
1139 9 : if (poGeom)
1140 : {
1141 8 : aeGeomTypesTmp[field.nIdx] = poGeom->getGeometryType();
1142 8 : osKeyTmp += poGeom->getGeometryName();
1143 8 : if (poGeom->Is3D())
1144 2 : osKeyTmp += 'Z';
1145 8 : if (poGeom->IsMeasured())
1146 2 : osKeyTmp += 'M';
1147 : }
1148 : else
1149 : {
1150 1 : aeGeomTypesTmp[field.nIdx] = wkbNone;
1151 1 : osKeyTmp += NULL_MARKER;
1152 : }
1153 : }
1154 10109 : else if (poFeature->IsFieldSetAndNotNull(field.nIdx))
1155 : {
1156 10081 : if (field.eType == OFTString)
1157 : {
1158 10068 : PercentEncode(
1159 : osKeyTmp,
1160 10068 : poFeature->GetFieldAsStringUnsafe(field.nIdx));
1161 : }
1162 13 : else if (field.eType == OFTInteger)
1163 : {
1164 : osKeyTmp += CPLSPrintf(
1165 : "%d",
1166 12 : poFeature->GetFieldAsIntegerUnsafe(field.nIdx));
1167 : }
1168 : else
1169 : {
1170 : osKeyTmp += CPLSPrintf(
1171 : CPL_FRMT_GIB,
1172 1 : poFeature->GetFieldAsInteger64Unsafe(field.nIdx));
1173 : }
1174 : }
1175 : else
1176 : {
1177 28 : osKeyTmp += NULL_MARKER;
1178 : }
1179 : }
1180 10118 : return {osKeyTmp, aeGeomTypesTmp};
1181 61 : };
1182 :
1183 61 : std::set<std::string> oSetKeys;
1184 61 : if (!bFormatSupportsAppend)
1185 : {
1186 2 : CPLDebug(
1187 : "GDAL",
1188 : "First pass to determine all distinct partitioned values...");
1189 :
1190 2 : if (asFields.size() == 1 && !asFields[0].bIsGeom)
1191 : {
1192 2 : std::string osSQL = "SELECT DISTINCT \"";
1193 2 : osSQL += CPLString(m_fields[0]).replaceAll('"', "\"\"");
1194 2 : osSQL += "\" FROM \"";
1195 2 : osSQL += CPLString(poSrcLayer->GetDescription())
1196 2 : .replaceAll('"', "\"\"");
1197 2 : osSQL += '"';
1198 2 : if (!osAttrQueryString.empty())
1199 : {
1200 0 : osSQL += " WHERE ";
1201 0 : osSQL += osAttrQueryString;
1202 : }
1203 : auto poSQLLayer =
1204 2 : poSrcDS->ExecuteSQL(osSQL.c_str(), nullptr, nullptr);
1205 2 : if (!poSQLLayer)
1206 0 : return false;
1207 8 : std::vector<Field> asSingleField{asFields[0]};
1208 2 : asSingleField[0].nIdx = 0;
1209 5 : for (auto &poFeature : *poSQLLayer)
1210 : {
1211 3 : const auto sPair = BuildKey(asFields, poFeature.get());
1212 3 : const std::string &osKey = sPair.first;
1213 3 : oSetKeys.insert(osKey);
1214 : #ifdef DEBUG_VERBOSE
1215 : CPLDebug("GDAL", "Found %s", osKey.c_str());
1216 : #endif
1217 : }
1218 2 : poSrcDS->ReleaseResultSet(poSQLLayer);
1219 :
1220 2 : if (!osAttrQueryString.empty())
1221 : {
1222 0 : poSrcLayer->SetAttributeFilter(osAttrQueryString.c_str());
1223 : }
1224 : }
1225 : else
1226 : {
1227 0 : for (auto &poFeature : *poSrcLayer)
1228 : {
1229 0 : const auto sPair = BuildKey(asFields, poFeature.get());
1230 0 : const std::string &osKey = sPair.first;
1231 0 : if (oSetKeys.insert(osKey).second)
1232 : {
1233 : #ifdef DEBUG_VERBOSE
1234 : CPLDebug("GDAL", "Found %s", osKey.c_str());
1235 : #endif
1236 : }
1237 : }
1238 : }
1239 2 : CPLDebug("GDAL",
1240 : "End of first pass: %d unique partitioning keys found -> "
1241 : "%d pass(es) needed",
1242 2 : static_cast<int>(oSetKeys.size()),
1243 2 : static_cast<int>((oSetKeys.size() + m_maxCacheSize - 1) /
1244 2 : m_maxCacheSize));
1245 :
1246 : // If we have less distinct values as the maximum cache size, we
1247 : // can do a single iteration.
1248 2 : if (oSetKeys.size() <= static_cast<size_t>(m_maxCacheSize))
1249 2 : oSetKeys.clear();
1250 : }
1251 :
1252 61 : std::set<std::string> oSetOutputDatasets;
1253 61 : auto oSetKeysIter = oSetKeys.begin();
1254 : while (true)
1255 : {
1256 : // Determine which keys are allowed for the current pass
1257 61 : std::set<std::string> oSetKeysAllowedInThisPass;
1258 61 : if (!oSetKeys.empty())
1259 : {
1260 0 : while (oSetKeysAllowedInThisPass.size() <
1261 0 : static_cast<size_t>(m_maxCacheSize) &&
1262 0 : oSetKeysIter != oSetKeys.end())
1263 : {
1264 0 : oSetKeysAllowedInThisPass.insert(*oSetKeysIter);
1265 0 : ++oSetKeysIter;
1266 : }
1267 0 : if (oSetKeysAllowedInThisPass.empty())
1268 0 : break;
1269 : }
1270 :
1271 10168 : for (auto &poFeature : *poSrcLayer)
1272 : {
1273 10115 : const auto sPair = BuildKey(asFields, poFeature.get());
1274 10115 : const std::string &osKey = sPair.first;
1275 10115 : const auto &aeGeomTypes = sPair.second;
1276 :
1277 10115 : if (!oSetKeysAllowedInThisPass.empty() &&
1278 0 : !cpl::contains(oSetKeysAllowedInThisPass, osKey))
1279 : {
1280 1 : continue;
1281 : }
1282 :
1283 20230 : if (!GetCurrentOutputLayer(
1284 : this, poSrcFeatureDefn, poSrcLayer, osKey, aeGeomTypes,
1285 10115 : osLayerDir, m_scheme, m_pattern,
1286 10115 : m_partDigitLeadingZeroes, m_partDigitCount,
1287 10115 : m_featureLimit, m_maxFileSize, m_omitPartitionedFields,
1288 : abPartitionedFields, abPartitionedGeomFields,
1289 : pszExtension, poOutDriver, datasetCreationOptions,
1290 : layerCreationOptions,
1291 10115 : poFeatureDefnWithoutPartitionedFields.get(),
1292 10115 : poFeature->GetGeometryRef()
1293 : ? nSpatialIndexPerFeatureConstant
1294 : : 0,
1295 : nSpatialIndexPerLog2FeatureCountConstant,
1296 : bUseTransactions, oCacheOutputLayer, outputLayer))
1297 : {
1298 8 : return false;
1299 : }
1300 :
1301 10109 : if (bParquetOutput)
1302 : {
1303 : oSetOutputDatasets.insert(
1304 14 : outputLayer->poDS->GetDescription());
1305 : }
1306 :
1307 10109 : if (m_appendLayer)
1308 8 : poFeature->SetFID(OGRNullFID);
1309 :
1310 : OGRErr eErr;
1311 20200 : if (m_omitPartitionedFields ||
1312 0 : std::find(aeGeomTypes.begin(), aeGeomTypes.end(),
1313 20200 : wkbNone) != aeGeomTypes.end())
1314 : {
1315 19 : OGRFeature oFeat(outputLayer->poLayer->GetLayerDefn());
1316 19 : oFeat.SetFrom(poFeature.get(), anMapForSetFrom.data());
1317 19 : oFeat.SetFID(poFeature->GetFID());
1318 19 : eErr = outputLayer->poLayer->CreateFeature(&oFeat);
1319 : }
1320 : else
1321 : {
1322 20180 : poFeature->SetFDefnUnsafe(
1323 10090 : outputLayer->poLayer->GetLayerDefn());
1324 10090 : eErr = outputLayer->poLayer->CreateFeature(poFeature.get());
1325 : }
1326 10109 : if (eErr != OGRERR_NONE)
1327 : {
1328 2 : ReportError(m_skipErrors ? CE_Warning : CE_Failure,
1329 : CPLE_AppDefined,
1330 : "Cannot insert feature " CPL_FRMT_GIB,
1331 : poFeature->GetFID());
1332 2 : if (m_skipErrors)
1333 1 : continue;
1334 1 : return false;
1335 : }
1336 10107 : ++outputLayer->nFeatureCount;
1337 :
1338 20194 : if (bUseTransactions &&
1339 10087 : (outputLayer->nFeatureCount % m_transactionSize) == 0)
1340 : {
1341 8 : if (outputLayer->poDS->CommitTransaction() != OGRERR_NONE ||
1342 4 : outputLayer->poDS->StartTransaction() != OGRERR_NONE)
1343 : {
1344 0 : return false;
1345 : }
1346 : }
1347 :
1348 : // Compute a rough estimate of the space taken by the feature
1349 10107 : if (m_maxFileSize > 0)
1350 : {
1351 10000 : outputLayer->nFileSize += GetEstimatedFeatureSize(
1352 10000 : poFeature.get(), abPartitionedFields,
1353 10000 : m_omitPartitionedFields, aeSrcFieldTypes,
1354 : bOutputFormatIsBinary);
1355 : }
1356 :
1357 10107 : ++nFeatureIter;
1358 10116 : if (ctxt.m_pfnProgress &&
1359 9 : !ctxt.m_pfnProgress(
1360 10107 : std::min(1.0, static_cast<double>(nFeatureIter) *
1361 9 : dfInvTotalFeatures),
1362 : "", ctxt.m_pProgressData))
1363 : {
1364 1 : ReportError(CE_Failure, CPLE_UserInterrupt,
1365 : "Interrupted by user");
1366 1 : return false;
1367 : }
1368 : }
1369 :
1370 53 : if (oSetKeysIter == oSetKeys.end())
1371 53 : break;
1372 0 : }
1373 :
1374 53 : const auto nCounter = CPLGetErrorCounter();
1375 53 : outputLayer.reset();
1376 53 : oCacheOutputLayer.clear();
1377 53 : if (CPLGetErrorCounter() != nCounter)
1378 1 : return false;
1379 :
1380 : // For Parquet output, create special "_metadata" file that contains
1381 : // the schema and references the individual files
1382 52 : if (bParquetOutput && !oSetOutputDatasets.empty())
1383 : {
1384 : auto poAlg =
1385 7 : GDALGlobalAlgorithmRegistry::GetSingleton().Instantiate(
1386 14 : "driver", "parquet", "create-metadata-file");
1387 7 : if (poAlg)
1388 : {
1389 7 : auto inputArg = poAlg->GetArg(GDAL_ARG_NAME_INPUT);
1390 7 : auto outputArg = poAlg->GetArg(GDAL_ARG_NAME_OUTPUT);
1391 7 : if (inputArg && inputArg->GetType() == GAAT_DATASET_LIST &&
1392 14 : outputArg && outputArg->GetType() == GAAT_DATASET)
1393 : {
1394 7 : std::vector<std::string> asInputFilenames;
1395 7 : asInputFilenames.insert(asInputFilenames.end(),
1396 : oSetOutputDatasets.begin(),
1397 14 : oSetOutputDatasets.end());
1398 7 : inputArg->Set(asInputFilenames);
1399 7 : outputArg->Set(CPLFormFilenameSafe(osLayerDir.c_str(),
1400 : "_metadata", nullptr));
1401 7 : if (!poAlg->Run())
1402 0 : return false;
1403 : }
1404 : }
1405 : }
1406 : }
1407 :
1408 28 : return true;
1409 : }
1410 :
1411 : /************************************************************************/
1412 : /* GDALVectorPartitionAlgorithm::RunImpl() */
1413 : /************************************************************************/
1414 :
1415 45 : bool GDALVectorPartitionAlgorithm::RunImpl(GDALProgressFunc pfnProgress,
1416 : void *pProgressData)
1417 : {
1418 45 : GDALPipelineStepRunContext stepCtxt;
1419 45 : stepCtxt.m_pfnProgress = pfnProgress;
1420 45 : stepCtxt.m_pProgressData = pProgressData;
1421 90 : return RunStep(stepCtxt);
1422 : }
1423 :
1424 : GDALVectorPartitionAlgorithmStandalone::
1425 : ~GDALVectorPartitionAlgorithmStandalone() = default;
1426 : //! @endcond
|
