Line data Source code
1 : /******************************************************************************
2 : *
3 : * Project: GDAL
4 : * Purpose: Arrow Database Connectivity driver
5 : * Author: Even Rouault, <even dot rouault at spatialys.com>
6 : *
7 : ******************************************************************************
8 : * Copyright (c) 2024, Even Rouault <even dot rouault at spatialys.com>
9 : * Copyright (c) 2024, Dewey Dunnington <dewey@voltrondata.com>
10 : *
11 : * SPDX-License-Identifier: MIT
12 : ****************************************************************************/
13 :
14 : #include "ogr_adbc.h"
15 : #include "ogr_spatialref.h"
16 : #include "ogr_p.h"
17 : #include "cpl_json.h"
18 :
19 : #include <cmath>
20 : #include <limits>
21 : #include <map>
22 : #include <set>
23 :
24 : #define ADBC_CALL(func, ...) m_poDS->m_driver.func(__VA_ARGS__)
25 :
26 0 : OGRArrowArrayToOGRFeatureAdapterLayer::~OGRArrowArrayToOGRFeatureAdapterLayer()
27 : {
28 0 : m_poLayerDefn->Release();
29 0 : }
30 :
31 : /************************************************************************/
32 : /* GetGeometryTypeFromString() */
33 : /************************************************************************/
34 :
35 0 : static OGRwkbGeometryType GetGeometryTypeFromString(const std::string &osType)
36 : {
37 0 : OGRwkbGeometryType eGeomType = wkbUnknown;
38 0 : OGRReadWKTGeometryType(osType.c_str(), &eGeomType);
39 0 : if (eGeomType == wkbUnknown && !osType.empty())
40 : {
41 0 : CPLDebug("ADBC", "Unknown geometry type: %s", osType.c_str());
42 : }
43 0 : return eGeomType;
44 : }
45 :
46 : /************************************************************************/
47 : /* OGRADBCLayer() */
48 : /************************************************************************/
49 :
50 0 : OGRADBCLayer::OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
51 : const char *pszStatement,
52 : std::unique_ptr<AdbcStatement> poStatement,
53 : std::unique_ptr<OGRArrowArrayStream> poStream,
54 0 : ArrowSchema *schema, bool bInternalUse)
55 : : m_poDS(poDS), m_osBaseStatement(pszStatement),
56 0 : m_osModifiedBaseStatement(m_osBaseStatement),
57 0 : m_statement(std::move(poStatement)), m_stream(std::move(poStream))
58 : {
59 0 : SetDescription(pszName);
60 :
61 0 : memcpy(&m_schema, schema, sizeof(m_schema));
62 0 : schema->release = nullptr;
63 :
64 0 : BuildLayerDefn(bInternalUse);
65 0 : }
66 :
67 : /************************************************************************/
68 : /* ParseGeometryColumnCovering() */
69 : /************************************************************************/
70 :
71 : //! Parse bounding box column definition
72 0 : static bool ParseGeometryColumnCovering(const CPLJSONObject &oJSONDef,
73 : std::string &osBBOXColumn,
74 : std::string &osXMin,
75 : std::string &osYMin,
76 : std::string &osXMax,
77 : std::string &osYMax)
78 : {
79 0 : const auto oCovering = oJSONDef["covering"];
80 0 : if (oCovering.IsValid() &&
81 0 : oCovering.GetType() == CPLJSONObject::Type::Object)
82 : {
83 0 : const auto oBBOX = oCovering["bbox"];
84 0 : if (oBBOX.IsValid() && oBBOX.GetType() == CPLJSONObject::Type::Object)
85 : {
86 0 : const auto oXMin = oBBOX["xmin"];
87 0 : const auto oYMin = oBBOX["ymin"];
88 0 : const auto oXMax = oBBOX["xmax"];
89 0 : const auto oYMax = oBBOX["ymax"];
90 0 : if (oXMin.IsValid() && oYMin.IsValid() && oXMax.IsValid() &&
91 0 : oYMax.IsValid() &&
92 0 : oXMin.GetType() == CPLJSONObject::Type::Array &&
93 0 : oYMin.GetType() == CPLJSONObject::Type::Array &&
94 0 : oXMax.GetType() == CPLJSONObject::Type::Array &&
95 0 : oYMax.GetType() == CPLJSONObject::Type::Array)
96 : {
97 0 : const auto osXMinArray = oXMin.ToArray();
98 0 : const auto osYMinArray = oYMin.ToArray();
99 0 : const auto osXMaxArray = oXMax.ToArray();
100 0 : const auto osYMaxArray = oYMax.ToArray();
101 0 : if (osXMinArray.Size() == 2 && osYMinArray.Size() == 2 &&
102 0 : osXMaxArray.Size() == 2 && osYMaxArray.Size() == 2 &&
103 0 : osXMinArray[0].GetType() == CPLJSONObject::Type::String &&
104 0 : osXMinArray[1].GetType() == CPLJSONObject::Type::String &&
105 0 : osYMinArray[0].GetType() == CPLJSONObject::Type::String &&
106 0 : osYMinArray[1].GetType() == CPLJSONObject::Type::String &&
107 0 : osXMaxArray[0].GetType() == CPLJSONObject::Type::String &&
108 0 : osXMaxArray[1].GetType() == CPLJSONObject::Type::String &&
109 0 : osYMaxArray[0].GetType() == CPLJSONObject::Type::String &&
110 0 : osYMaxArray[1].GetType() == CPLJSONObject::Type::String &&
111 0 : osXMinArray[0].ToString() == osYMinArray[0].ToString() &&
112 0 : osXMinArray[0].ToString() == osXMaxArray[0].ToString() &&
113 0 : osXMinArray[0].ToString() == osYMaxArray[0].ToString())
114 : {
115 0 : osBBOXColumn = osXMinArray[0].ToString();
116 0 : osXMin = osXMinArray[1].ToString();
117 0 : osYMin = osYMinArray[1].ToString();
118 0 : osXMax = osXMaxArray[1].ToString();
119 0 : osYMax = osYMaxArray[1].ToString();
120 0 : return true;
121 : }
122 : }
123 : }
124 : }
125 0 : return false;
126 : }
127 :
128 : /************************************************************************/
129 : /* ParseGeoParquetColumn() */
130 : /************************************************************************/
131 :
132 0 : static void ParseGeoParquetColumn(
133 : const CPLJSONObject &oColumn,
134 : std::map<std::string, OGRwkbGeometryType> &oMapType,
135 : std::map<std::string, OGREnvelope3D> &oMapExtent,
136 : std::map<std::string, OGRADBCLayer::GeomColBBOX>
137 : &oMapGeomColumnToCoveringBBOXColumn,
138 : std::map<std::string, std::unique_ptr<OGRSpatialReference>>
139 : &oMapGeomColumnsFromGeoParquet,
140 : std::set<std::string> &oSetCoveringBBoxColumn)
141 : {
142 0 : auto oCrs = oColumn.GetObj("crs");
143 0 : if (!oCrs.IsValid())
144 : {
145 : // WGS 84 is implied if no crs member is found.
146 0 : auto poSRS = std::make_unique<OGRSpatialReference>();
147 0 : poSRS->SetAxisMappingStrategy(OAMS_TRADITIONAL_GIS_ORDER);
148 0 : poSRS->importFromEPSG(4326);
149 0 : oMapGeomColumnsFromGeoParquet[oColumn.GetName()] = std::move(poSRS);
150 : }
151 0 : else if (oCrs.GetType() == CPLJSONObject::Type::Object)
152 : {
153 : // CRS encoded as PROJJSON (extension)
154 0 : const auto oType = oCrs["type"];
155 0 : if (oType.IsValid() && oType.GetType() == CPLJSONObject::Type::String)
156 : {
157 0 : const auto osType = oType.ToString();
158 0 : if (osType.find("CRS") != std::string::npos)
159 : {
160 0 : auto poSRS = std::make_unique<OGRSpatialReference>();
161 0 : poSRS->SetAxisMappingStrategy(OAMS_TRADITIONAL_GIS_ORDER);
162 :
163 0 : if (poSRS->SetFromUserInput(oCrs.ToString().c_str()) ==
164 : OGRERR_NONE)
165 : {
166 0 : oMapGeomColumnsFromGeoParquet[oColumn.GetName()] =
167 0 : std::move(poSRS);
168 : }
169 : }
170 : }
171 : }
172 : else
173 : {
174 0 : oMapGeomColumnsFromGeoParquet[oColumn.GetName()] = nullptr;
175 : }
176 :
177 0 : OGRwkbGeometryType eGeomType = wkbUnknown;
178 0 : auto oType = oColumn.GetObj("geometry_types");
179 0 : if (oType.GetType() == CPLJSONObject::Type::Array)
180 : {
181 0 : const auto oTypeArray = oType.ToArray();
182 0 : if (oTypeArray.Size() == 1)
183 : {
184 0 : eGeomType = GetGeometryTypeFromString(oTypeArray[0].ToString());
185 : }
186 0 : else if (oTypeArray.Size() > 1)
187 : {
188 0 : const auto PromoteToCollection = [](OGRwkbGeometryType eType)
189 : {
190 0 : if (eType == wkbPoint)
191 0 : return wkbMultiPoint;
192 0 : if (eType == wkbLineString)
193 0 : return wkbMultiLineString;
194 0 : if (eType == wkbPolygon)
195 0 : return wkbMultiPolygon;
196 0 : return eType;
197 : };
198 0 : bool bMixed = false;
199 0 : bool bHasMulti = false;
200 0 : bool bHasZ = false;
201 0 : bool bHasM = false;
202 0 : const auto eFirstType = OGR_GT_Flatten(
203 0 : GetGeometryTypeFromString(oTypeArray[0].ToString()));
204 0 : const auto eFirstTypeCollection = PromoteToCollection(eFirstType);
205 0 : for (int i = 0; i < oTypeArray.Size(); ++i)
206 : {
207 : const auto eThisGeom =
208 0 : GetGeometryTypeFromString(oTypeArray[i].ToString());
209 0 : if (PromoteToCollection(OGR_GT_Flatten(eThisGeom)) !=
210 : eFirstTypeCollection)
211 : {
212 0 : bMixed = true;
213 0 : break;
214 : }
215 0 : bHasZ |= OGR_GT_HasZ(eThisGeom) != FALSE;
216 0 : bHasM |= OGR_GT_HasM(eThisGeom) != FALSE;
217 0 : bHasMulti |= (PromoteToCollection(OGR_GT_Flatten(eThisGeom)) ==
218 0 : OGR_GT_Flatten(eThisGeom));
219 : }
220 0 : if (!bMixed)
221 : {
222 0 : if (eFirstTypeCollection == wkbMultiPolygon ||
223 : eFirstTypeCollection == wkbMultiLineString)
224 : {
225 0 : if (bHasMulti)
226 0 : eGeomType = OGR_GT_SetModifier(eFirstTypeCollection,
227 : bHasZ, bHasM);
228 : else
229 : eGeomType =
230 0 : OGR_GT_SetModifier(eFirstType, bHasZ, bHasM);
231 : }
232 : }
233 : }
234 : }
235 :
236 0 : oMapType[oColumn.GetName()] = eGeomType;
237 :
238 0 : OGREnvelope3D sExtent;
239 0 : const auto oBBox = oColumn.GetArray("bbox");
240 0 : if (oBBox.IsValid() && oBBox.Size() == 4)
241 : {
242 0 : sExtent.MinX = oBBox[0].ToDouble();
243 0 : sExtent.MinY = oBBox[1].ToDouble();
244 0 : sExtent.MinZ = std::numeric_limits<double>::infinity();
245 0 : sExtent.MaxX = oBBox[2].ToDouble();
246 0 : sExtent.MaxY = oBBox[3].ToDouble();
247 0 : sExtent.MaxZ = -std::numeric_limits<double>::infinity();
248 0 : if (sExtent.MinX <= sExtent.MaxX)
249 : {
250 0 : oMapExtent[oColumn.GetName()] = sExtent;
251 : }
252 : }
253 0 : else if (oBBox.IsValid() && oBBox.Size() == 6)
254 : {
255 0 : sExtent.MinX = oBBox[0].ToDouble();
256 0 : sExtent.MinY = oBBox[1].ToDouble();
257 0 : sExtent.MinZ = oBBox[2].ToDouble();
258 0 : sExtent.MaxX = oBBox[3].ToDouble();
259 0 : sExtent.MaxY = oBBox[4].ToDouble();
260 0 : sExtent.MaxZ = oBBox[5].ToDouble();
261 0 : if (sExtent.MinX <= sExtent.MaxX)
262 : {
263 0 : oMapExtent[oColumn.GetName()] = sExtent;
264 : }
265 : }
266 :
267 0 : std::string osBBOXColumn;
268 0 : std::string osXMin, osYMin, osXMax, osYMax;
269 0 : if (ParseGeometryColumnCovering(oColumn, osBBOXColumn, osXMin, osYMin,
270 : osXMax, osYMax))
271 : {
272 0 : OGRADBCLayer::GeomColBBOX geomColBBOX;
273 : const std::string osPrefix =
274 0 : std::string("\"")
275 0 : .append(OGRDuplicateCharacter(osBBOXColumn, '"'))
276 0 : .append("\".\"");
277 :
278 0 : const auto BuildColName = [&osPrefix](const std::string &s)
279 : {
280 0 : return std::string(osPrefix)
281 0 : .append(OGRDuplicateCharacter(s, '"'))
282 0 : .append("\"");
283 0 : };
284 :
285 0 : geomColBBOX.osXMin = BuildColName(osXMin);
286 0 : geomColBBOX.osYMin = BuildColName(osYMin);
287 0 : geomColBBOX.osXMax = BuildColName(osXMax);
288 0 : geomColBBOX.osYMax = BuildColName(osYMax);
289 0 : oMapGeomColumnToCoveringBBOXColumn[oColumn.GetName()] =
290 0 : std::move(geomColBBOX);
291 0 : oSetCoveringBBoxColumn.insert(std::move(osBBOXColumn));
292 : }
293 0 : }
294 :
295 : /************************************************************************/
296 : /* BuildLayerDefn() */
297 : /************************************************************************/
298 :
299 0 : void OGRADBCLayer::BuildLayerDefn(bool bInternalUse)
300 : {
301 : // Identify geometry columns for Parquet files, and query them with
302 : // ST_AsWKB() to avoid getting duckdb_spatial own geometry encoding
303 : // (https://github.com/duckdb/duckdb_spatial/blob/a60aa3733741a99c49baaf33390c0f7c8a9598a3/spatial/src/spatial/core/geometry/geometry_serialization.cpp#L11)
304 0 : std::map<std::string, std::unique_ptr<OGRSpatialReference>> oMapGeomColumns;
305 0 : std::map<std::string, OGRwkbGeometryType> oMapType;
306 0 : std::map<std::string, OGREnvelope3D> oMapExtent;
307 0 : std::map<std::string, GeomColBBOX> oMapGeomColumnToCoveringBBOXColumn;
308 0 : if (!bInternalUse && STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT ") &&
309 0 : (m_poDS->m_bIsDuckDBDriver ||
310 0 : (!m_poDS->m_osParquetFilename.empty() &&
311 0 : CPLString(m_osBaseStatement)
312 0 : .ifind(std::string(" FROM '").append(OGRDuplicateCharacter(
313 0 : m_poDS->m_osParquetFilename, '\''))) !=
314 : std::string::npos)))
315 : {
316 : // Try to read GeoParquet 'geo' metadata
317 : std::map<std::string, std::unique_ptr<OGRSpatialReference>>
318 0 : oMapGeomColumnsFromGeoParquet;
319 0 : std::set<std::string> oSetCoveringBBoxColumn;
320 :
321 0 : std::string osGeoParquetMD;
322 0 : if (!m_poDS->m_osParquetFilename.empty())
323 : {
324 0 : auto poMetadataLayer = m_poDS->CreateInternalLayer(
325 0 : std::string("SELECT value FROM parquet_kv_metadata('")
326 0 : .append(OGRDuplicateCharacter(m_poDS->m_osParquetFilename,
327 0 : '\''))
328 0 : .append("') WHERE key = 'geo'")
329 0 : .c_str());
330 0 : if (poMetadataLayer)
331 : {
332 : auto f = std::unique_ptr<OGRFeature>(
333 0 : poMetadataLayer->GetNextFeature());
334 0 : if (f)
335 : {
336 0 : int nBytes = 0;
337 0 : const GByte *pabyData = f->GetFieldAsBinary(0, &nBytes);
338 : osGeoParquetMD.assign(
339 0 : reinterpret_cast<const char *>(pabyData), nBytes);
340 : // CPLDebug("ADBC", "%s", osGeoParquetMD.c_str());
341 : }
342 : }
343 : }
344 0 : CPLJSONDocument oDoc;
345 0 : if (!osGeoParquetMD.empty() && oDoc.LoadMemory(osGeoParquetMD))
346 : {
347 0 : const auto oColumns = oDoc.GetRoot().GetObj("columns");
348 0 : for (const auto &oColumn : oColumns.GetChildren())
349 : {
350 0 : if (oColumn.GetString("encoding") == "WKB")
351 : {
352 0 : ParseGeoParquetColumn(oColumn, oMapType, oMapExtent,
353 : oMapGeomColumnToCoveringBBOXColumn,
354 : oMapGeomColumnsFromGeoParquet,
355 : oSetCoveringBBoxColumn);
356 : }
357 : }
358 : }
359 :
360 0 : auto poDescribeLayer = m_poDS->CreateInternalLayer(
361 0 : std::string("DESCRIBE ").append(m_osBaseStatement).c_str());
362 0 : std::string osNewStatement;
363 0 : bool bNewStatement = false;
364 0 : if (poDescribeLayer && (m_poDS->m_bIsDuckDBDriver ||
365 : // cppcheck-suppress knownConditionTrueFalse
366 0 : !oMapGeomColumnsFromGeoParquet.empty()))
367 : {
368 0 : for (auto &&f : *poDescribeLayer)
369 : {
370 0 : const char *pszColName = f->GetFieldAsString("column_name");
371 0 : if (cpl::contains(oSetCoveringBBoxColumn, pszColName))
372 : {
373 0 : bNewStatement = true;
374 0 : continue;
375 : }
376 :
377 : // f->DumpReadable(stdout);
378 0 : if (osNewStatement.empty())
379 0 : osNewStatement = "SELECT ";
380 : else
381 0 : osNewStatement += ", ";
382 :
383 0 : auto oIter = oMapGeomColumnsFromGeoParquet.find(pszColName);
384 0 : if (oIter != oMapGeomColumnsFromGeoParquet.end())
385 : {
386 0 : oMapGeomColumns[pszColName] = std::move(oIter->second);
387 : }
388 0 : if (EQUAL(f->GetFieldAsString("column_type"), "GEOMETRY") &&
389 0 : m_poDS->m_bSpatialLoaded)
390 : {
391 0 : bNewStatement = true;
392 0 : osNewStatement += "ST_AsWKB(\"";
393 0 : osNewStatement += OGRDuplicateCharacter(pszColName, '"');
394 0 : osNewStatement += "\") AS ";
395 0 : if (oIter == oMapGeomColumnsFromGeoParquet.end())
396 0 : oMapGeomColumns[pszColName] = nullptr;
397 : }
398 0 : osNewStatement += '"';
399 0 : osNewStatement += OGRDuplicateCharacter(pszColName, '"');
400 0 : osNewStatement += '"';
401 : }
402 0 : m_osModifiedSelect = osNewStatement;
403 0 : osNewStatement += " FROM (";
404 0 : osNewStatement += m_osBaseStatement;
405 0 : osNewStatement += " )";
406 : }
407 :
408 0 : if (bNewStatement)
409 : {
410 : // CPLDebug("ADBC", "%s -> %s", m_osBaseStatement.c_str(), osNewStatement.c_str());
411 0 : if (ReplaceStatement(osNewStatement.c_str()))
412 : {
413 0 : m_osModifiedBaseStatement = std::move(osNewStatement);
414 : }
415 : else
416 : {
417 0 : m_osModifiedSelect.clear();
418 0 : oMapGeomColumns.clear();
419 : }
420 : }
421 : }
422 :
423 0 : m_poAdapterLayer = std::make_unique<OGRArrowArrayToOGRFeatureAdapterLayer>(
424 0 : GetDescription());
425 :
426 0 : for (int i = 0; i < m_schema.n_children; ++i)
427 : {
428 0 : const char *pszColName = m_schema.children[i]->name;
429 0 : auto oIter = oMapGeomColumns.find(pszColName);
430 0 : if (oIter != oMapGeomColumns.end())
431 : {
432 0 : OGRGeomFieldDefn oGeomFieldDefn(pszColName, oMapType[pszColName]);
433 0 : auto poSRS = std::move(oIter->second).release();
434 0 : if (poSRS)
435 : {
436 0 : oGeomFieldDefn.SetSpatialRef(poSRS);
437 0 : poSRS->Release();
438 : }
439 0 : m_poAdapterLayer->m_poLayerDefn->AddGeomFieldDefn(&oGeomFieldDefn);
440 :
441 0 : m_extents.push_back(oMapExtent[pszColName]);
442 0 : m_geomColBBOX.push_back(
443 0 : oMapGeomColumnToCoveringBBOXColumn[pszColName]);
444 : }
445 : else
446 : {
447 0 : m_poAdapterLayer->CreateFieldFromArrowSchema(m_schema.children[i]);
448 : }
449 : }
450 0 : }
451 :
452 : /************************************************************************/
453 : /* ~OGRADBCLayer() */
454 : /************************************************************************/
455 :
456 0 : OGRADBCLayer::~OGRADBCLayer()
457 : {
458 0 : OGRADBCError error;
459 0 : if (m_statement)
460 0 : ADBC_CALL(StatementRelease, m_statement.get(), error);
461 0 : if (m_schema.release)
462 0 : m_schema.release(&m_schema);
463 0 : }
464 :
465 : /************************************************************************/
466 : /* ReplaceStatement() */
467 : /************************************************************************/
468 :
469 0 : bool OGRADBCLayer::ReplaceStatement(const char *pszNewStatement)
470 : {
471 : // CPLDebug("ADBC", "%s", pszNewStatement);
472 0 : OGRADBCError error;
473 0 : auto statement = std::make_unique<AdbcStatement>();
474 0 : if (ADBC_CALL(StatementNew, m_poDS->m_connection.get(), statement.get(),
475 0 : error) != ADBC_STATUS_OK)
476 : {
477 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcStatementNew() failed: %s",
478 : error.message());
479 0 : ADBC_CALL(StatementRelease, statement.get(), error);
480 : }
481 0 : else if (ADBC_CALL(StatementSetSqlQuery, statement.get(), pszNewStatement,
482 0 : error) != ADBC_STATUS_OK)
483 : {
484 0 : CPLError(CE_Failure, CPLE_AppDefined,
485 : "AdbcStatementSetSqlQuery() failed: %s", error.message());
486 0 : error.clear();
487 0 : ADBC_CALL(StatementRelease, statement.get(), error);
488 : }
489 : else
490 : {
491 0 : auto stream = std::make_unique<OGRArrowArrayStream>();
492 0 : int64_t rows_affected = -1;
493 : ArrowSchema newSchema;
494 0 : memset(&newSchema, 0, sizeof(newSchema));
495 0 : if (ADBC_CALL(StatementExecuteQuery, statement.get(), stream->get(),
496 0 : &rows_affected, error) != ADBC_STATUS_OK)
497 : {
498 0 : CPLError(CE_Failure, CPLE_AppDefined,
499 : "AdbcStatementExecuteQuery() failed: %s", error.message());
500 0 : error.clear();
501 0 : ADBC_CALL(StatementRelease, statement.get(), error);
502 : }
503 0 : else if (stream->get_schema(&newSchema) != 0)
504 : {
505 0 : CPLError(CE_Failure, CPLE_AppDefined, "get_schema() failed");
506 0 : ADBC_CALL(StatementRelease, statement.get(), error);
507 : }
508 : else
509 : {
510 0 : if (m_schema.release)
511 0 : m_schema.release(&m_schema);
512 0 : memcpy(&m_schema, &newSchema, sizeof(newSchema));
513 :
514 0 : if (m_statement)
515 0 : ADBC_CALL(StatementRelease, m_statement.get(), error);
516 0 : m_statement = std::move(statement);
517 :
518 0 : m_stream = std::move(stream);
519 :
520 0 : return true;
521 : }
522 : }
523 0 : return false;
524 : }
525 :
526 : /************************************************************************/
527 : /* GetNextRawFeature() */
528 : /************************************************************************/
529 :
530 0 : OGRFeature *OGRADBCLayer::GetNextRawFeature()
531 : {
532 0 : if (m_bEOF)
533 0 : return nullptr;
534 :
535 0 : if (m_nIdx == m_poAdapterLayer->m_apoFeatures.size())
536 : {
537 0 : m_nIdx = 0;
538 0 : m_poAdapterLayer->m_apoFeatures.clear();
539 :
540 0 : if (!m_stream)
541 : {
542 0 : auto stream = std::make_unique<OGRArrowArrayStream>();
543 0 : if (!GetArrowStreamInternal(stream->get()))
544 : {
545 0 : m_bEOF = true;
546 0 : return nullptr;
547 : }
548 0 : m_stream = std::move(stream);
549 : }
550 :
551 : struct ArrowArray array;
552 0 : memset(&array, 0, sizeof(array));
553 0 : if (m_stream->get_next(&array) != 0)
554 : {
555 0 : m_bEOF = true;
556 0 : return nullptr;
557 : }
558 : const bool bOK =
559 0 : array.length
560 0 : ? m_poAdapterLayer->WriteArrowBatch(&m_schema, &array, nullptr)
561 0 : : false;
562 0 : if (array.release)
563 0 : array.release(&array);
564 0 : if (!bOK)
565 : {
566 0 : m_bEOF = true;
567 0 : return nullptr;
568 : }
569 : }
570 :
571 0 : auto poFeature = m_poAdapterLayer->m_apoFeatures[m_nIdx++].release();
572 : const int nGeomFieldCount =
573 0 : m_poAdapterLayer->m_poLayerDefn->GetFieldCount();
574 0 : for (int i = 0; i < nGeomFieldCount; ++i)
575 : {
576 0 : auto poGeom = poFeature->GetGeomFieldRef(i);
577 0 : if (poGeom)
578 0 : poGeom->assignSpatialReference(
579 0 : m_poAdapterLayer->m_poLayerDefn->GetGeomFieldDefn(i)
580 0 : ->GetSpatialRef());
581 : }
582 0 : poFeature->SetFID(m_nFeatureID++);
583 0 : return poFeature;
584 : }
585 :
586 : /************************************************************************/
587 : /* ResetReading() */
588 : /************************************************************************/
589 :
590 0 : void OGRADBCLayer::ResetReading()
591 : {
592 0 : if (m_nIdx > 0 || m_bEOF)
593 : {
594 0 : m_poAdapterLayer->m_apoFeatures.clear();
595 0 : m_stream.reset();
596 0 : m_bEOF = false;
597 0 : m_nIdx = 0;
598 0 : m_nFeatureID = 0;
599 : }
600 0 : }
601 :
602 : /************************************************************************/
603 : /* IGetExtent() */
604 : /************************************************************************/
605 :
606 0 : OGRErr OGRADBCLayer::IGetExtent(int iGeomField, OGREnvelope *psExtent,
607 : bool bForce)
608 : {
609 0 : *psExtent = m_extents[iGeomField];
610 0 : if (psExtent->IsInit())
611 0 : return OGRERR_NONE;
612 :
613 0 : return OGRLayer::IGetExtent(iGeomField, psExtent, bForce);
614 : }
615 :
616 : /************************************************************************/
617 : /* IGetExtent3D() */
618 : /************************************************************************/
619 :
620 0 : OGRErr OGRADBCLayer::IGetExtent3D(int iGeomField, OGREnvelope3D *psExtent,
621 : bool bForce)
622 : {
623 0 : *psExtent = m_extents[iGeomField];
624 0 : if (psExtent->IsInit())
625 0 : return OGRERR_NONE;
626 :
627 0 : return OGRLayer::IGetExtent3D(iGeomField, psExtent, bForce);
628 : }
629 :
630 : /************************************************************************/
631 : /* GetCurrentStatement() */
632 : /************************************************************************/
633 :
634 0 : std::string OGRADBCLayer::GetCurrentStatement() const
635 : {
636 0 : if (!m_osModifiedSelect.empty() &&
637 0 : STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT * FROM ") &&
638 0 : (!m_osAttributeFilter.empty() ||
639 0 : (m_poFilterGeom &&
640 0 : (!m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty() ||
641 0 : m_poDS->m_bSpatialLoaded))))
642 : {
643 0 : std::string osStatement(m_osModifiedSelect);
644 0 : osStatement.append(" FROM (")
645 0 : .append(m_osBaseStatement)
646 0 : .append(") WHERE ");
647 :
648 0 : bool bAddAnd = false;
649 0 : if (m_poFilterGeom)
650 : {
651 0 : const double dfMinX = std::isinf(m_sFilterEnvelope.MinX)
652 0 : ? -std::numeric_limits<double>::max()
653 0 : : m_sFilterEnvelope.MinX;
654 0 : const double dfMinY = std::isinf(m_sFilterEnvelope.MinY)
655 0 : ? -std::numeric_limits<double>::max()
656 0 : : m_sFilterEnvelope.MinY;
657 0 : const double dfMaxX = std::isinf(m_sFilterEnvelope.MaxX)
658 0 : ? std::numeric_limits<double>::max()
659 0 : : m_sFilterEnvelope.MaxX;
660 0 : const double dfMaxY = std::isinf(m_sFilterEnvelope.MaxY)
661 0 : ? std::numeric_limits<double>::max()
662 0 : : m_sFilterEnvelope.MaxY;
663 0 : if (!m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty())
664 : {
665 0 : bAddAnd = true;
666 0 : osStatement.append(m_geomColBBOX[m_iGeomFieldFilter].osXMin)
667 0 : .append(" <= ")
668 0 : .append(CPLSPrintf("%.17g", dfMaxX))
669 0 : .append(" AND ")
670 0 : .append(m_geomColBBOX[m_iGeomFieldFilter].osXMax)
671 0 : .append(" >= ")
672 0 : .append(CPLSPrintf("%.17g", dfMinX))
673 0 : .append(" AND ")
674 0 : .append(m_geomColBBOX[m_iGeomFieldFilter].osYMin)
675 0 : .append(" <= ")
676 0 : .append(CPLSPrintf("%.17g", dfMaxY))
677 0 : .append(" AND ")
678 0 : .append(m_geomColBBOX[m_iGeomFieldFilter].osYMax)
679 0 : .append(" >= ")
680 0 : .append(CPLSPrintf("%.17g", dfMinY));
681 : }
682 0 : if (m_poDS->m_bSpatialLoaded)
683 : {
684 0 : if (bAddAnd)
685 0 : osStatement.append(" AND ");
686 0 : bAddAnd = true;
687 0 : osStatement.append("ST_Intersects(\"")
688 0 : .append(OGRDuplicateCharacter(
689 0 : m_poAdapterLayer->m_poLayerDefn
690 0 : ->GetGeomFieldDefn(m_iGeomFieldFilter)
691 : ->GetNameRef(),
692 0 : '"'))
693 : .append(CPLSPrintf(
694 : "\", ST_MakeEnvelope(%.17g,%.17g,%.17g,%.17g))", dfMinX,
695 0 : dfMinY, dfMaxX, dfMaxY));
696 : }
697 : }
698 0 : if (!m_osAttributeFilter.empty())
699 : {
700 0 : if (bAddAnd)
701 0 : osStatement.append(" AND ");
702 0 : osStatement.append("(");
703 0 : osStatement.append(m_osAttributeFilter);
704 0 : osStatement.append(")");
705 : }
706 :
707 0 : return osStatement;
708 : }
709 : else
710 : {
711 0 : return m_osModifiedBaseStatement;
712 : }
713 : }
714 :
715 : /************************************************************************/
716 : /* UpdateStatement() */
717 : /************************************************************************/
718 :
719 0 : bool OGRADBCLayer::UpdateStatement()
720 : {
721 0 : return ReplaceStatement(GetCurrentStatement().c_str());
722 : }
723 :
724 : /***********************************************************************/
725 : /* SetAttributeFilter() */
726 : /***********************************************************************/
727 :
728 0 : OGRErr OGRADBCLayer::SetAttributeFilter(const char *pszFilter)
729 : {
730 0 : if (!m_osModifiedSelect.empty() &&
731 0 : STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT * FROM "))
732 : {
733 0 : m_osAttributeFilter = pszFilter ? pszFilter : "";
734 0 : return UpdateStatement() ? OGRERR_NONE : OGRERR_FAILURE;
735 : }
736 : else
737 : {
738 0 : return OGRLayer::SetAttributeFilter(pszFilter);
739 : }
740 : }
741 :
742 : /************************************************************************/
743 : /* ISetSpatialFilter() */
744 : /************************************************************************/
745 :
746 0 : OGRErr OGRADBCLayer::ISetSpatialFilter(int iGeomField,
747 : const OGRGeometry *poGeomIn)
748 :
749 : {
750 0 : if (iGeomField < GetLayerDefn()->GetGeomFieldCount())
751 : {
752 0 : m_iGeomFieldFilter = iGeomField;
753 0 : if (InstallFilter(poGeomIn))
754 0 : ResetReading();
755 0 : UpdateStatement();
756 : }
757 0 : return OGRERR_NONE;
758 : }
759 :
760 : /************************************************************************/
761 : /* TestCapability() */
762 : /************************************************************************/
763 :
764 0 : int OGRADBCLayer::TestCapability(const char *pszCap)
765 : {
766 0 : if (EQUAL(pszCap, OLCFastGetArrowStream))
767 : {
768 0 : return !m_poFilterGeom && !m_poAttrQuery && m_osAttributeFilter.empty();
769 : }
770 0 : else if (EQUAL(pszCap, OLCFastFeatureCount))
771 : {
772 0 : return !m_poFilterGeom && !m_poAttrQuery &&
773 0 : m_osAttributeFilter.empty() && m_bIsParquetLayer;
774 : }
775 0 : else if (EQUAL(pszCap, OLCFastGetExtent))
776 : {
777 0 : return !m_extents.empty() && m_extents[0].IsInit();
778 : }
779 0 : else if (EQUAL(pszCap, OLCFastSpatialFilter) && m_iGeomFieldFilter >= 0 &&
780 0 : m_iGeomFieldFilter < GetLayerDefn()->GetGeomFieldCount())
781 : {
782 0 : if (m_poDS->m_bSpatialLoaded && m_poDS->m_bIsDuckDBDataset)
783 : {
784 : const char *pszGeomColName =
785 0 : m_poAdapterLayer->m_poLayerDefn
786 0 : ->GetGeomFieldDefn(m_iGeomFieldFilter)
787 0 : ->GetNameRef();
788 0 : auto poTmpLayer = m_poDS->CreateInternalLayer(CPLSPrintf(
789 : "SELECT 1 FROM sqlite_master WHERE tbl_name = '%s' AND type = "
790 : "'index' AND (sql LIKE '%%USING RTREE (%s)%%' OR sql LIKE "
791 : "'%%USING RTREE (\"%s\")%%')",
792 0 : OGRDuplicateCharacter(GetDescription(), '\'').c_str(),
793 : pszGeomColName,
794 0 : OGRDuplicateCharacter(pszGeomColName, '"').c_str()));
795 0 : return poTmpLayer &&
796 0 : std::unique_ptr<OGRFeature>(poTmpLayer->GetNextFeature());
797 : }
798 0 : else if (!m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty())
799 : {
800 : // Let's assume that the presence of a geometry bounding box
801 : // column is sufficient enough to pretend to have fast spatial
802 : // filter capabilities
803 0 : return true;
804 : }
805 : }
806 :
807 0 : return false;
808 : }
809 :
810 : /************************************************************************/
811 : /* GetDataset() */
812 : /************************************************************************/
813 :
814 0 : GDALDataset *OGRADBCLayer::GetDataset()
815 : {
816 0 : return m_poDS;
817 : }
818 :
819 : /************************************************************************/
820 : /* GetArrowStream() */
821 : /************************************************************************/
822 :
823 0 : bool OGRADBCLayer::GetArrowStream(struct ArrowArrayStream *out_stream,
824 : CSLConstList papszOptions)
825 : {
826 0 : if (m_poFilterGeom || m_poAttrQuery ||
827 0 : CPLFetchBool(papszOptions, GAS_OPT_DATETIME_AS_STRING, false))
828 : {
829 0 : return OGRLayer::GetArrowStream(out_stream, papszOptions);
830 : }
831 :
832 0 : if (m_stream)
833 : {
834 0 : memcpy(out_stream, m_stream->get(), sizeof(*out_stream));
835 0 : memset(m_stream->get(), 0, sizeof(*out_stream));
836 0 : m_stream.reset();
837 : }
838 :
839 0 : return GetArrowStreamInternal(out_stream);
840 : }
841 :
842 : /************************************************************************/
843 : /* GetArrowStreamInternal() */
844 : /************************************************************************/
845 :
846 0 : bool OGRADBCLayer::GetArrowStreamInternal(struct ArrowArrayStream *out_stream)
847 : {
848 0 : OGRADBCError error;
849 0 : int64_t rows_affected = -1;
850 0 : if (ADBC_CALL(StatementExecuteQuery, m_statement.get(), out_stream,
851 0 : &rows_affected, error) != ADBC_STATUS_OK)
852 : {
853 0 : CPLError(CE_Failure, CPLE_AppDefined,
854 : "AdbcStatementExecuteQuery() failed: %s", error.message());
855 0 : return false;
856 : }
857 :
858 0 : return true;
859 : }
860 :
861 : /************************************************************************/
862 : /* GetFeatureCount() */
863 : /************************************************************************/
864 :
865 0 : GIntBig OGRADBCLayer::GetFeatureCount(int bForce)
866 : {
867 0 : if (m_poFilterGeom || m_poAttrQuery || !m_osAttributeFilter.empty())
868 : {
869 0 : if (!m_osModifiedSelect.empty() &&
870 0 : STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT * FROM ") &&
871 0 : (!m_poFilterGeom ||
872 0 : !m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty() ||
873 0 : m_poDS->m_bSpatialLoaded))
874 : {
875 0 : const std::string osCurStatement = GetCurrentStatement();
876 0 : auto poCountLayer = m_poDS->CreateInternalLayer(
877 0 : std::string("SELECT COUNT(*) FROM (")
878 0 : .append(osCurStatement)
879 0 : .append(")")
880 0 : .c_str());
881 0 : if (poCountLayer &&
882 0 : poCountLayer->GetLayerDefn()->GetFieldCount() == 1)
883 : {
884 : auto poFeature =
885 0 : std::unique_ptr<OGRFeature>(poCountLayer->GetNextFeature());
886 0 : if (poFeature)
887 0 : return poFeature->GetFieldAsInteger64(0);
888 : }
889 : }
890 :
891 0 : return OGRLayer::GetFeatureCount(bForce);
892 : }
893 :
894 0 : if (m_bIsParquetLayer)
895 : {
896 0 : return GetFeatureCountParquet();
897 : }
898 :
899 0 : if (m_nIdx > 0 || m_bEOF)
900 0 : m_stream.reset();
901 :
902 0 : if (!m_stream)
903 : {
904 0 : auto stream = std::make_unique<OGRArrowArrayStream>();
905 0 : if (!GetArrowStreamInternal(stream->get()))
906 : {
907 0 : return -1;
908 : }
909 0 : m_stream = std::move(stream);
910 : }
911 :
912 0 : GIntBig nTotal = 0;
913 : while (true)
914 : {
915 : struct ArrowArray array;
916 0 : memset(&array, 0, sizeof(array));
917 0 : if (m_stream->get_next(&array) != 0)
918 : {
919 0 : m_stream.reset();
920 0 : return -1;
921 : }
922 0 : const bool bStop = array.length == 0;
923 0 : nTotal += array.length;
924 0 : if (array.release)
925 0 : array.release(&array);
926 0 : if (bStop)
927 0 : break;
928 0 : }
929 0 : m_stream.reset();
930 0 : return nTotal;
931 : }
932 :
933 : /************************************************************************/
934 : /* GetFeatureCountParquet() */
935 : /************************************************************************/
936 :
937 0 : GIntBig OGRADBCLayer::GetFeatureCountParquet()
938 : {
939 : const std::string osSQL(CPLSPrintf(
940 : "SELECT CAST(SUM(num_rows) AS BIGINT) FROM parquet_file_metadata('%s')",
941 0 : OGRDuplicateCharacter(m_poDS->m_osParquetFilename, '\'').c_str()));
942 0 : auto poCountLayer = m_poDS->CreateInternalLayer(osSQL.c_str());
943 0 : if (poCountLayer && poCountLayer->GetLayerDefn()->GetFieldCount() == 1)
944 : {
945 : auto poFeature =
946 0 : std::unique_ptr<OGRFeature>(poCountLayer->GetNextFeature());
947 0 : if (poFeature)
948 0 : return poFeature->GetFieldAsInteger64(0);
949 : }
950 :
951 0 : return -1;
952 : }
953 :
954 : #undef ADBC_CALL
|
