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