Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

GPKG: implement SetNextByIndex() #8306

Merged
merged 6 commits into from
Sep 5, 2023
Merged

GPKG: implement SetNextByIndex() #8306

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
44252ab
GPKG: remove duplicate call to CancelAsyncNextArrowArray() in DeleteF…
rouault Sep 2, 2023
5b84080
OGRGeoPackageLayer: prefix all member variables with m_, and use inli…
rouault Sep 2, 2023
9aaf4e8
GPKG: implement SetNextByIndex() on table layers (by appending 'OFFSE…
rouault Sep 2, 2023
15ad2b4
TileDB ArrowArray: make sure ResetReading() does not invalidate array
rouault Sep 2, 2023
248cf60
OGRLayer::GetArrowStream(): do not issue ResetReading() at beginning …
rouault Sep 2, 2023
ecf2f4d
GPKG: make GetArrowStream() honour SetNextByIndex()
rouault Sep 2, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 13 additions & 0 deletions autotest/ogr/ogr_gpkg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7687,6 +7687,19 @@ def test_ogr_gpkg_arrow_stream_numpy():
assert batch["int16"][0] == 123
assert len(batch["fid"]) == 1

assert lyr.SetNextByIndex(1) == ogr.OGRERR_NONE
stream = lyr.GetArrowStreamAsNumPy(options=["USE_MASKED_ARRAYS=NO"])
batches = [batch for batch in stream]
assert len(batches) == 1
assert list(batches[0]["fid"]) == [2, 3]

with ds.ExecuteSQL("SELECT * FROM test") as sql_lyr:
assert sql_lyr.SetNextByIndex(1) == ogr.OGRERR_NONE
stream = sql_lyr.GetArrowStreamAsNumPy(options=["USE_MASKED_ARRAYS=NO"])
batches = [batch for batch in stream]
assert len(batches) == 1
assert list(batches[0]["fid"]) == [2, 3]

with lyr.GetArrowStreamAsNumPy(options=["MAX_FEATURES_IN_BATCH=1"]) as stream:
batches = [batch for batch in stream]
assert len(batches) == 3
Expand Down
3 changes: 2 additions & 1 deletion frmts/tiledb/tiledbheaders.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,7 @@ class OGRTileDBLayer final : public OGRLayer,
bool m_bInitialized = false;
OGRFeatureDefn *m_poFeatureDefn = nullptr;
std::string m_osFIDColumn{};
GIntBig m_nNextFID = -1;
GIntBig m_nNextFID = 1;
int64_t m_nTotalFeatureCount = -1;
bool m_bStats = false;
bool m_bQueryComplete = false;
Expand Down Expand Up @@ -345,6 +345,7 @@ class OGRTileDBLayer final : public OGRLayer,
const char *GetDatabaseGeomColName();
void InitializeSchemaAndArray();
void FlushArrays();
void AllocateNewBuffers();
void ResetBuffers();
void SwitchToReadingMode();
void SwitchToWritingMode();
Expand Down
63 changes: 37 additions & 26 deletions frmts/tiledb/tiledbsparse.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -511,7 +511,6 @@ OGRLayer *OGRTileDBDataset::ICreateLayer(const char *pszName,
poLayer->m_poFeatureDefn->GetGeomFieldDefn(0)->SetName(pszGeomColName);

poLayer->m_eCurrentMode = OGRTileDBLayer::CurrentMode::WriteInProgress;
poLayer->m_nNextFID = 1;

const char *pszTileDBStringType =
CSLFetchNameValue(papszOptions, "TILEDB_STRING_TYPE");
Expand Down Expand Up @@ -1550,24 +1549,28 @@ template <class T> struct ResetArray
};
} // namespace

