Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: Feather Translator
4 : * Purpose: Implements OGRFeatherDriver.
5 : * Author: Even Rouault, <even.rouault at spatialys.com>
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2022, Planet Labs
9 : *
10 : * Permission is hereby granted, free of charge, to any person obtaining a
11 : * copy of this software and associated documentation files (the "Software"),
12 : * to deal in the Software without restriction, including without limitation
13 : * the rights to use, copy, modify, merge, publish, distribute, sublicense,
14 : * and/or sell copies of the Software, and to permit persons to whom the
15 : * Software is furnished to do so, subject to the following conditions:
16 : *
17 : * The above copyright notice and this permission notice shall be included
18 : * in all copies or substantial portions of the Software.
19 : *
20 : * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
21 : * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22 : * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
23 : * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24 : * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
25 : * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
26 : * DEALINGS IN THE SOFTWARE.
27 : ****************************************************************************/
28 :
29 : #include "cpl_json.h"
30 : #include "cpl_time.h"
31 : #include "gdal_pam.h"
32 : #include "ogrsf_frmts.h"
33 : #include "ogr_p.h"
34 :
35 : #include <cinttypes>
36 : #include <limits>
37 : #include <map>
38 : #include <set>
39 : #include <utility>
40 :
41 : #include "ogr_feather.h"
42 :
43 : #include "../arrow_common/ograrrowlayer.hpp"
44 : #include "../arrow_common/ograrrowdataset.hpp"
45 :
46 : /************************************************************************/
47 : /* OGRFeatherLayer() */
48 : /************************************************************************/
49 :
50 531 : OGRFeatherLayer::OGRFeatherLayer(
51 : OGRFeatherDataset *poDS, const char *pszLayerName,
52 531 : std::shared_ptr<arrow::ipc::RecordBatchFileReader> &poRecordBatchFileReader)
53 : : OGRArrowLayer(poDS, pszLayerName), m_poDS(poDS),
54 531 : m_poRecordBatchFileReader(poRecordBatchFileReader)
55 : {
56 531 : EstablishFeatureDefn();
57 531 : CPLAssert(static_cast<int>(m_aeGeomEncoding.size()) ==
58 : m_poFeatureDefn->GetGeomFieldCount());
59 531 : }
60 :
61 : /************************************************************************/
62 : /* OGRFeatherLayer() */
63 : /************************************************************************/
64 :
65 9 : OGRFeatherLayer::OGRFeatherLayer(
66 : OGRFeatherDataset *poDS, const char *pszLayerName,
67 : std::shared_ptr<arrow::io::RandomAccessFile> poFile, bool bSeekable,
68 : const arrow::ipc::IpcReadOptions &oOptions,
69 : std::shared_ptr<arrow::ipc::RecordBatchStreamReader>
70 9 : &poRecordBatchStreamReader)
71 : : OGRArrowLayer(poDS, pszLayerName), m_poDS(poDS),
72 9 : m_poFile(std::move(poFile)), m_bSeekable(bSeekable), m_oOptions(oOptions),
73 9 : m_poRecordBatchReader(poRecordBatchStreamReader)
74 : {
75 9 : EstablishFeatureDefn();
76 9 : CPLAssert(static_cast<int>(m_aeGeomEncoding.size()) ==
77 : m_poFeatureDefn->GetGeomFieldCount());
78 9 : }
79 :
80 : /************************************************************************/
81 : /* GetDataset() */
82 : /************************************************************************/
83 :
84 112 : GDALDataset *OGRFeatherLayer::GetDataset()
85 : {
86 112 : return m_poDS;
87 : }
88 :
89 : /************************************************************************/
90 : /* LoadGeoMetadata() */
91 : /************************************************************************/
92 :
93 540 : void OGRFeatherLayer::LoadGeoMetadata(
94 : const arrow::KeyValueMetadata *kv_metadata, const std::string &key)
95 : {
96 540 : if (kv_metadata && kv_metadata->Contains(key))
97 : {
98 1060 : auto geo = kv_metadata->Get(key);
99 530 : if (geo.ok())
100 : {
101 1060 : CPLJSONDocument oDoc;
102 530 : if (oDoc.LoadMemory(*geo))
103 : {
104 1060 : auto oRoot = oDoc.GetRoot();
105 1590 : const auto osVersion = oRoot.GetString("schema_version");
106 530 : if (key != GDAL_GEO_FOOTER_KEY && osVersion != "0.1.0")
107 : {
108 396 : CPLDebug("FEATHER",
109 : "schema_version = %s not explicitly handled by "
110 : "the driver",
111 : osVersion.c_str());
112 : }
113 1590 : auto oColumns = oRoot.GetObj("columns");
114 530 : if (oColumns.IsValid())
115 : {
116 1060 : for (const auto &oColumn : oColumns.GetChildren())
117 : {
118 530 : m_oMapGeometryColumns[oColumn.GetName()] = oColumn;
119 : }
120 : }
121 : }
122 : else
123 : {
124 0 : CPLError(CE_Warning, CPLE_AppDefined,
125 : "Cannot parse 'geo' metadata");
126 : }
127 : }
128 : }
129 540 : }
130 :
131 : /************************************************************************/
132 : /* EstablishFeatureDefn() */
133 : /************************************************************************/
134 :
135 540 : void OGRFeatherLayer::EstablishFeatureDefn()
136 : {
137 1080 : m_poSchema = m_poRecordBatchFileReader ? m_poRecordBatchFileReader->schema()
138 1080 : : m_poRecordBatchReader->schema();
139 540 : const auto &kv_metadata = m_poSchema->metadata();
140 :
141 : #ifdef DEBUG
142 540 : if (kv_metadata)
143 : {
144 1076 : for (const auto &keyValue : kv_metadata->sorted_pairs())
145 : {
146 542 : CPLDebug("FEATHER", "%s = %s", keyValue.first.c_str(),
147 : keyValue.second.c_str());
148 : }
149 : }
150 : #endif
151 :
152 : auto poFooterMetadata = m_poRecordBatchFileReader
153 531 : ? m_poRecordBatchFileReader->metadata()
154 1611 : : nullptr;
155 665 : if (poFooterMetadata && poFooterMetadata->Contains(GDAL_GEO_FOOTER_KEY) &&
156 125 : CPLTestBool(CPLGetConfigOption("OGR_ARROW_READ_GDAL_FOOTER", "YES")))
157 : {
158 125 : LoadGeoMetadata(poFooterMetadata.get(), GDAL_GEO_FOOTER_KEY);
159 : }
160 : else
161 : {
162 415 : LoadGeoMetadata(kv_metadata.get(), "geo");
163 : }
164 : const auto oMapFieldNameToGDALSchemaFieldDefn =
165 1080 : LoadGDALSchema(kv_metadata.get());
166 :
167 540 : const auto &fields = m_poSchema->fields();
168 3987 : for (int i = 0; i < m_poSchema->num_fields(); ++i)
169 : {
170 3447 : const auto &field = fields[i];
171 3447 : const auto &fieldName = field->name();
172 :
173 3447 : const auto &field_kv_metadata = field->metadata();
174 3447 : std::string osExtensionName;
175 3447 : if (field_kv_metadata)
176 : {
177 : auto extension_name =
178 1046 : field_kv_metadata->Get("ARROW:extension:name");
179 523 : if (extension_name.ok())
180 : {
181 520 : osExtensionName = *extension_name;
182 : }
183 : #ifdef DEBUG
184 523 : CPLDebug("FEATHER", "Metadata field %s:", fieldName.c_str());
185 1431 : for (const auto &keyValue : field_kv_metadata->sorted_pairs())
186 : {
187 908 : CPLDebug("FEATHER", " %s = %s", keyValue.first.c_str(),
188 : keyValue.second.c_str());
189 : }
190 : #endif
191 : }
192 :
193 3447 : if (!m_osFIDColumn.empty() && fieldName == m_osFIDColumn)
194 : {
195 7 : m_iFIDArrowColumn = i;
196 7 : continue;
197 : }
198 :
199 3440 : bool bRegularField = true;
200 3440 : auto oIter = m_oMapGeometryColumns.find(fieldName);
201 3440 : if (oIter != m_oMapGeometryColumns.end() || !osExtensionName.empty())
202 : {
203 1068 : CPLJSONObject oJSONDef;
204 534 : if (oIter != m_oMapGeometryColumns.end())
205 530 : oJSONDef = oIter->second;
206 1602 : auto osEncoding = oJSONDef.GetString("encoding");
207 534 : if (osEncoding.empty() && !osExtensionName.empty())
208 4 : osEncoding = std::move(osExtensionName);
209 :
210 534 : OGRwkbGeometryType eGeomType = wkbUnknown;
211 534 : auto eGeomEncoding = OGRArrowGeomEncoding::WKB;
212 534 : if (IsValidGeometryEncoding(field, osEncoding,
213 1068 : oIter != m_oMapGeometryColumns.end(),
214 : eGeomType, eGeomEncoding))
215 : {
216 531 : bRegularField = false;
217 1062 : OGRGeomFieldDefn oField(fieldName.c_str(), wkbUnknown);
218 :
219 1593 : const auto osWKT = oJSONDef.GetString("crs");
220 531 : if (osWKT.empty())
221 : {
222 : #if 0
223 : CPLError(CE_Warning, CPLE_AppDefined,
224 : "Missing required 'crs' field for geometry column %s",
225 : fieldName.c_str());
226 : #endif
227 : }
228 : else
229 : {
230 136 : OGRSpatialReference *poSRS = new OGRSpatialReference();
231 136 : poSRS->SetAxisMappingStrategy(OAMS_TRADITIONAL_GIS_ORDER);
232 :
233 136 : if (poSRS->importFromWkt(osWKT.c_str()) == OGRERR_NONE)
234 : {
235 136 : const double dfCoordEpoch = oJSONDef.GetDouble("epoch");
236 136 : if (dfCoordEpoch > 0)
237 2 : poSRS->SetCoordinateEpoch(dfCoordEpoch);
238 :
239 136 : oField.SetSpatialRef(poSRS);
240 : }
241 136 : poSRS->Release();
242 : }
243 :
244 : // m_aeGeomEncoding be filled before calling
245 : // ComputeGeometryColumnType()
246 531 : m_aeGeomEncoding.push_back(eGeomEncoding);
247 531 : if (eGeomType == wkbUnknown)
248 : {
249 705 : auto osType = oJSONDef.GetString("geometry_type");
250 235 : if (osType.empty())
251 235 : osType = oJSONDef.GetString("gdal:geometry_type");
252 470 : if (m_bSeekable && osType.empty() &&
253 235 : CPLTestBool(CPLGetConfigOption(
254 : "OGR_ARROW_COMPUTE_GEOMETRY_TYPE", "YES")))
255 : {
256 235 : eGeomType = ComputeGeometryColumnType(
257 235 : m_poFeatureDefn->GetGeomFieldCount(), i);
258 235 : if (m_poRecordBatchReader)
259 0 : ResetRecordBatchReader();
260 : }
261 : else
262 0 : eGeomType = GetGeometryTypeFromString(osType);
263 : }
264 :
265 531 : oField.SetType(eGeomType);
266 531 : oField.SetNullable(field->nullable());
267 531 : m_poFeatureDefn->AddGeomFieldDefn(&oField);
268 531 : m_anMapGeomFieldIndexToArrowColumn.push_back(i);
269 : }
270 : }
271 :
272 3440 : if (bRegularField)
273 : {
274 2909 : CreateFieldFromSchema(field, {i},
275 : oMapFieldNameToGDALSchemaFieldDefn);
276 : }
277 : }
278 :
279 540 : CPLAssert(static_cast<int>(m_anMapFieldIndexToArrowColumn.size()) ==
280 : m_poFeatureDefn->GetFieldCount());
281 540 : CPLAssert(static_cast<int>(m_anMapGeomFieldIndexToArrowColumn.size()) ==
282 : m_poFeatureDefn->GetGeomFieldCount());
283 540 : }
284 :
285 : /************************************************************************/
286 : /* ResetRecordBatchReader() */
287 : /************************************************************************/
288 :
289 12 : bool OGRFeatherLayer::ResetRecordBatchReader()
290 : {
291 12 : const auto nPos = *(m_poFile->Tell());
292 12 : CPL_IGNORE_RET_VAL(m_poFile->Seek(0));
293 : auto result =
294 24 : arrow::ipc::RecordBatchStreamReader::Open(m_poFile, m_oOptions);
295 12 : if (!result.ok())
296 : {
297 0 : CPLError(CE_Failure, CPLE_AppDefined,
298 : "RecordBatchStreamReader::Open() failed with %s",
299 0 : result.status().message().c_str());
300 0 : CPL_IGNORE_RET_VAL(m_poFile->Seek(nPos));
301 0 : return false;
302 : }
303 : else
304 : {
305 12 : m_poRecordBatchReader = *result;
306 12 : return true;
307 : }
308 : }
309 :
310 : /************************************************************************/
311 : /* ComputeGeometryColumnType() */
312 : /************************************************************************/
313 :
314 235 : OGRwkbGeometryType OGRFeatherLayer::ComputeGeometryColumnType(int iGeomCol,
315 : int iCol) const
316 : {
317 : // Compute type of geometry column by iterating over each geometry, and
318 : // looking at the WKB geometry type in the first 5 bytes of each geometry.
319 :
320 235 : OGRwkbGeometryType eGeomType = wkbNone;
321 :
322 235 : if (m_poRecordBatchReader != nullptr)
323 : {
324 0 : std::shared_ptr<arrow::RecordBatch> poBatch;
325 : while (true)
326 : {
327 0 : auto status = m_poRecordBatchReader->ReadNext(&poBatch);
328 0 : if (!status.ok())
329 : {
330 0 : CPLError(CE_Failure, CPLE_AppDefined, "ReadNext() failed: %s",
331 0 : status.message().c_str());
332 0 : break;
333 : }
334 0 : else if (!poBatch)
335 0 : break;
336 0 : eGeomType = ComputeGeometryColumnTypeProcessBatch(poBatch, iGeomCol,
337 : iCol, eGeomType);
338 0 : if (eGeomType == wkbUnknown)
339 0 : break;
340 0 : }
341 : }
342 : else
343 : {
344 470 : for (int iBatch = 0;
345 470 : iBatch < m_poRecordBatchFileReader->num_record_batches(); ++iBatch)
346 : {
347 235 : auto result = m_poRecordBatchFileReader->ReadRecordBatch(iBatch);
348 235 : if (!result.ok())
349 : {
350 0 : CPLError(CE_Failure, CPLE_AppDefined,
351 : "ReadRecordBatch() failed: %s",
352 0 : result.status().message().c_str());
353 0 : break;
354 : }
355 235 : eGeomType = ComputeGeometryColumnTypeProcessBatch(*result, iGeomCol,
356 : iCol, eGeomType);
357 235 : if (eGeomType == wkbUnknown)
358 0 : break;
359 : }
360 : }
361 :
362 235 : return eGeomType == wkbNone ? wkbUnknown : eGeomType;
363 : }
364 :
365 : /************************************************************************/
366 : /* BuildDomain() */
367 : /************************************************************************/
368 :
369 : std::unique_ptr<OGRFieldDomain>
370 19 : OGRFeatherLayer::BuildDomain(const std::string &osDomainName,
371 : int iFieldIndex) const
372 : {
373 19 : const int iArrowCol = m_anMapFieldIndexToArrowColumn[iFieldIndex][0];
374 19 : CPLAssert(m_poSchema->fields()[iArrowCol]->type()->id() ==
375 : arrow::Type::DICTIONARY);
376 :
377 19 : if (m_poRecordBatchReader)
378 : {
379 6 : if (m_poBatch)
380 : {
381 6 : return BuildDomainFromBatch(osDomainName, m_poBatch, iArrowCol);
382 : }
383 : }
384 13 : else if (m_poRecordBatchFileReader)
385 : {
386 13 : auto result = m_poRecordBatchFileReader->ReadRecordBatch(0);
387 13 : if (!result.ok())
388 : {
389 0 : CPLError(CE_Failure, CPLE_AppDefined,
390 : "ReadRecordBatch() failed: %s",
391 0 : result.status().message().c_str());
392 : }
393 13 : auto poBatch = *result;
394 13 : if (poBatch)
395 : {
396 13 : return BuildDomainFromBatch(osDomainName, poBatch, iArrowCol);
397 : }
398 : }
399 :
400 0 : return nullptr;
401 : }
402 :
403 : /************************************************************************/
404 : /* ResetReading() */
405 : /************************************************************************/
406 :
407 927 : void OGRFeatherLayer::ResetReading()
408 : {
409 927 : if (m_poRecordBatchReader != nullptr && m_iRecordBatch > 0)
410 : {
411 17 : if (m_iRecordBatch == 1 && m_poBatchIdx1)
412 : {
413 : // do nothing
414 : }
415 : else
416 : {
417 16 : m_bResetRecordBatchReaderAsked = true;
418 : }
419 : }
420 927 : OGRArrowLayer::ResetReading();
421 927 : }
422 :
423 : /************************************************************************/
424 : /* ReadNextBatch() */
425 : /************************************************************************/
426 :
427 1004 : bool OGRFeatherLayer::ReadNextBatch()
428 : {
429 1004 : if (m_poRecordBatchFileReader == nullptr)
430 : {
431 118 : return ReadNextBatchStream();
432 : }
433 : else
434 : {
435 886 : return ReadNextBatchFile();
436 : }
437 : }
438 :
439 : /************************************************************************/
440 : /* ReadNextBatchFile() */
441 : /************************************************************************/
442 :
443 886 : bool OGRFeatherLayer::ReadNextBatchFile()
444 : {
445 : while (true)
446 : {
447 886 : ++m_iRecordBatch;
448 886 : if (m_iRecordBatch == m_poRecordBatchFileReader->num_record_batches())
449 : {
450 449 : if (m_iRecordBatch == 1)
451 446 : m_iRecordBatch = 0;
452 : else
453 3 : m_poBatch.reset();
454 449 : return false;
455 : }
456 :
457 437 : m_nIdxInBatch = 0;
458 :
459 : auto result =
460 437 : m_poRecordBatchFileReader->ReadRecordBatch(m_iRecordBatch);
461 437 : if (!result.ok())
462 : {
463 0 : CPLError(CE_Failure, CPLE_AppDefined,
464 : "ReadRecordBatch() failed: %s",
465 0 : result.status().message().c_str());
466 0 : m_poBatch.reset();
467 0 : return false;
468 : }
469 437 : if ((*result)->num_rows() != 0)
470 : {
471 437 : SetBatch(*result);
472 437 : break;
473 : }
474 0 : }
475 :
476 437 : return true;
477 : }
478 :
479 : /************************************************************************/
480 : /* ReadNextBatchStream() */
481 : /************************************************************************/
482 :
483 153 : bool OGRFeatherLayer::ReadNextBatchStream()
484 : {
485 153 : m_nIdxInBatch = 0;
486 :
487 306 : std::shared_ptr<arrow::RecordBatch> poNextBatch;
488 0 : do
489 : {
490 153 : if (m_iRecordBatch == 0 && m_poBatchIdx0)
491 : {
492 1 : SetBatch(m_poBatchIdx0);
493 1 : m_iRecordBatch = 1;
494 101 : return true;
495 : }
496 :
497 152 : else if (m_iRecordBatch == 1 && m_poBatchIdx1)
498 : {
499 1 : SetBatch(m_poBatchIdx1);
500 1 : m_iRecordBatch = 2;
501 1 : return true;
502 : }
503 :
504 151 : else if (m_bSingleBatch)
505 : {
506 81 : CPLAssert(m_iRecordBatch == 0);
507 81 : CPLAssert(m_poBatch != nullptr);
508 81 : return false;
509 : }
510 :
511 70 : if (m_bResetRecordBatchReaderAsked)
512 : {
513 13 : if (!m_bSeekable)
514 : {
515 1 : CPLError(CE_Failure, CPLE_NotSupported,
516 : "Attempting to rewind non-seekable stream");
517 1 : return false;
518 : }
519 12 : if (!ResetRecordBatchReader())
520 0 : return false;
521 12 : m_bResetRecordBatchReaderAsked = false;
522 : }
523 :
524 69 : CPLAssert(m_poRecordBatchReader);
525 :
526 69 : ++m_iRecordBatch;
527 :
528 69 : poNextBatch.reset();
529 69 : auto status = m_poRecordBatchReader->ReadNext(&poNextBatch);
530 69 : if (!status.ok())
531 : {
532 0 : CPLError(CE_Failure, CPLE_AppDefined, "ReadNext() failed: %s",
533 0 : status.message().c_str());
534 0 : poNextBatch.reset();
535 : }
536 69 : if (poNextBatch == nullptr)
537 : {
538 17 : if (m_iRecordBatch == 1)
539 : {
540 3 : m_iRecordBatch = 0;
541 3 : m_bSingleBatch = true;
542 : }
543 : else
544 : {
545 14 : m_poBatch.reset();
546 14 : m_poBatchColumns.clear();
547 : }
548 17 : return false;
549 : }
550 52 : } while (poNextBatch->num_rows() == 0);
551 :
552 52 : SetBatch(poNextBatch);
553 :
554 52 : return true;
555 : }
556 :
557 : /************************************************************************/
558 : /* TryToCacheFirstTwoBatches() */
559 : /************************************************************************/
560 :
561 1 : void OGRFeatherLayer::TryToCacheFirstTwoBatches()
562 : {
563 2 : if (m_poRecordBatchReader != nullptr && m_iRecordBatch <= 0 &&
564 2 : !m_bSingleBatch && m_poBatchIdx0 == nullptr)
565 : {
566 1 : ResetReading();
567 1 : if (!m_poBatch)
568 : {
569 0 : CPL_IGNORE_RET_VAL(ReadNextBatchStream());
570 : }
571 1 : if (m_poBatch)
572 : {
573 2 : auto poBatchIdx0 = m_poBatch;
574 1 : if (ReadNextBatchStream())
575 : {
576 1 : CPLAssert(m_iRecordBatch == 1);
577 1 : m_poBatchIdx0 = poBatchIdx0;
578 1 : m_poBatchIdx1 = m_poBatch;
579 1 : SetBatch(poBatchIdx0);
580 1 : ResetReading();
581 : }
582 1 : ResetReading();
583 : }
584 : }
585 1 : }
586 :
587 : /************************************************************************/
588 : /* CanPostFilterArrowArray() */
589 : /************************************************************************/
590 :
591 20 : bool OGRFeatherLayer::CanPostFilterArrowArray(
592 : const struct ArrowSchema *schema) const
593 : {
594 20 : if (m_poRecordBatchReader)
595 10 : return false;
596 10 : return OGRArrowLayer::CanPostFilterArrowArray(schema);
597 : }
598 :
599 : /************************************************************************/
600 : /* InvalidateCachedBatches() */
601 : /************************************************************************/
602 :
603 109 : void OGRFeatherLayer::InvalidateCachedBatches()
604 : {
605 109 : if (m_poRecordBatchFileReader)
606 : {
607 63 : m_iRecordBatch = -1;
608 63 : ResetReading();
609 : }
610 109 : }
611 :
612 : /************************************************************************/
613 : /* GetFeatureCount() */
614 : /************************************************************************/
615 :
616 336 : GIntBig OGRFeatherLayer::GetFeatureCount(int bForce)
617 : {
618 628 : if (m_poRecordBatchFileReader != nullptr && m_poAttrQuery == nullptr &&
619 292 : m_poFilterGeom == nullptr)
620 : {
621 288 : auto result = m_poRecordBatchFileReader->CountRows();
622 288 : if (result.ok())
623 288 : return *result;
624 : }
625 48 : else if (m_poRecordBatchReader != nullptr)
626 : {
627 36 : if (!m_bSeekable && !bForce)
628 : {
629 1 : if (m_poAttrQuery == nullptr && m_poFilterGeom == nullptr)
630 : {
631 1 : TryToCacheFirstTwoBatches();
632 : }
633 :
634 1 : if (!m_bSingleBatch)
635 : {
636 1 : CPLError(
637 : CE_Failure, CPLE_AppDefined,
638 : "GetFeatureCount() cannot be run in non-forced mode on "
639 : "a non-seekable file made of several batches");
640 1 : return -1;
641 : }
642 : }
643 :
644 35 : if (m_poAttrQuery == nullptr && m_poFilterGeom == nullptr)
645 : {
646 23 : GIntBig nFeatures = 0;
647 23 : ResetReading();
648 23 : if (!m_poBatch)
649 3 : ReadNextBatchStream();
650 31 : while (m_poBatch)
651 : {
652 31 : nFeatures += m_poBatch->num_rows();
653 31 : if (!ReadNextBatchStream())
654 23 : break;
655 : }
656 23 : ResetReading();
657 23 : return nFeatures;
658 : }
659 : }
660 24 : return OGRLayer::GetFeatureCount(bForce);
661 : }
662 :
663 : /************************************************************************/
664 : /* CanRunNonForcedGetExtent() */
665 : /************************************************************************/
666 :
667 0 : bool OGRFeatherLayer::CanRunNonForcedGetExtent()
668 : {
669 0 : if (m_bSeekable)
670 0 : return true;
671 0 : TryToCacheFirstTwoBatches();
672 0 : if (!m_bSingleBatch)
673 : {
674 0 : CPLError(CE_Failure, CPLE_AppDefined,
675 : "GetExtent() cannot be run in non-forced mode on "
676 : "a non-seekable file made of several batches");
677 0 : return false;
678 : }
679 0 : return true;
680 : }
681 :
682 : /************************************************************************/
683 : /* TestCapability() */
684 : /************************************************************************/
685 :
686 292 : int OGRFeatherLayer::TestCapability(const char *pszCap)
687 : {
688 292 : if (EQUAL(pszCap, OLCFastFeatureCount))
689 : {
690 28 : return m_bSeekable && m_poAttrQuery == nullptr &&
691 28 : m_poFilterGeom == nullptr;
692 : }
693 :
694 274 : if (EQUAL(pszCap, OLCMeasuredGeometries))
695 16 : return true;
696 258 : if (EQUAL(pszCap, OLCZGeometries))
697 12 : return true;
698 :
699 246 : return OGRArrowLayer::TestCapability(pszCap);
700 : }
701 :
702 : /************************************************************************/
703 : /* GetMetadataItem() */
704 : /************************************************************************/
705 :
706 259 : const char *OGRFeatherLayer::GetMetadataItem(const char *pszName,
707 : const char *pszDomain)
708 : {
709 : // Mostly for unit test purposes
710 259 : if (pszDomain != nullptr && EQUAL(pszDomain, "_ARROW_"))
711 : {
712 9 : if (EQUAL(pszName, "FORMAT"))
713 : {
714 5 : return m_poRecordBatchFileReader ? "FILE" : "STREAM";
715 : }
716 4 : if (m_poRecordBatchFileReader != nullptr)
717 : {
718 4 : int iBatch = -1;
719 4 : if (EQUAL(pszName, "NUM_RECORD_BATCHES"))
720 : {
721 1 : return CPLSPrintf(
722 5 : "%d", m_poRecordBatchFileReader->num_record_batches());
723 : }
724 6 : else if (sscanf(pszName, "RECORD_BATCHES[%d]", &iBatch) == 1 &&
725 3 : strstr(pszName, ".NUM_ROWS"))
726 : {
727 : auto result =
728 6 : m_poRecordBatchFileReader->ReadRecordBatch(iBatch);
729 3 : if (!result.ok())
730 : {
731 0 : return nullptr;
732 : }
733 3 : return CPLSPrintf("%" PRId64, (*result)->num_rows());
734 : }
735 : }
736 0 : return nullptr;
737 : }
738 250 : else if (pszDomain != nullptr && EQUAL(pszDomain, "_ARROW_METADATA_"))
739 : {
740 : const auto kv_metadata =
741 5 : (m_poRecordBatchFileReader ? m_poRecordBatchFileReader->schema()
742 8 : : m_poRecordBatchReader->schema())
743 10 : ->metadata();
744 5 : if (kv_metadata && kv_metadata->Contains(pszName))
745 : {
746 5 : auto metadataItem = kv_metadata->Get(pszName);
747 5 : if (metadataItem.ok())
748 : {
749 5 : return CPLSPrintf("%s", metadataItem->c_str());
750 : }
751 : }
752 0 : return nullptr;
753 : }
754 476 : else if (m_poRecordBatchFileReader != nullptr && pszDomain != nullptr &&
755 231 : EQUAL(pszDomain, "_ARROW_FOOTER_METADATA_"))
756 : {
757 2 : const auto kv_metadata = m_poRecordBatchFileReader->metadata();
758 1 : if (kv_metadata && kv_metadata->Contains(pszName))
759 : {
760 1 : auto metadataItem = kv_metadata->Get(pszName);
761 1 : if (metadataItem.ok())
762 : {
763 1 : return CPLSPrintf("%s", metadataItem->c_str());
764 : }
765 : }
766 0 : return nullptr;
767 : }
768 244 : return OGRLayer::GetMetadataItem(pszName, pszDomain);
769 : }
770 :
771 : /************************************************************************/
772 : /* GetMetadata() */
773 : /************************************************************************/
774 :
775 34 : char **OGRFeatherLayer::GetMetadata(const char *pszDomain)
776 : {
777 : // Mostly for unit test purposes
778 34 : if (pszDomain != nullptr && EQUAL(pszDomain, "_ARROW_METADATA_"))
779 : {
780 5 : m_aosFeatherMetadata.Clear();
781 : const auto kv_metadata =
782 5 : (m_poRecordBatchFileReader ? m_poRecordBatchFileReader->schema()
783 8 : : m_poRecordBatchReader->schema())
784 10 : ->metadata();
785 5 : if (kv_metadata)
786 : {
787 13 : for (const auto &kv : kv_metadata->sorted_pairs())
788 : {
789 : m_aosFeatherMetadata.SetNameValue(kv.first.c_str(),
790 8 : kv.second.c_str());
791 : }
792 : }
793 5 : return m_aosFeatherMetadata.List();
794 : }
795 41 : if (m_poRecordBatchFileReader != nullptr && pszDomain != nullptr &&
796 12 : EQUAL(pszDomain, "_ARROW_FOOTER_METADATA_"))
797 : {
798 2 : m_aosFeatherMetadata.Clear();
799 4 : const auto kv_metadata = m_poRecordBatchFileReader->metadata();
800 2 : if (kv_metadata)
801 : {
802 3 : for (const auto &kv : kv_metadata->sorted_pairs())
803 : {
804 : m_aosFeatherMetadata.SetNameValue(kv.first.c_str(),
805 1 : kv.second.c_str());
806 : }
807 : }
808 2 : return m_aosFeatherMetadata.List();
809 : }
810 27 : return OGRLayer::GetMetadata(pszDomain);
811 : }
|
