Skip to content

Commit 67b97e2

Browse files
aheevaheev
committed
fix tests
1 parent a753b25 commit 67b97e2

12 files changed

Lines changed: 990 additions & 592 deletions

‎lance/src/include/lance_node_table.h‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -25,10 +25,7 @@ struct LanceBatchData {
2525
std::memset(&schema, 0, sizeof(schema));
2626
std::memset(&array, 0, sizeof(array));
2727
}
28-
~LanceBatchData() {
29-
if (array.release) array.release(&array);
30-
// schema.release is intentionally left null (owned by shared state)
31-
}
28+
~LanceBatchData() = default;
3229

3330
LanceBatchData(const LanceBatchData&) = delete;
3431
LanceBatchData& operator=(const LanceBatchData&) = delete;
@@ -57,7 +54,8 @@ struct LanceNodeTableScanSharedState final : storage::ColumnarNodeTableScanShare
5754
}
5855

5956
~LanceNodeTableScanSharedState() override {
60-
if (stream_.release) stream_.release(&stream_);
57+
if (stream_.release)
58+
stream_.release(&stream_);
6159
}
6260

6361
void reset(ArrowArrayStream newStream);
@@ -87,10 +85,10 @@ struct LanceNodeTableScanSharedState final : storage::ColumnarNodeTableScanShare
8785
class LanceNodeTable final : public storage::ColumnarNodeTableBase {
8886
public:
8987
LanceNodeTable(const storage::StorageManager* storageManager,
90-
const catalog::NodeTableCatalogEntry* nodeTableEntry,
91-
storage::MemoryManager* memoryManager, main::ClientContext* context);
88+
const catalog::NodeTableCatalogEntry* nodeTableEntry, storage::MemoryManager* memoryManager,
89+
main::ClientContext* context);
9290

93-
~LanceNodeTable() override = default;
91+
~LanceNodeTable() override;
9492

9593
void initializeScanCoordination(const transaction::Transaction* transaction) override;
9694