bool OGRTileDBLayer::SetupQuery(tiledb::QueryCondition *queryCondition)
void OGRTileDBLayer::AllocateNewBuffers()
{
if (!m_bArrowBatchReleased)
m_anFIDs = std::make_shared<std::vector<int64_t>>();
m_adfXs = std::make_shared<std::vector<double>>();
m_adfYs = std::make_shared<std::vector<double>>();
m_adfZs = std::make_shared<std::vector<double>>();
m_abyGeometries = std::make_shared<std::vector<unsigned char>>();
m_anGeometryOffsets = std::make_shared<std::vector<uint64_t>>();

for (int i = 0; i < m_poFeatureDefn->GetFieldCount(); i++)
{
m_anFIDs = std::make_shared<std::vector<int64_t>>();
m_adfXs = std::make_shared<std::vector<double>>();
m_adfYs = std::make_shared<std::vector<double>>();
m_adfZs = std::make_shared<std::vector<double>>();
m_abyGeometries = std::make_shared<std::vector<unsigned char>>();
m_anGeometryOffsets = std::make_shared<std::vector<uint64_t>>();
ProcessField<ResetArray>::exec(m_aeFieldTypes[i], m_aFieldValues[i]);

for (int i = 0; i < m_poFeatureDefn->GetFieldCount(); i++)
{
ProcessField<ResetArray>::exec(m_aeFieldTypes[i],
m_aFieldValues[i]);
m_aFieldValueOffsets[i] = std::make_shared<std::vector<uint64_t>>();
}
}

m_aFieldValueOffsets[i] = std::make_shared<std::vector<uint64_t>>();
}
bool OGRTileDBLayer::SetupQuery(tiledb::QueryCondition *queryCondition)
{
if (!m_bArrowBatchReleased)
{
AllocateNewBuffers();
}

m_anFIDs->clear();
Expand Down Expand Up @@ -4479,18 +4482,26 @@ template <class T> struct ClearArray

void OGRTileDBLayer::ResetBuffers()
{
// Reset buffers
m_anFIDs->clear();
m_adfXs->clear();
m_adfYs->clear();
m_adfZs->clear();
m_abyGeometries->clear();
m_anGeometryOffsets->clear();
for (int i = 0; i < m_poFeatureDefn->GetFieldCount(); i++)
if (!m_bArrowBatchReleased)
{
m_aFieldValueOffsets[i]->clear();
m_aFieldValidity[i].clear();
ProcessField<ClearArray>::exec(m_aeFieldTypes[i], m_aFieldValues[i]);
AllocateNewBuffers();
}
else
{
// Reset buffers
m_anFIDs->clear();
m_adfXs->clear();
m_adfYs->clear();
m_adfZs->clear();
m_abyGeometries->clear();
m_anGeometryOffsets->clear();
for (int i = 0; i < m_poFeatureDefn->GetFieldCount(); i++)
{
m_aFieldValueOffsets[i]->clear();
m_aFieldValidity[i].clear();
ProcessField<ClearArray>::exec(m_aeFieldTypes[i],
m_aFieldValues[i]);
}
}
}

Expand Down
4 changes: 2 additions & 2 deletions ogr/ogrsf_frmts/generic/ogrlayerarrow.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1729,6 +1729,8 @@ void OGRLayer::ReleaseStream(struct ArrowArrayStream *stream)
static_cast<ArrowArrayStreamPrivateDataSharedDataWrapper *>(
stream->private_data);
poPrivate->poShared->m_bArrowArrayStreamInProgress = false;
if (poPrivate->poShared->m_poLayer)
poPrivate->poShared->m_poLayer->ResetReading();
delete poPrivate;
stream->private_data = nullptr;
stream->release = nullptr;
Expand Down Expand Up @@ -1865,8 +1867,6 @@ bool OGRLayer::GetArrowStream(struct ArrowArrayStream *out_stream,
}
m_aosArrowArrayStreamOptions.Assign(CSLDuplicate(papszOptions), true);

ResetReading();

out_stream->get_schema = OGRLayer::StaticGetArrowSchema;
out_stream->get_next = OGRLayer::StaticGetNextArrowArray;
out_stream->get_last_error = OGRLayer::GetLastErrorArrowArrayStream;
Expand Down
23 changes: 13 additions & 10 deletions ogr/ogrsf_frmts/gpkg/ogr_geopackage.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -539,20 +539,20 @@ class OGRGeoPackageLayer CPL_NON_FINAL : public OGRLayer,
int /*argc*/,
sqlite3_value **argv);

GDALGeoPackageDataset *m_poDS;
GDALGeoPackageDataset *m_poDS = nullptr;

OGRFeatureDefn *m_poFeatureDefn;
GIntBig iNextShapeId;
OGRFeatureDefn *m_poFeatureDefn = nullptr;
GIntBig m_iNextShapeId = 0;

sqlite3_stmt *m_poQueryStatement;
bool bDoStep;
sqlite3_stmt *m_poQueryStatement = nullptr;
bool m_bDoStep = true;
bool m_bEOF = false;

char *m_pszFidColumn;
char *m_pszFidColumn = nullptr;

int iFIDCol;
int iGeomCol;
int *panFieldOrdinals;
int m_iFIDCol = -1;
int m_iGeomCol = -1;
std::vector<int> m_anFieldOrdinals{};

void ClearStatement();
virtual OGRErr ResetStatement() = 0;
Expand Down Expand Up @@ -667,6 +667,7 @@ class OGRGeoPackageTableLayer final : public OGRGeoPackageLayer
std::set<OGRwkbGeometryType> m_eSetBadGeomTypeWarned{};

int m_nIsCompatOfOptimizedGetNextArrowArray = -1;
bool m_bGetNextArrowArrayCalledSinceResetReading = false;

int m_nCountInsertInTransactionThreshold = -1;
GIntBig m_nCountInsertInTransaction = 0;
Expand Down Expand Up @@ -707,6 +708,7 @@ class OGRGeoPackageTableLayer final : public OGRGeoPackageLayer
void RemoveAsyncRTreeTempDB();
void AsyncRTreeThreadFunction();

OGRErr ResetStatementInternal(GIntBig nStartIndex);
virtual OGRErr ResetStatement() override;

void BuildWhere();
Expand Down Expand Up @@ -803,6 +805,7 @@ class OGRGeoPackageTableLayer final : public OGRGeoPackageLayer
int nFlagsIn) override;
virtual OGRErr ReorderFields(int *panMap) override;
void ResetReading() override;
OGRErr SetNextByIndex(GIntBig nIndex) override;
OGRErr ICreateFeature(OGRFeature *poFeature) override;
OGRErr ISetFeature(OGRFeature *poFeature) override;
OGRErr IUpsertFeature(OGRFeature *poFeature) override;
Expand Down Expand Up @@ -1032,7 +1035,7 @@ class OGRGeoPackageSelectLayer final : public OGRGeoPackageLayer,
}
virtual int HasReadFeature() override
{
return iNextShapeId > 0;
return m_iNextShapeId > 0;
}
virtual void BaseResetReading() override
{
Expand Down
Loading