From 80b73c716b4e8ffad99cd2391658fed302c71f13 Mon Sep 17 00:00:00 2001 From: Even Rouault Date: Thu, 19 Oct 2023 22:11:46 +0200 Subject: [PATCH 1/5] TileDB: fix attribute filtering with GetArrowStream() --- frmts/tiledb/tiledbsparse.cpp | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/frmts/tiledb/tiledbsparse.cpp b/frmts/tiledb/tiledbsparse.cpp index a30c67a6789e..d8c5bb4aabee 100644 --- a/frmts/tiledb/tiledbsparse.cpp +++ b/frmts/tiledb/tiledbsparse.cpp @@ -5104,7 +5104,7 @@ void OGRTileDBLayer::FillBoolListArray( /* GetNextArrowArray() */ /************************************************************************/ -int OGRTileDBLayer::GetNextArrowArray(struct ArrowArrayStream *, +int OGRTileDBLayer::GetNextArrowArray(struct ArrowArrayStream *stream, struct ArrowArray *out_array) { memset(out_array, 0, sizeof(*out_array)); @@ -5505,6 +5505,22 @@ int OGRTileDBLayer::GetNextArrowArray(struct ArrowArrayStream *, } } CPL_IGNORE_RET_VAL(iSchemaChild); + + if (m_poAttrQuery && + (!m_poQueryCondition || m_bAttributeFilterPartiallyTranslated)) + { + struct ArrowSchema schema; + stream->get_schema(stream, &schema); + CPLAssert(schema.release != nullptr); + CPLAssert(schema.n_children == out_array->n_children); + // Spatial filter already evaluated + auto poFilterGeomBackup = m_poFilterGeom; + m_poFilterGeom = nullptr; + if (CanPostFilterArrowArray(&schema)) + PostFilterArrowArray(&schema, out_array, nullptr); + schema.release(&schema); + m_poFilterGeom = poFilterGeomBackup; + } } catch (const std::exception &e) { From 6496fc69117987d4ab301fde9b2d818a3f9cf37e Mon Sep 17 00:00:00 2001 From: Even Rouault Date: Thu, 19 Oct 2023 22:12:08 +0200 Subject: [PATCH 2/5] Arrow/Parquet: emit error when SetSpatialFilter() called on invalid geom field index --- ogr/ogrsf_frmts/arrow_common/ograrrowlayer.hpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/ogr/ogrsf_frmts/arrow_common/ograrrowlayer.hpp b/ogr/ogrsf_frmts/arrow_common/ograrrowlayer.hpp index d8bf784dc0f9..a90848022ab9 100644 --- a/ogr/ogrsf_frmts/arrow_common/ograrrowlayer.hpp +++ b/ogr/ogrsf_frmts/arrow_common/ograrrowlayer.hpp @@ -3414,6 +3414,14 @@ inline void OGRArrowLayer::SetSpatialFilter(int iGeomField, OGRGeometry *poGeomIn) { + if (iGeomField < 0 || (iGeomField >= GetLayerDefn()->GetGeomFieldCount() && + !(iGeomField == 0 && poGeomIn == nullptr))) + { + CPLError(CE_Failure, CPLE_AppDefined, + "Invalid geometry field index : %d", iGeomField); + return; + } + // When changing filters, we need to invalidate cached batches, as // PostFilterArrowArray() has potentially modified array contents if (m_poFilterGeom) From 72968f9af98b16b9ce1b71b400b0d9bf30c33dad Mon Sep 17 00:00:00 2001 From: Even Rouault Date: Thu, 19 Oct 2023 23:42:05 +0200 Subject: [PATCH 3/5] OGRLayer::GetNextArrowArray(): make it robust against broken GetNextFeature() --- ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp | 15 ++++++++++++++- ogr/ogrsf_frmts/ogrsf_frmts.h | 1 + 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp b/ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp index 2d80eee0826c..21809acb49cc 100644 --- a/ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp +++ b/ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp @@ -1549,9 +1549,13 @@ FillDateTimeArray(struct ArrowArray *psChild, * * @since GDAL 3.6 */ -int OGRLayer::GetNextArrowArray(struct ArrowArrayStream *, +int OGRLayer::GetNextArrowArray(struct ArrowArrayStream *stream, struct ArrowArray *out_array) { + ArrowArrayStreamPrivateDataSharedDataWrapper *poPrivate = + static_cast( + stream->private_data); + const bool bIncludeFID = CPLTestBool( m_aosArrowArrayStreamOptions.FetchNameValueDef("INCLUDE_FID", "YES")); int nMaxBatchSize = atoi(m_aosArrowArrayStreamOptions.FetchNameValueDef( @@ -1573,6 +1577,11 @@ int OGRLayer::GetNextArrowArray(struct ArrowArrayStream *, } memset(out_array, 0, sizeof(*out_array)); + if (poPrivate->poShared->m_bEOF) + { + return 0; + } + auto poLayerDefn = GetLayerDefn(); const int nFieldCount = poLayerDefn->GetFieldCount(); const int nGeomFieldCount = poLayerDefn->GetGeomFieldCount(); @@ -1586,7 +1595,10 @@ int OGRLayer::GetNextArrowArray(struct ArrowArrayStream *, { auto poFeature = std::unique_ptr(GetNextFeature()); if (!poFeature) + { + poPrivate->poShared->m_bEOF = true; break; + } apoFeatures.emplace_back(std::move(poFeature)); } if (apoFeatures.empty()) @@ -1914,6 +1926,7 @@ void OGRLayer::ReleaseStream(struct ArrowArrayStream *stream) static_cast( stream->private_data); poPrivate->poShared->m_bArrowArrayStreamInProgress = false; + poPrivate->poShared->m_bEOF = false; if (poPrivate->poShared->m_poLayer) poPrivate->poShared->m_poLayer->ResetReading(); delete poPrivate; diff --git a/ogr/ogrsf_frmts/ogrsf_frmts.h b/ogr/ogrsf_frmts/ogrsf_frmts.h index 75b5a0f1b707..c1cf326640fc 100644 --- a/ogr/ogrsf_frmts/ogrsf_frmts.h +++ b/ogr/ogrsf_frmts/ogrsf_frmts.h @@ -130,6 +130,7 @@ class CPL_DLL OGRLayer : public GDALMajorObject struct ArrowArrayStreamPrivateData { bool m_bArrowArrayStreamInProgress = false; + bool m_bEOF = false; OGRLayer *m_poLayer = nullptr; }; std::shared_ptr From 0df559784fa3ff1fe384888c866b8dba2ad7f349 Mon Sep 17 00:00:00 2001 From: Even Rouault Date: Fri, 20 Oct 2023 01:52:34 +0200 Subject: [PATCH 4/5] OGREditableLayer: do not forward GetArrowStream() to decorated layer --- ogr/ogrsf_frmts/generic/ogreditablelayer.cpp | 10 ++++++++++ ogr/ogrsf_frmts/generic/ogreditablelayer.h | 2 ++ 2 files changed, 12 insertions(+) diff --git a/ogr/ogrsf_frmts/generic/ogreditablelayer.cpp b/ogr/ogrsf_frmts/generic/ogreditablelayer.cpp index 44b9f8dac2e1..22b82cdeb6ac 100644 --- a/ogr/ogrsf_frmts/generic/ogreditablelayer.cpp +++ b/ogr/ogrsf_frmts/generic/ogreditablelayer.cpp @@ -552,6 +552,16 @@ OGRErr OGREditableLayer::SetAttributeFilter(const char *poAttrFilter) return OGRLayer::SetAttributeFilter(poAttrFilter); } +/************************************************************************/ +/* GetArrowStream() */ +/************************************************************************/ + +bool OGREditableLayer::GetArrowStream(struct ArrowArrayStream *out_stream, + CSLConstList papszOptions) +{ + return OGRLayer::GetArrowStream(out_stream, papszOptions); +} + /************************************************************************/ /* SetSpatialFilter() */ /************************************************************************/ diff --git a/ogr/ogrsf_frmts/generic/ogreditablelayer.h b/ogr/ogrsf_frmts/generic/ogreditablelayer.h index b852ac776a5b..33b567a61cbf 100644 --- a/ogr/ogrsf_frmts/generic/ogreditablelayer.h +++ b/ogr/ogrsf_frmts/generic/ogreditablelayer.h @@ -91,6 +91,8 @@ class CPL_DLL OGREditableLayer : public OGRLayerDecorator double dfMaxY) override; virtual OGRErr SetAttributeFilter(const char *) override; + virtual bool GetArrowStream(struct ArrowArrayStream *out_stream, + CSLConstList papszOptions = nullptr) override; virtual void ResetReading() override; virtual OGRFeature *GetNextFeature() override; From 1fb26344c09da97f9db63b2f826f9c350dff542e Mon Sep 17 00:00:00 2001 From: Even Rouault Date: Thu, 19 Oct 2023 22:12:21 +0200 Subject: [PATCH 5/5] test_ogrsf: add compliance tests for GetArrowStream() --- apps/test_ogrsf.cpp | 528 ++++++++++++++++++++++++++++++++++++++------ 1 file changed, 465 insertions(+), 63 deletions(-) diff --git a/apps/test_ogrsf.cpp b/apps/test_ogrsf.cpp index e08e61cbac58..933730663bbf 100644 --- a/apps/test_ogrsf.cpp +++ b/apps/test_ogrsf.cpp @@ -34,8 +34,10 @@ #include "ogr_p.h" #include "ogrsf_frmts.h" #include "ogr_swq.h" +#include "ogr_recordbatch.h" #include "commonutils.h" +#include #include #include @@ -2619,25 +2621,43 @@ static int TestSpatialFilter(OGRLayer *poLayer) } /************************************************************************/ -/* TestAttributeFilter() */ -/* */ -/* This is intended to be a simple test of the attribute */ -/* filtering. We read the first feature. Then construct a */ -/* attribute filter which includes it, install and */ -/* verify that we get the feature. Next install a attribute */ -/* filter that doesn't include this feature, and test again. */ +/* GetQuotedIfNeededIdentifier() */ /************************************************************************/ -static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) +static std::string GetQuotedIfNeededIdentifier(const char *pszFieldName) +{ + std::string osIdentifier; + const bool bMustQuoteAttrName = + pszFieldName[0] == '\0' || strchr(pszFieldName, '_') || + strchr(pszFieldName, ' ') || swq_is_reserved_keyword(pszFieldName); + if (bMustQuoteAttrName) + { + osIdentifier = "\""; + osIdentifier += pszFieldName; + osIdentifier += "\""; + } + else + { + osIdentifier = pszFieldName; + } + return osIdentifier; +} + +/************************************************************************/ +/* GetAttributeFilters() */ +/************************************************************************/ +static bool GetAttributeFilters(OGRLayer *poLayer, + std::unique_ptr &poTargetFeature, + std::string &osInclusiveFilter, + std::string &osExclusiveFilter) { - int bRet = TRUE; /* -------------------------------------------------------------------- */ /* Read the target feature. */ /* -------------------------------------------------------------------- */ LOG_ACTION(poLayer->ResetReading()); - OGRFeature *poTargetFeature = LOG_ACTION(poLayer->GetNextFeature()); + poTargetFeature.reset(LOG_ACTION(poLayer->GetNextFeature())); if (poTargetFeature == nullptr) { @@ -2647,7 +2667,7 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) " No features in layer.\n", poLayer->GetName()); } - return bRet; + return false; } int i = 0; @@ -2669,11 +2689,10 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) " Could not find non NULL field.\n", poLayer->GetName()); } - DestroyFeatureAndNullify(poTargetFeature); - return bRet; + return false; } - const char *pszFieldName = + const std::string osFieldName = poTargetFeature->GetFieldDefnRef(i)->GetNameRef(); CPLString osValue = poTargetFeature->GetFieldAsString(i); if (eType == OFTReal) @@ -2694,32 +2713,64 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) /* -------------------------------------------------------------------- */ /* Construct inclusive filter. */ /* -------------------------------------------------------------------- */ + osInclusiveFilter = GetQuotedIfNeededIdentifier(osFieldName.c_str()); + osInclusiveFilter += " = "; + if (eType == OFTString) + osInclusiveFilter += "'"; + osInclusiveFilter += osValue; + if (eType == OFTString) + osInclusiveFilter += "'"; + /* Make sure that the literal will be recognized as a float value */ + /* to avoid int underflow/overflow */ + else if (eType == OFTReal && strchr(osValue, '.') == nullptr) + osInclusiveFilter += "."; - CPLString osAttributeFilter; - const bool bMustQuoteAttrName = - pszFieldName[0] == '\0' || strchr(pszFieldName, '_') || - strchr(pszFieldName, ' ') || swq_is_reserved_keyword(pszFieldName); - if (bMustQuoteAttrName) - { - osAttributeFilter = "\""; - osAttributeFilter += pszFieldName; - osAttributeFilter += "\""; - } - else - { - osAttributeFilter = pszFieldName; - } - osAttributeFilter += " = "; + /* -------------------------------------------------------------------- */ + /* Construct exclusive filter. */ + /* -------------------------------------------------------------------- */ + osExclusiveFilter = GetQuotedIfNeededIdentifier(osFieldName.c_str()); + osExclusiveFilter += " <> "; if (eType == OFTString) - osAttributeFilter += "'"; - osAttributeFilter += osValue; + osExclusiveFilter += "'"; + osExclusiveFilter += osValue; if (eType == OFTString) - osAttributeFilter += "'"; + osExclusiveFilter += "'"; /* Make sure that the literal will be recognized as a float value */ /* to avoid int underflow/overflow */ else if (eType == OFTReal && strchr(osValue, '.') == nullptr) - osAttributeFilter += "."; - LOG_ACTION(poLayer->SetAttributeFilter(osAttributeFilter)); + osExclusiveFilter += "."; + + return true; +} + +/************************************************************************/ +/* TestAttributeFilter() */ +/* */ +/* This is intended to be a simple test of the attribute */ +/* filtering. We read the first feature. Then construct a */ +/* attribute filter which includes it, install and */ +/* verify that we get the feature. Next install a attribute */ +/* filter that doesn't include this feature, and test again. */ +/************************************************************************/ + +static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) + +{ + int bRet = TRUE; + + std::unique_ptr poTargetFeature; + std::string osInclusiveFilter; + std::string osExclusiveFilter; + if (!GetAttributeFilters(poLayer, poTargetFeature, osInclusiveFilter, + osExclusiveFilter)) + { + return true; + } + + /* -------------------------------------------------------------------- */ + /* Apply inclusive filter. */ + /* -------------------------------------------------------------------- */ + LOG_ACTION(poLayer->SetAttributeFilter(osInclusiveFilter.c_str())); /* -------------------------------------------------------------------- */ /* Verify that we can find the target feature. */ @@ -2730,7 +2781,7 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) OGRFeature *poFeature = nullptr; while ((poFeature = LOG_ACTION(poLayer->GetNextFeature())) != nullptr) { - if (poFeature->Equal(poTargetFeature)) + if (poFeature->Equal(poTargetFeature.get())) { bFoundFeature = true; DestroyFeatureAndNullify(poFeature); @@ -2755,29 +2806,9 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) const GIntBig nInclusiveCount = LOG_ACTION(poLayer->GetFeatureCount()); /* -------------------------------------------------------------------- */ - /* Construct exclusive filter. */ + /* Apply exclusive filter. */ /* -------------------------------------------------------------------- */ - if (bMustQuoteAttrName) - { - osAttributeFilter = "\""; - osAttributeFilter += pszFieldName; - osAttributeFilter += "\""; - } - else - { - osAttributeFilter = pszFieldName; - } - osAttributeFilter += " <> "; - if (eType == OFTString) - osAttributeFilter += "'"; - osAttributeFilter += osValue; - if (eType == OFTString) - osAttributeFilter += "'"; - /* Make sure that the literal will be recognized as a float value */ - /* to avoid int underflow/overflow */ - else if (eType == OFTReal && strchr(osValue, '.') == nullptr) - osAttributeFilter += "."; - LOG_ACTION(poLayer->SetAttributeFilter(osAttributeFilter)); + LOG_ACTION(poLayer->SetAttributeFilter(osExclusiveFilter.c_str())); /* -------------------------------------------------------------------- */ /* Verify that we can find the target feature. */ @@ -2787,7 +2818,7 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) GIntBig nExclusiveCountWhileIterating = 0; while ((poFeature = LOG_ACTION(poLayer->GetNextFeature())) != nullptr) { - if (poFeature->Equal(poTargetFeature)) + if (poFeature->Equal(poTargetFeature.get())) { DestroyFeatureAndNullify(poFeature); break; @@ -2803,13 +2834,13 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) // Check that GetFeature() ignores the attribute filter OGRFeature *poFeature2 = - LOG_ACTION(poLayer->GetFeature(poTargetFeature->GetFID())); + LOG_ACTION(poLayer->GetFeature(poTargetFeature.get()->GetFID())); poLayer->ResetReading(); OGRFeature *poFeature3 = nullptr; while ((poFeature3 = LOG_ACTION(poLayer->GetNextFeature())) != nullptr) { - if (poFeature3->Equal(poTargetFeature)) + if (poFeature3->Equal(poTargetFeature.get())) { DestroyFeatureAndNullify(poFeature3); break; @@ -2846,7 +2877,7 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) printf("INFO: Attribute filter exclusion seems to work.\n"); } - if (poFeature2 == nullptr || !poFeature2->Equal(poTargetFeature)) + if (poFeature2 == nullptr || !poFeature2->Equal(poTargetFeature.get())) { bRet = FALSE; printf("ERROR: Attribute filter has been taken into account " @@ -2868,8 +2899,6 @@ static int TestAttributeFilter(CPL_UNUSED GDALDataset *poDS, OGRLayer *poLayer) if (poFeature2 != nullptr) DestroyFeatureAndNullify(poFeature2); - DestroyFeatureAndNullify(poTargetFeature); - return bRet; } @@ -3782,6 +3811,374 @@ static int TestLayerSQL(GDALDataset *poDS, OGRLayer *poLayer) return bRet; } +/************************************************************************/ +/* CountFeaturesUsingArrowStream() */ +/************************************************************************/ + +static int64_t CountFeaturesUsingArrowStream(OGRLayer *poLayer, + int64_t nExpectedFID, + int64_t nUnexpectedFID, bool &bOK) +{ + struct ArrowArrayStream stream; + if (!LOG_ACTION(poLayer->GetArrowStream(&stream))) + { + printf("ERROR: GetArrowStream() failed\n"); + return -1; + } + struct ArrowSchema schema; + if (stream.get_schema(&stream, &schema) != 0) + { + printf("ERROR: stream.get_schema() failed\n"); + stream.release(&stream); + return -1; + } + int iFIDColumn = -1; + if (schema.n_children > 0 && + (strcmp(schema.children[0]->name, "OGC_FID") == 0 || + strcmp(schema.children[0]->name, poLayer->GetFIDColumn()) == 0) && + strcmp(schema.children[0]->format, "l") == 0) + { + iFIDColumn = 0; + } + schema.release(&schema); + int64_t nFeatureCountFiltered = 0; + bool bExpectedFIDFound = false; + bool bUnexpectedFIDFound = false; + while (true) + { + struct ArrowArray array; + if (stream.get_next(&stream, &array) != 0) + { + printf("ERROR: stream.get_next() is NULL\n"); + stream.release(&stream); + return -1; + } + if (!array.release) + break; + if (iFIDColumn >= 0 && (nExpectedFID >= 0 || nUnexpectedFID >= 0)) + { + const int64_t *panIds = + static_cast( + array.children[iFIDColumn]->buffers[1]) + + array.children[iFIDColumn]->offset; + for (int64_t i = 0; i < array.length; ++i) + { + if (nExpectedFID >= 0 && panIds[i] == nExpectedFID) + bExpectedFIDFound = true; + if (nUnexpectedFID >= 0 && panIds[i] == nUnexpectedFID) + bUnexpectedFIDFound = true; + } + } + nFeatureCountFiltered += array.length; + array.release(&array); + } + if (iFIDColumn >= 0) + { + if (nExpectedFID >= 0 && !bExpectedFIDFound) + { + bOK = false; + printf("ERROR: expected to find feature of id %" PRId64 + ", but did not get it\n", + nExpectedFID); + } + if (nUnexpectedFID >= 0 && bUnexpectedFIDFound) + { + bOK = false; + printf("ERROR: expected *not* to find feature of id %" PRId64 + ", but did get it\n", + nUnexpectedFID); + } + } + stream.release(&stream); + return nFeatureCountFiltered; +} + +/************************************************************************/ +/* TestLayerGetArrowStream() */ +/************************************************************************/ + +static int TestLayerGetArrowStream(OGRLayer *poLayer) +{ + LOG_ACTION(poLayer->SetSpatialFilter(nullptr)); + LOG_ACTION(poLayer->SetAttributeFilter(nullptr)); + LOG_ACTION(poLayer->ResetReading()); + + struct ArrowArrayStream stream; + if (!LOG_ACTION(poLayer->GetArrowStream(&stream))) + { + printf("ERROR: GetArrowStream() failed\n"); + return false; + } + + if (!stream.release) + { + printf("ERROR: stream.release is NULL\n"); + return false; + } + + struct ArrowSchema schema; + if (stream.get_schema(&stream, &schema) != 0) + { + printf("ERROR: stream.get_schema() failed\n"); + stream.release(&stream); + return false; + } + + if (!schema.release) + { + printf("ERROR: schema.release is NULL\n"); + stream.release(&stream); + return false; + } + + if (strcmp(schema.format, "+s") != 0) + { + printf("ERROR: expected schema.format to be '+s'. Got '%s'\n", + schema.format); + schema.release(&schema); + stream.release(&stream); + return false; + } + + int64_t nFeatureCount = 0; + while (true) + { + struct ArrowArray array; + if (stream.get_next(&stream, &array) != 0) + { + printf("ERROR: stream.get_next() is NULL\n"); + schema.release(&schema); + stream.release(&stream); + return false; + } + if (array.release == nullptr) + { + break; + } + + if (array.n_children != schema.n_children) + { + printf("ERROR: expected array.n_children (=%d) to be " + "schema.n_children (=%d)\n", + int(array.n_children), int(schema.n_children)); + array.release(&array); + schema.release(&schema); + stream.release(&stream); + return false; + } + + int bRet = true; + for (int i = 0; i < array.n_children; ++i) + { + if (array.children[i]->length != array.length) + { + bRet = false; + printf("ERROR: expected array.children[i]->length (=%d) to be " + "array.length (=%d)\n", + int(array.children[i]->length), int(array.length)); + } + } + if (!bRet) + { + array.release(&array); + schema.release(&schema); + stream.release(&stream); + return false; + } + + nFeatureCount += array.length; + + array.release(&array); + + if (array.release) + { + printf("ERROR: array.release should be NULL after release\n"); + schema.release(&schema); + stream.release(&stream); + return false; + } + } + + bool bRet = true; + // We are a bit in non-specified behaviour below by calling get_next() + // after end of iteration. + { + struct ArrowArray array; + if (stream.get_next(&stream, &array) == 0) + { + if (array.length != 0) + { + bRet = false; + printf("ERROR: get_next() return an array with length != 0 " + "after end of iteration\n"); + } + if (array.release) + array.release(&array); + } + } + + schema.release(&schema); + if (schema.release) + { + printf("ERROR: schema.release should be NULL after release\n"); + stream.release(&stream); + return false; + } + + stream.release(&stream); + if (stream.release) + { + printf("ERROR: stream.release should be NULL after release\n"); + return false; + } + + const int64_t nFCClassic = poLayer->GetFeatureCount(true); + if (nFeatureCount != nFCClassic) + { + printf("ERROR: GetArrowStream() returned %" PRId64 + " features, whereas GetFeatureCount() returned %" PRId64 "\n", + nFeatureCount, nFCClassic); + bRet = false; + } + + { + LOG_ACTION(poLayer->SetAttributeFilter("1=1")); + const auto nFeatureCountFiltered = + CountFeaturesUsingArrowStream(poLayer, -1, -1, bRet); + LOG_ACTION(poLayer->SetAttributeFilter(nullptr)); + if (nFeatureCount != nFeatureCountFiltered) + { + printf("ERROR: GetArrowStream() with 1=1 filter returned %" PRId64 + " features, whereas %" PRId64 " expected\n", + nFeatureCountFiltered, nFeatureCount); + bRet = false; + } + } + + { + LOG_ACTION(poLayer->SetAttributeFilter("1=0")); + const auto nFeatureCountFiltered = + CountFeaturesUsingArrowStream(poLayer, -1, -1, bRet); + LOG_ACTION(poLayer->SetAttributeFilter(nullptr)); + if (nFeatureCountFiltered != 0) + { + printf("ERROR: GetArrowStream() with 1=0 filter returned %" PRId64 + " features, whereas 0 expected\n", + nFeatureCountFiltered); + bRet = false; + } + } + + std::unique_ptr poTargetFeature; + std::string osInclusiveFilter; + std::string osExclusiveFilter; + if (GetAttributeFilters(poLayer, poTargetFeature, osInclusiveFilter, + osExclusiveFilter)) + { + { + LOG_ACTION(poLayer->SetAttributeFilter(osInclusiveFilter.c_str())); + const auto nFeatureCountFiltered = CountFeaturesUsingArrowStream( + poLayer, poTargetFeature->GetFID(), -1, bRet); + LOG_ACTION(poLayer->SetAttributeFilter(nullptr)); + if (nFeatureCountFiltered == 0) + { + printf( + "ERROR: GetArrowStream() with %s filter returned %" PRId64 + " features, whereas at least one expected\n", + osInclusiveFilter.c_str(), nFeatureCountFiltered); + bRet = false; + } + else if (bVerbose) + { + printf("INFO: Attribute filter inclusion with GetArrowStream " + "seems to work.\n"); + } + } + + { + LOG_ACTION(poLayer->SetAttributeFilter(osExclusiveFilter.c_str())); + const auto nFeatureCountFiltered = + CountFeaturesUsingArrowStream(poLayer, -1, -1, bRet); + LOG_ACTION(poLayer->SetAttributeFilter(nullptr)); + if (nFeatureCountFiltered >= nFCClassic) + { + printf( + "ERROR: GetArrowStream() with %s filter returned %" PRId64 + " features, whereas less than %" PRId64 " expected\n", + osExclusiveFilter.c_str(), nFeatureCountFiltered, + nFCClassic); + bRet = false; + } + else if (bVerbose) + { + printf("INFO: Attribute filter exclusion with GetArrowStream " + "seems to work.\n"); + } + } + + auto poGeom = poTargetFeature->GetGeometryRef(); + if (poGeom && !poGeom->IsEmpty()) + { + OGREnvelope sEnvelope; + poGeom->getEnvelope(&sEnvelope); + + OGREnvelope sLayerExtent; + double epsilon = 10.0; + if (LOG_ACTION(poLayer->TestCapability(OLCFastGetExtent)) && + LOG_ACTION(poLayer->GetExtent(&sLayerExtent)) == OGRERR_NONE && + sLayerExtent.MinX < sLayerExtent.MaxX && + sLayerExtent.MinY < sLayerExtent.MaxY) + { + epsilon = std::min(sLayerExtent.MaxX - sLayerExtent.MinX, + sLayerExtent.MaxY - sLayerExtent.MinY) / + 10.0; + } + + /* -------------------------------------------------------------------- */ + /* Construct inclusive filter. */ + /* -------------------------------------------------------------------- */ + + OGRLinearRing oRing; + oRing.setPoint(0, sEnvelope.MinX - 2 * epsilon, + sEnvelope.MinY - 2 * epsilon); + oRing.setPoint(1, sEnvelope.MinX - 2 * epsilon, + sEnvelope.MaxY + 1 * epsilon); + oRing.setPoint(2, sEnvelope.MaxX + 1 * epsilon, + sEnvelope.MaxY + 1 * epsilon); + oRing.setPoint(3, sEnvelope.MaxX + 1 * epsilon, + sEnvelope.MinY - 2 * epsilon); + oRing.setPoint(4, sEnvelope.MinX - 2 * epsilon, + sEnvelope.MinY - 2 * epsilon); + + OGRPolygon oInclusiveFilter; + oInclusiveFilter.addRing(&oRing); + + LOG_ACTION(poLayer->SetSpatialFilter(&oInclusiveFilter)); + const auto nFeatureCountFiltered = CountFeaturesUsingArrowStream( + poLayer, poTargetFeature->GetFID(), -1, bRet); + LOG_ACTION(poLayer->SetSpatialFilter(nullptr)); + if (nFeatureCountFiltered == 0) + { + printf("ERROR: GetArrowStream() with inclusive spatial filter " + "returned %" PRId64 + " features, whereas at least 1 expected\n", + nFeatureCountFiltered); + bRet = false; + } + else if (bVerbose) + { + printf("INFO: Spatial filter inclusion with GetArrowStream " + "seems to work.\n"); + } + } + } + + if (bRet && bVerbose) + printf("INFO: TestLayerGetArrowStream passed.\n"); + + return bRet; +} + /************************************************************************/ /* TestOGRLayer() */ /************************************************************************/ @@ -3830,6 +4227,11 @@ static int TestOGRLayer(GDALDataset *poDS, OGRLayer *poLayer, int bIsSQLLayer) /* -------------------------------------------------------------------- */ bRet &= TestGetExtent(poLayer); + /* -------------------------------------------------------------------- */ + /* Test GetArrowStream() interface */ + /* -------------------------------------------------------------------- */ + bRet &= TestLayerGetArrowStream(poLayer); + /* -------------------------------------------------------------------- */ /* Test random reading. */ /* -------------------------------------------------------------------- */