@@ -125,6 +123,8 @@ class LanceNodeTable final : public storage::ColumnarNodeTableBase {
125123
common::row_idx_t getTotalRowCount(const transaction::Transaction* transaction) const override;
126124

127125
private:
126+
void ensureDatasetLoaded() const;
127+
128128
std::vector<int64_t> getOutputToLanceColumnIdx(
129129
const std::vector<common::column_id_t>& columnIDs) const;
130130

@@ -135,9 +135,10 @@ class LanceNodeTable final : public storage::ColumnarNodeTableBase {
135135
private:
136136
std::string datasetPath;
137137
mutable uint64_t cachedTotalRows = common::INVALID_ROW_IDX;
138-
uint32_t numLanceColumns = 0;
139-
ArrowSchema cachedSchema_;
140-
bool schemaCached = false;
138+
mutable uint32_t numLanceColumns = 0;
139+
mutable ArrowSchema cachedSchema_;
140+
mutable bool schemaCached = false;
141+
mutable std::mutex metadataMtx;
141142

142143
static constexpr size_t kDefaultMorselSize = 2048;
143144
};

‎lance/src/include/lance_rel_table.h‎

Lines changed: 12 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,9 @@ struct LanceBatchData;
1818

1919
/// Per-thread scan state for LanceRelTable.
2020
struct LanceRelTableScanState final : storage::RelTableScanState {
21-
std::shared_ptr<LanceBatchData> cachedBatchData;
22-
uint64_t currentBatchStartOffset = 0;
21+
/// Index into the LanceRelTable's flat edge cache for resumable morsel scans.
2322
uint64_t currentLocalRowIdx = 0;
2423

25-
std::unordered_map<common::offset_t, common::sel_t> boundNodeOffsets;
26-
27-
ArrowArrayStream stream;
28-
bool streamExhausted = false;
29-
bool streamInitialized = false;
30-
31-
ArrowSchema streamSchema;
32-
bool schemaFetched = false;
33-
3424
LanceRelTableScanState(storage::MemoryManager& mm, common::ValueVector* nodeIDVector,
3525
std::vector<common::ValueVector*> outputVectors,
3626
std::shared_ptr<common::DataChunkState> outChunkState);
@@ -42,8 +32,6 @@ struct LanceRelTableScanState final : storage::RelTableScanState {
4232
std::vector<storage::ColumnPredicateSet> columnPredicateSets_,
4333
common::RelDataDirection direction_) override;
4434

45-
void reset(std::unordered_map<common::offset_t, common::sel_t> boundNodeOffsets_);
46-
4735
LanceRelTableScanState(const LanceRelTableScanState&) = delete;
4836
LanceRelTableScanState& operator=(const LanceRelTableScanState&) = delete;
4937
};
@@ -57,6 +45,10 @@ class LanceRelTable final : public storage::ColumnarRelTableBase {
5745

5846
~LanceRelTable() override = default;
5947

48+
std::unique_ptr<storage::RelTableScanState> createScanState(common::ValueVector* nodeIDVector,
49+
const std::vector<common::ValueVector*>& outVectors,
50+
storage::MemoryManager* memoryManager) const override;
51+
6052
void initScanState(transaction::Transaction* transaction, storage::TableScanState& scanState,
6153
bool resetCachedBoundNodeSelVec = true) const override;
6254

@@ -78,12 +70,17 @@ class LanceRelTable final : public storage::ColumnarRelTableBase {
7870
common::idx_t k) const override;
7971

8072
private:
73+
void ensureDatasetLoaded() const;
74+
8175
bool scanFlat(transaction::Transaction* transaction, LanceRelTableScanState& scanState);
8276

83-
int32_t fromColumnIdx = -1;
84-
int32_t toColumnIdx = -1;
77+
mutable int32_t fromColumnIdx = -1;
78+
mutable int32_t toColumnIdx = -1;
8579
std::string datasetPath;
8680
mutable uint64_t cachedTotalRows = common::INVALID_ROW_IDX;
81+
mutable std::vector<std::pair<common::offset_t, common::offset_t>> edgeCache;
82+
mutable bool edgeCacheLoaded = false;
83+
mutable std::mutex metadataMtx;
8784
};
8885

8986
} // namespace lance_extension

‎lance/src/lance_extension.cpp‎

