Merge pull request OSGeo#11536 from rouault/adbc_parquet_sql

ADBC: make ogrinfo ADBC: -oo SQL="SELECT * FROM '/path/to/my.parquet [LIMIT x]" work
innFactory · Dec 25, 2024 · 39d8ed6 · 39d8ed6
2 parents f314d4b + 7796bf7
commit 39d8ed6
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 6 deletions.
diff --git a/autotest/ogr/ogr_adbc.py b/autotest/ogr/ogr_adbc.py
@@ -306,6 +306,29 @@ def test_ogr_adbc_duckdb_parquet_with_spatial(OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL)
 ###############################################################################
 
 
+@pytest.mark.parametrize("OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", ["ON", "OFF"])
+def test_ogr_adbc_duckdb_parquet_with_spatial_and_SQL_open_optoin(
+    OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL,
+):
+
+    if not _has_libduckdb():
+        pytest.skip("libduckdb.so missing")
+
+    with gdal.config_option(
+        "OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL
+    ):
+        open_options = ["SQL=SELECT * FROM 'data/parquet/poly.parquet' LIMIT 1"]
+        if OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL == "ON":
+            open_options += ["PRELUDE_STATEMENTS=INSTALL spatial"]
+        with gdal.OpenEx("ADBC:", gdal.OF_VECTOR, open_options=open_options) as ds:
+            lyr = ds.GetLayer(0)
+            assert lyr.GetGeomType() == ogr.wkbPolygon
+            assert lyr.GetFeatureCount() == 1
+
+
+###############################################################################
+
+
 @pytest.mark.parametrize("OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", ["ON", "OFF"])
 def test_ogr_adbc_duckdb_with_spatial_index(OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL):
 

diff --git a/ogr/ogrsf_frmts/adbc/ogradbcdataset.cpp b/ogr/ogrsf_frmts/adbc/ogradbcdataset.cpp
@@ -218,8 +218,28 @@ bool OGRADBCDataset::Open(const GDALOpenInfo *poOpenInfo)
     const bool bIsSQLite3 =
         (pszADBCDriverName && EQUAL(pszADBCDriverName, "adbc_driver_sqlite")) ||
         OGRADBCDriverIsSQLite3(poOpenInfo);
-    const bool bIsParquet = OGRADBCDriverIsParquet(poOpenInfo) ||
-                            EQUAL(CPLGetExtension(pszFilename), "parquet");
+    bool bIsParquet = OGRADBCDriverIsParquet(poOpenInfo) ||
+                      EQUAL(CPLGetExtension(pszFilename), "parquet");
+    const char *pszSQL = CSLFetchNameValue(poOpenInfo->papszOpenOptions, "SQL");
+    if (!bIsParquet && pszSQL)
+    {
+        CPLString osSQL(pszSQL);
+        auto iPos = osSQL.find("FROM '");
+        if (iPos != std::string::npos)
+        {
+            iPos += strlen("FROM '");
+            const auto iPos2 = osSQL.find("'", iPos);
+            if (iPos2 != std::string::npos)
+            {
+                const std::string osFilename = osSQL.substr(iPos, iPos2 - iPos);
+                if (EQUAL(CPLGetExtension(osFilename.c_str()), "parquet"))
+                {
+                    m_osParquetFilename = osFilename;
+                    bIsParquet = true;
+                }
+            }
+        }
+    }
     const bool bIsPostgreSQL = STARTS_WITH(pszFilename, "postgresql://");
 
     if (!pszADBCDriverName)
@@ -376,14 +396,15 @@ bool OGRADBCDataset::Open(const GDALOpenInfo *poOpenInfo)
 
     std::string osLayerName = "RESULTSET";
     std::string osSQL;
-    const char *pszSQL = CSLFetchNameValue(poOpenInfo->papszOpenOptions, "SQL");
     bool bIsParquetLayer = false;
     if (bIsParquet)
     {
-        m_osParquetFilename = pszFilename;
-        osLayerName = CPLGetBasename(pszFilename);
+        if (m_osParquetFilename.empty())
+            m_osParquetFilename = pszFilename;
+        osLayerName = CPLGetBasename(m_osParquetFilename.c_str());
         if (osLayerName == "*")
-            osLayerName = CPLGetBasename(CPLGetDirname(pszFilename));
+            osLayerName =
+                CPLGetBasename(CPLGetDirname(m_osParquetFilename.c_str()));
         if (!pszSQL)
         {
             osSQL =