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