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