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