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