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 : /* IGetExtent() */
599 : /************************************************************************/
600 :
601 0 : OGRErr OGRADBCLayer::IGetExtent(int iGeomField, OGREnvelope *psExtent,
602 : bool bForce)
603 : {
604 0 : *psExtent = m_extents[iGeomField];
605 0 : if (psExtent->IsInit())
606 0 : return OGRERR_NONE;
607 :
608 0 : return OGRLayer::IGetExtent(iGeomField, psExtent, bForce);
609 : }
610 :
611 : /************************************************************************/
612 : /* IGetExtent3D() */
613 : /************************************************************************/
614 :
615 0 : OGRErr OGRADBCLayer::IGetExtent3D(int iGeomField, OGREnvelope3D *psExtent,
616 : bool bForce)
617 : {
618 0 : *psExtent = m_extents[iGeomField];
619 0 : if (psExtent->IsInit())
620 0 : return OGRERR_NONE;
621 :
622 0 : return OGRLayer::IGetExtent3D(iGeomField, psExtent, bForce);
623 : }
624 :
625 : /************************************************************************/
626 : /* GetCurrentStatement() */
627 : /************************************************************************/
628 :
629 0 : std::string OGRADBCLayer::GetCurrentStatement() const
630 : {
631 0 : if (!m_osModifiedSelect.empty() &&
632 0 : STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT * FROM ") &&
633 0 : (!m_osAttributeFilter.empty() ||
634 0 : (m_poFilterGeom &&
635 0 : (!m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty() ||
636 0 : m_poDS->m_bSpatialLoaded))))
637 : {
638 0 : std::string osStatement(m_osModifiedSelect);
639 0 : osStatement.append(" FROM (")
640 0 : .append(m_osBaseStatement)
641 0 : .append(") WHERE ");
642 :
643 0 : bool bAddAnd = false;
644 0 : if (m_poFilterGeom)
645 : {
646 0 : const double dfMinX = std::isinf(m_sFilterEnvelope.MinX)
647 0 : ? -std::numeric_limits<double>::max()
648 0 : : m_sFilterEnvelope.MinX;
649 0 : const double dfMinY = std::isinf(m_sFilterEnvelope.MinY)
650 0 : ? -std::numeric_limits<double>::max()
651 0 : : m_sFilterEnvelope.MinY;
652 0 : const double dfMaxX = std::isinf(m_sFilterEnvelope.MaxX)
653 0 : ? std::numeric_limits<double>::max()
654 0 : : m_sFilterEnvelope.MaxX;
655 0 : const double dfMaxY = std::isinf(m_sFilterEnvelope.MaxY)
656 0 : ? std::numeric_limits<double>::max()
657 0 : : m_sFilterEnvelope.MaxY;
658 0 : if (!m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty())
659 : {
660 0 : bAddAnd = true;
661 0 : osStatement.append(m_geomColBBOX[m_iGeomFieldFilter].osXMin)
662 0 : .append(" <= ")
663 0 : .append(CPLSPrintf("%.17g", dfMaxX))
664 0 : .append(" AND ")
665 0 : .append(m_geomColBBOX[m_iGeomFieldFilter].osXMax)
666 0 : .append(" >= ")
667 0 : .append(CPLSPrintf("%.17g", dfMinX))
668 0 : .append(" AND ")
669 0 : .append(m_geomColBBOX[m_iGeomFieldFilter].osYMin)
670 0 : .append(" <= ")
671 0 : .append(CPLSPrintf("%.17g", dfMaxY))
672 0 : .append(" AND ")
673 0 : .append(m_geomColBBOX[m_iGeomFieldFilter].osYMax)
674 0 : .append(" >= ")
675 0 : .append(CPLSPrintf("%.17g", dfMinY));
676 : }
677 0 : if (m_poDS->m_bSpatialLoaded)
678 : {
679 0 : if (bAddAnd)
680 0 : osStatement.append(" AND ");
681 0 : bAddAnd = true;
682 0 : osStatement.append("ST_Intersects(\"")
683 0 : .append(OGRDuplicateCharacter(
684 0 : m_poAdapterLayer->m_poLayerDefn
685 0 : ->GetGeomFieldDefn(m_iGeomFieldFilter)
686 : ->GetNameRef(),
687 0 : '"'))
688 : .append(CPLSPrintf(
689 : "\", ST_MakeEnvelope(%.17g,%.17g,%.17g,%.17g))", dfMinX,
690 0 : dfMinY, dfMaxX, dfMaxY));
691 : }
692 : }
693 0 : if (!m_osAttributeFilter.empty())
694 : {
695 0 : if (bAddAnd)
696 0 : osStatement.append(" AND ");
697 0 : osStatement.append("(");
698 0 : osStatement.append(m_osAttributeFilter);
699 0 : osStatement.append(")");
700 : }
701 :
702 0 : return osStatement;
703 : }
704 : else
705 : {
706 0 : return m_osModifiedBaseStatement;
707 : }
708 : }
709 :
710 : /************************************************************************/
711 : /* UpdateStatement() */
712 : /************************************************************************/
713 :
714 0 : bool OGRADBCLayer::UpdateStatement()
715 : {
716 0 : return ReplaceStatement(GetCurrentStatement().c_str());
717 : }
718 :
719 : /***********************************************************************/
720 : /* SetAttributeFilter() */
721 : /***********************************************************************/
722 :
723 0 : OGRErr OGRADBCLayer::SetAttributeFilter(const char *pszFilter)
724 : {
725 0 : if (!m_osModifiedSelect.empty() &&
726 0 : STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT * FROM "))
727 : {
728 0 : m_osAttributeFilter = pszFilter ? pszFilter : "";
729 0 : return UpdateStatement() ? OGRERR_NONE : OGRERR_FAILURE;
730 : }
731 : else
732 : {
733 0 : return OGRLayer::SetAttributeFilter(pszFilter);
734 : }
735 : }
736 :
737 : /************************************************************************/
738 : /* ISetSpatialFilter() */
739 : /************************************************************************/
740 :
741 0 : OGRErr OGRADBCLayer::ISetSpatialFilter(int iGeomField,
742 : const OGRGeometry *poGeomIn)
743 :
744 : {
745 0 : if (iGeomField < GetLayerDefn()->GetGeomFieldCount())
746 : {
747 0 : m_iGeomFieldFilter = iGeomField;
748 0 : if (InstallFilter(poGeomIn))
749 0 : ResetReading();
750 0 : UpdateStatement();
751 : }
752 0 : return OGRERR_NONE;
753 : }
754 :
755 : /************************************************************************/
756 : /* TestCapability() */
757 : /************************************************************************/
758 :
759 0 : int OGRADBCLayer::TestCapability(const char *pszCap)
760 : {
761 0 : if (EQUAL(pszCap, OLCFastGetArrowStream))
762 : {
763 0 : return !m_poFilterGeom && !m_poAttrQuery && m_osAttributeFilter.empty();
764 : }
765 0 : else if (EQUAL(pszCap, OLCFastFeatureCount))
766 : {
767 0 : return !m_poFilterGeom && !m_poAttrQuery &&
768 0 : m_osAttributeFilter.empty() && m_bIsParquetLayer;
769 : }
770 0 : else if (EQUAL(pszCap, OLCFastGetExtent))
771 : {
772 0 : return !m_extents.empty() && m_extents[0].IsInit();
773 : }
774 0 : else if (EQUAL(pszCap, OLCFastSpatialFilter) && m_iGeomFieldFilter >= 0 &&
775 0 : m_iGeomFieldFilter < GetLayerDefn()->GetGeomFieldCount())
776 : {
777 0 : if (m_poDS->m_bSpatialLoaded && m_poDS->m_bIsDuckDBDataset)
778 : {
779 : const char *pszGeomColName =
780 0 : m_poAdapterLayer->m_poLayerDefn
781 0 : ->GetGeomFieldDefn(m_iGeomFieldFilter)
782 0 : ->GetNameRef();
783 0 : auto poTmpLayer = m_poDS->CreateInternalLayer(CPLSPrintf(
784 : "SELECT 1 FROM sqlite_master WHERE tbl_name = '%s' AND type = "
785 : "'index' AND (sql LIKE '%%USING RTREE (%s)%%' OR sql LIKE "
786 : "'%%USING RTREE (\"%s\")%%')",
787 0 : OGRDuplicateCharacter(GetDescription(), '\'').c_str(),
788 : pszGeomColName,
789 0 : OGRDuplicateCharacter(pszGeomColName, '"').c_str()));
790 0 : return poTmpLayer &&
791 0 : std::unique_ptr<OGRFeature>(poTmpLayer->GetNextFeature());
792 : }
793 0 : else if (!m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty())
794 : {
795 : // Let's assume that the presence of a geometry bounding box
796 : // column is sufficient enough to pretend to have fast spatial
797 : // filter capabilities
798 0 : return true;
799 : }
800 : }
801 :
802 0 : return false;
803 : }
804 :
805 : /************************************************************************/
806 : /* GetDataset() */
807 : /************************************************************************/
808 :
809 0 : GDALDataset *OGRADBCLayer::GetDataset()
810 : {
811 0 : return m_poDS;
812 : }
813 :
814 : /************************************************************************/
815 : /* GetArrowStream() */
816 : /************************************************************************/
817 :
818 0 : bool OGRADBCLayer::GetArrowStream(struct ArrowArrayStream *out_stream,
819 : CSLConstList papszOptions)
820 : {
821 0 : if (m_poFilterGeom || m_poAttrQuery ||
822 0 : CPLFetchBool(papszOptions, GAS_OPT_DATETIME_AS_STRING, false))
823 : {
824 0 : return OGRLayer::GetArrowStream(out_stream, papszOptions);
825 : }
826 :
827 0 : if (m_stream)
828 : {
829 0 : memcpy(out_stream, m_stream->get(), sizeof(*out_stream));
830 0 : memset(m_stream->get(), 0, sizeof(*out_stream));
831 0 : m_stream.reset();
832 : }
833 :
834 0 : return GetArrowStreamInternal(out_stream);
835 : }
836 :
837 : /************************************************************************/
838 : /* GetArrowStreamInternal() */
839 : /************************************************************************/
840 :
841 0 : bool OGRADBCLayer::GetArrowStreamInternal(struct ArrowArrayStream *out_stream)
842 : {
843 0 : OGRADBCError error;
844 0 : int64_t rows_affected = -1;
845 0 : if (ADBC_CALL(StatementExecuteQuery, m_statement.get(), out_stream,
846 0 : &rows_affected, error) != ADBC_STATUS_OK)
847 : {
848 0 : CPLError(CE_Failure, CPLE_AppDefined,
849 : "AdbcStatementExecuteQuery() failed: %s", error.message());
850 0 : return false;
851 : }
852 :
853 0 : return true;
854 : }
855 :
856 : /************************************************************************/
857 : /* GetFeatureCount() */
858 : /************************************************************************/
859 :
860 0 : GIntBig OGRADBCLayer::GetFeatureCount(int bForce)
861 : {
862 0 : if (m_poFilterGeom || m_poAttrQuery || !m_osAttributeFilter.empty())
863 : {
864 0 : if (!m_osModifiedSelect.empty() &&
865 0 : STARTS_WITH_CI(m_osBaseStatement.c_str(), "SELECT * FROM ") &&
866 0 : (!m_poFilterGeom ||
867 0 : !m_geomColBBOX[m_iGeomFieldFilter].osXMin.empty() ||
868 0 : m_poDS->m_bSpatialLoaded))
869 : {
870 0 : const std::string osCurStatement = GetCurrentStatement();
871 0 : auto poCountLayer = m_poDS->CreateInternalLayer(
872 0 : std::string("SELECT COUNT(*) FROM (")
873 0 : .append(osCurStatement)
874 0 : .append(")")
875 0 : .c_str());
876 0 : if (poCountLayer &&
877 0 : poCountLayer->GetLayerDefn()->GetFieldCount() == 1)
878 : {
879 : auto poFeature =
880 0 : std::unique_ptr<OGRFeature>(poCountLayer->GetNextFeature());
881 0 : if (poFeature)
882 0 : return poFeature->GetFieldAsInteger64(0);
883 : }
884 : }
885 :
886 0 : return OGRLayer::GetFeatureCount(bForce);
887 : }
888 :
889 0 : if (m_bIsParquetLayer)
890 : {
891 0 : return GetFeatureCountParquet();
892 : }
893 :
894 0 : if (m_nIdx > 0 || m_bEOF)
895 0 : m_stream.reset();
896 :
897 0 : if (!m_stream)
898 : {
899 0 : auto stream = std::make_unique<OGRArrowArrayStream>();
900 0 : if (!GetArrowStreamInternal(stream->get()))
901 : {
902 0 : return -1;
903 : }
904 0 : m_stream = std::move(stream);
905 : }
906 :
907 0 : GIntBig nTotal = 0;
908 : while (true)
909 : {
910 : struct ArrowArray array;
911 0 : memset(&array, 0, sizeof(array));
912 0 : if (m_stream->get_next(&array) != 0)
913 : {
914 0 : m_stream.reset();
915 0 : return -1;
916 : }
917 0 : const bool bStop = array.length == 0;
918 0 : nTotal += array.length;
919 0 : if (array.release)
920 0 : array.release(&array);
921 0 : if (bStop)
922 0 : break;
923 0 : }
924 0 : m_stream.reset();
925 0 : return nTotal;
926 : }
927 :
928 : /************************************************************************/
929 : /* GetFeatureCountParquet() */
930 : /************************************************************************/
931 :
932 0 : GIntBig OGRADBCLayer::GetFeatureCountParquet()
933 : {
934 : const std::string osSQL(CPLSPrintf(
935 : "SELECT CAST(SUM(num_rows) AS BIGINT) FROM parquet_file_metadata('%s')",
936 0 : OGRDuplicateCharacter(m_poDS->m_osParquetFilename, '\'').c_str()));
937 0 : auto poCountLayer = m_poDS->CreateInternalLayer(osSQL.c_str());
938 0 : if (poCountLayer && poCountLayer->GetLayerDefn()->GetFieldCount() == 1)
939 : {
940 : auto poFeature =
941 0 : std::unique_ptr<OGRFeature>(poCountLayer->GetNextFeature());
942 0 : if (poFeature)
943 0 : return poFeature->GetFieldAsInteger64(0);
944 : }
945 :
946 0 : return -1;
947 : }
948 :
949 : #undef ADBC_CALL
|
