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