Lines changed: 11 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -18,26 +18,23 @@ void LanceExtension::load(main::ClientContext* context) {
1818
auto* storageManager = storage::StorageManager::Get(*context);
1919

2020
// ── Register table factories ────────────────────────────────────────────
21-
storage::NodeTableFactory nodeFactory = [](const storage::StorageManager* sm,
22-
const catalog::NodeTableCatalogEntry* entry,
23-
storage::MemoryManager* mm,
24-
main::ClientContext* ctx)
25-
-> std::unique_ptr<storage::Table> {
21+
storage::NodeTableFactory nodeFactory =
22+
[](const storage::StorageManager* sm, const catalog::NodeTableCatalogEntry* entry,
23+
storage::MemoryManager* mm,
24+
main::ClientContext* ctx) -> std::unique_ptr<storage::Table> {
2625
return std::make_unique<LanceNodeTable>(sm, entry, mm, ctx);
2726
};
2827

29-
storage::RelTableFactory relFactory = [](catalog::RelGroupCatalogEntry* entry,
30-
common::table_id_t fromTableID,
31-
common::table_id_t toTableID,
32-
const storage::StorageManager* sm,
33-
storage::MemoryManager* mm,
34-
main::ClientContext* ctx)
35-
-> std::unique_ptr<storage::Table> {
28+
storage::RelTableFactory relFactory =
29+
[](catalog::RelGroupCatalogEntry* entry, common::table_id_t fromTableID,
30+
common::table_id_t toTableID, const storage::StorageManager* sm,
31+
storage::MemoryManager* mm,
32+
main::ClientContext* ctx) -> std::unique_ptr<storage::Table> {
3633
return std::make_unique<LanceRelTable>(entry, fromTableID, toTableID, sm, mm, ctx);
3734
};
3835

39-
storageManager->registerStorageFormatHandler(
40-
common::StorageFormat::LANCE, std::move(nodeFactory), std::move(relFactory));
36+
storageManager->registerStorageFormatHandler(common::StorageFormat::LANCE,
37+
std::move(nodeFactory), std::move(relFactory));
4138

4239
// ── Register search functions ───────────────────────────────────────────
4340
ExtensionUtils::addTableFunc<LanceVectorSearchFunction>(db);

‎lance/src/lance_functions.cpp‎

Lines changed: 45 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,10 @@
1515
#include "function/table/bind_data.h"
1616
#include "function/table/bind_input.h"
1717
#include "function/table/table_function.h"
18+
#include "lance/lance.hpp"
1819
#include "main/client_context.h"
1920
#include "processor/execution_context.h"
2021

21-
#include "lance/lance.hpp"
22-
2322
namespace lbug {
2423
namespace lance_extension {
2524

@@ -37,8 +36,7 @@ struct LanceVectorSearchBindData : TableFuncBindData {
3736
uint32_t nprobes;
3837

3938
LanceVectorSearchBindData(std::string path, std::string col, std::vector<float> query,
40-
uint32_t k, std::string metric, uint32_t nprobes,
41-
binder::expression_vector columns)
39+
uint32_t k, std::string metric, uint32_t nprobes, binder::expression_vector columns)
4240
: TableFuncBindData{std::move(columns)}, datasetPath{std::move(path)},
4341
columnName{std::move(col)}, queryVector{std::move(query)}, k{k},
4442
metric{std::move(metric)}, nprobes{nprobes} {}
@@ -57,7 +55,8 @@ struct LanceSearchSharedState : TableFuncSharedState {
5755
LanceSearchSharedState() { std::memset(&stream, 0, sizeof(stream)); }
5856

5957
~LanceSearchSharedState() override {
60-
if (stream.release) stream.release(&stream);
58+
if (stream.release)
59+
stream.release(&stream);
6160
}
6261
};
6362

@@ -78,9 +77,8 @@ static std::unique_ptr<TableFuncBindData> bindVectorSearch(main::ClientContext*
7877
}
7978
auto k = static_cast<uint32_t>(input->getLiteralVal<int64_t>(3));
8079
std::string metric = (input->params.size() > 4) ? input->getLiteralVal<std::string>(4) : "l2";
81-
uint32_t nprobes = (input->params.size() > 5)
82-
? static_cast<uint32_t>(input->getLiteralVal<int64_t>(5))
83-
: 1;
80+
uint32_t nprobes =
81+
(input->params.size() > 5) ? static_cast<uint32_t>(input->getLiteralVal<int64_t>(5)) : 1;
8482

8583
// Open dataset to discover schema
8684
auto resolvedPath = VirtualFileSystem::resolvePath(context, datasetPath);
@@ -93,15 +91,17 @@ static std::unique_ptr<TableFuncBindData> bindVectorSearch(main::ClientContext*
9391
std::vector<LogicalType> returnTypes;
9492
std::vector<std::string> returnNames;
9593
for (int32_t i = 0; i < schema.n_children; ++i) {
96-
if (!schema.children[i] || !schema.children[i]->name) continue;
94+
if (!schema.children[i] || !schema.children[i]->name)
95+
continue;
9796
returnNames.push_back(schema.children[i]->name);
9897
returnTypes.push_back(ArrowConverter::fromArrowSchema(schema.children[i]));
9998
}
10099
// Add the _distance column
101100
returnNames.push_back("_distance");
102101
returnTypes.push_back(LogicalType{LogicalTypeID::FLOAT});
103102

104-
if (schema.release) schema.release(&schema);
103+
if (schema.release)
104+
schema.release(&schema);
105105

106106
auto columns = input->binder->createVariables(returnNames, returnTypes);
107107
return std::make_unique<LanceVectorSearchBindData>(std::move(resolvedPath),
@@ -122,11 +122,14 @@ static std::shared_ptr<TableFuncSharedState> initVectorSearchSharedState(
122122
auto scanner = dataset.scan();
123123

124124
LanceMetricType metric = LANCE_METRIC_L2;
125-
if (bindData->metric == "cosine") metric = LANCE_METRIC_COSINE;
126-
else if (bindData->metric == "dot") metric = LANCE_METRIC_DOT;
125+
if (bindData->metric == "cosine")
126+
metric = LANCE_METRIC_COSINE;
127+
else if (bindData->metric == "dot")
128+
metric = LANCE_METRIC_DOT;
127129

128-
scanner.nearest(bindData->columnName, bindData->queryVector.data(),
129-
bindData->queryVector.size(), bindData->k)
130+
scanner
131+
.nearest(bindData->columnName, bindData->queryVector.data(),
132+
bindData->queryVector.size(), bindData->k)
130133
.nprobes(bindData->nprobes)
131134
.metric(metric);
132135

@@ -140,7 +143,8 @@ static std::shared_ptr<TableFuncSharedState> initVectorSearchSharedState(
140143

141144
static offset_t vectorSearchTableFunc(const TableFuncInput& input, TableFuncOutput& output) {
142145
auto* sharedState = input.sharedState->ptrCast<LanceSearchSharedState>();
143-
if (sharedState->exhausted) return 0;
146+
if (sharedState->exhausted)
147+
return 0;
144148

145149
ArrowArray batch;
146150
std::memset(&batch, 0, sizeof(batch));
@@ -160,19 +164,23 @@ static offset_t vectorSearchTableFunc(const TableFuncInput& input, TableFuncOutp
160164

161165
// Copy each column from the batch to the output data chunk
162166
for (uint64_t col = 0; col < static_cast<uint64_t>(output.dataChunk.getNumValueVectors()) &&
163-
col < static_cast<uint64_t>(batch.n_children) &&
164-
col < static_cast<uint64_t>(schema.n_children);
167+
col < static_cast<uint64_t>(batch.n_children) &&
168+
col < static_cast<uint64_t>(schema.n_children);
165169
++col) {
166170
auto* childArr = batch.children[col];
167171
auto* childSchema = schema.children[col];
168-
if (!childArr || !childSchema) continue;
172+
if (!childArr || !childSchema)
173+
continue;
169174
ArrowNullMaskTree nullMask(childSchema, childArr, childArr->offset, childArr->length);
170-
ArrowConverter::fromArrowArray(childSchema, childArr, output.dataChunk.getValueVectorMutable(col),
171-
&nullMask, static_cast<uint64_t>(childArr->offset), 0, batchLen);
175+
ArrowConverter::fromArrowArray(childSchema, childArr,
176+
output.dataChunk.getValueVectorMutable(col), &nullMask,
177+
static_cast<uint64_t>(childArr->offset), 0, batchLen);
172178
}
173179

174-
if (schema.release) schema.release(&schema);
175-
if (batch.release) batch.release(&batch);
180+
if (schema.release)
181+
schema.release(&schema);
182+
if (batch.release)
183+
batch.release(&batch);
176184
output.dataChunk.state->getSelVectorUnsafe().setSelSize(batchLen);
177185
return batchLen;
178186
}
@@ -221,9 +229,8 @@ static std::unique_ptr<TableFuncBindData> bindFTS(main::ClientContext* context,
221229
if (input->params.size() > 2) {
222230
searchCols.push_back(input->getLiteralVal<std::string>(2));
223231
}
224-
uint32_t maxFuzzy = (input->params.size() > 3)
225-
? static_cast<uint32_t>(input->getLiteralVal<int64_t>(3))
226-
: 0;
232+
uint32_t maxFuzzy =
233+
(input->params.size() > 3) ? static_cast<uint32_t>(input->getLiteralVal<int64_t>(3)) : 0;
227234

228235
auto resolvedPath = VirtualFileSystem::resolvePath(context, datasetPath);
229236
try {
@@ -235,14 +242,16 @@ static std::unique_ptr<TableFuncBindData> bindFTS(main::ClientContext* context,
235242
std::vector<LogicalType> returnTypes;
236243
std::vector<std::string> returnNames;
237244
for (int32_t i = 0; i < schema.n_children; ++i) {
238-
if (!schema.children[i] || !schema.children[i]->name) continue;
245+
if (!schema.children[i] || !schema.children[i]->name)
246+
continue;
239247
returnNames.push_back(schema.children[i]->name);
240248
returnTypes.push_back(ArrowConverter::fromArrowSchema(schema.children[i]));
241249
}
242250
returnNames.push_back("_score");
243251
returnTypes.push_back(LogicalType{LogicalTypeID::FLOAT});
244252

245-
if (schema.release) schema.release(&schema);
253+
if (schema.release)
254+
schema.release(&schema);
246255
auto columns = input->binder->createVariables(returnNames, returnTypes);
247256
return std::make_unique<LanceFTSBindData>(std::move(resolvedPath), std::move(query),
248257
std::move(searchCols), maxFuzzy, std::move(columns));
@@ -268,8 +277,8 @@ static std::shared_ptr<TableFuncSharedState> initFTSSharedState(
268277

269278
function_set LanceFTSFunction::getFunctionSet() {
270279
function_set functionSet;
271-
auto func = std::make_unique<TableFunction>(
272-
name, std::vector<LogicalTypeID>{LogicalTypeID::STRING, LogicalTypeID::STRING});
280+
auto func = std::make_unique<TableFunction>(name,
281+
std::vector<LogicalTypeID>{LogicalTypeID::STRING, LogicalTypeID::STRING});
273282
func->bindFunc = bindFTS;
274283
func->initSharedStateFunc = initFTSSharedState;
275284
func->initLocalStateFunc = TableFunction::initEmptyLocalState;
@@ -303,8 +312,8 @@ struct LanceHybridSearchBindData : TableFuncBindData {
303312
static std::unique_ptr<TableFuncBindData> bindHybridSearch(main::ClientContext* context,
304313
const TableFuncBindInput* input) {
305314
if (input->params.size() < 5) {
306-
throw RuntimeException(
307-
"LANCE_HYBRID_SEARCH requires 5 arguments: dataset_path, column, query_vector, k, fts_query");
315+
throw RuntimeException("LANCE_HYBRID_SEARCH requires 5 arguments: dataset_path, column, "
316+
"query_vector, k, fts_query");
308317
}
309318
auto datasetPath = input->getLiteralVal<std::string>(0);
310319
auto columnName = input->getLiteralVal<std::string>(1);
@@ -327,18 +336,19 @@ static std::unique_ptr<TableFuncBindData> bindHybridSearch(main::ClientContext*
327336
std::vector<LogicalType> returnTypes;
328337
std::vector<std::string> returnNames;
329338
for (int32_t i = 0; i < schema.n_children; ++i) {
330-
if (!schema.children[i] || !schema.children[i]->name) continue;
339+
if (!schema.children[i] || !schema.children[i]->name)
340+
continue;
331341
returnNames.push_back(schema.children[i]->name);
332342
returnTypes.push_back(ArrowConverter::fromArrowSchema(schema.children[i]));
333343
}
334344
returnNames.push_back("_score");
335345
returnTypes.push_back(LogicalType{LogicalTypeID::FLOAT});
336346

337-
if (schema.release) schema.release(&schema);
347+
if (schema.release)
348+
schema.release(&schema);
338349
auto columns = input->binder->createVariables(returnNames, returnTypes);
339350
return std::make_unique<LanceHybridSearchBindData>(std::move(resolvedPath),
340-
std::move(columnName), std::move(queryVec), k, std::move(ftsQuery),
341-
std::move(columns));
351+
std::move(columnName), std::move(queryVec), k, std::move(ftsQuery), std::move(columns));
342352
} catch (const lance::Error& e) {
343353
throw RuntimeException(std::string("LANCE_HYBRID_SEARCH bind failed: ") + e.what());
344354
}

0 commit comments

Comments
 (0)