diff --git a/docs/src/operations/ddl/.pages b/docs/src/operations/ddl/.pages index bd8c6daf1..f3ebe4923 100644 --- a/docs/src/operations/ddl/.pages +++ b/docs/src/operations/ddl/.pages @@ -11,6 +11,8 @@ nav: - show-tblproperties.md - drop-table.md - create-index.md + - refresh-index.md + - drop-index.md - show-indexes.md - create-branch.md - drop-branch.md diff --git a/docs/src/operations/ddl/create-index.md b/docs/src/operations/ddl/create-index.md index 03d1a1191..ecb26496a 100755 --- a/docs/src/operations/ddl/create-index.md +++ b/docs/src/operations/ddl/create-index.md @@ -53,7 +53,13 @@ The distributed build used by `zonemap`, `bitmap`, `label_list`, `ngram`, `bloom | Option | Type | Description | |----------------|---------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| -| `num_segments` | Integer | Target number of parallel build tasks (upper bound; clamped to fragment count when larger). Fragments are assigned by row count to balance estimated task workloads. Defaults to `min(fragment_count, spark.default.parallelism)`. | +| `num_segments` | Integer | Number of parallel build tasks, and so of index segments created (clamped to fragment count when larger). Each task takes a contiguous run of fragments, sized to balance estimated workloads by row count. Defaults to `min(fragment_count, spark.default.parallelism)`. | + +Option names are case-insensitive: `WITH (NUM_SEGMENTS = 8)` and `WITH (num_segments = 8)` are the same option. + +Contiguous coverage matters beyond parallelism: [OPTIMIZE](./optimize.md) can only group fragments +that the identical set of index segments covers, so segments whose fragment ids interleave leave it +nothing to coalesce. ### ZoneMap Options @@ -252,28 +258,30 @@ fill it in later, or when you intend to build it incrementally: A deferred index returns `fragments_indexed = 0` and is treated as fully unindexed: queries fall back to scanning the data until it is populated. There are two ways to populate it: -- **Full distributed build (recommended):** re-run `CREATE INDEX` with the same name. This uses the - normal distributed build across Spark tasks and atomically replaces the empty index: +- **[REFRESH INDEX](./refresh-index.md) (recommended):** builds only the fragments the index does not + cover, distributed across Spark tasks. For a deferred index that is the whole table; afterwards it + is whatever has been appended since. Repeat any method options here, since a deferred index holds + no built data to inherit them from: ```sql - ALTER TABLE lance.db.users CREATE INDEX idx_id USING zonemap (id); + ALTER TABLE lance.db.users REFRESH INDEX idx_id; ``` -- **Incremental build through the SDK:** when only some fragments are unindexed (for example after - appending data to an already-built index), `Dataset.optimizeIndices` indexes just the unindexed - fragments. This currently runs on a single node: +- **Re-run `CREATE INDEX`** with the same name. This rebuilds every fragment and atomically replaces + the existing index, which is what you want when you also need to change the index options or + consolidate accumulated segments: - ```java - dataset.optimizeIndices(OptimizeOptions.builder().build()); + ```sql + ALTER TABLE lance.db.users CREATE INDEX idx_id USING zonemap (id); ``` `train = false` is supported for all index methods. Because deferred index creation does not build -index data, `num_segments` cannot be combined with `train = false` — pass it on the eager build -that populates the index instead. +index data, `num_segments` cannot be combined with `train = false` — pass it on the build that +populates the index instead. Creating a scalar index on an empty table also registers an empty index with zero fragment coverage. The index is immediately visible through `SHOW INDEXES`. After data is appended, populate -it by re-running `CREATE INDEX` or calling `Dataset.optimizeIndices`. +it with `REFRESH INDEX`. ## Output @@ -305,4 +313,5 @@ The `CREATE INDEX` command operates as follows: - **Index Methods**: The `zonemap`, `bitmap`, `label_list`, `ngram`, `bloomfilter`, `rtree`, `btree`, and `fts` (or `inverted`) methods are supported for index creation. - **Indexed Column Count**: All supported index methods currently support exactly one indexed column. - **Index Replacement**: If you create an index with the same name as an existing one, the old index will be replaced by the new one. -- **Deferred Training**: With `train = false` the index is registered empty and is populated later, either by re-running `CREATE INDEX` (a full distributed build that replaces the empty index) or, for incremental coverage of newly appended fragments, by `Dataset.optimizeIndices` in the SDK. The SQL `OPTIMIZE` command compacts fragments and does not train deferred indexes. +- **Deferred Training**: With `train = false` the index is registered empty and is populated later, either by [REFRESH INDEX](./refresh-index.md) (a distributed build of the fragments the index does not cover) or by re-running `CREATE INDEX` (a full distributed rebuild that replaces the existing index). The SQL `OPTIMIZE` command compacts fragments and does not train deferred indexes. +- **Incremental Maintenance**: after appending data, [REFRESH INDEX](./refresh-index.md) indexes just the new fragments instead of rebuilding the table. diff --git a/docs/src/operations/ddl/refresh-index.md b/docs/src/operations/ddl/refresh-index.md new file mode 100644 index 000000000..62c724ed9 --- /dev/null +++ b/docs/src/operations/ddl/refresh-index.md @@ -0,0 +1,214 @@ +# REFRESH INDEX + +Index the fragments an existing index does not yet cover, without rebuilding the ones it does. + +!!! warning "Spark Extension Required" + This feature requires the Lance Spark SQL extension to be enabled. See [Spark SQL Extensions](../../config.md#spark-sql-extensions) for configuration details. + +## Overview + +Appending to a table creates new fragments, and an existing index does not cover them until it is +maintained. `REFRESH INDEX` builds index data only for those uncovered fragments and adds it to the +index, so cost scales with the newly appended data rather than with the table. + +This is the incremental counterpart to re-running [CREATE INDEX](./create-index.md), which rebuilds +every fragment. Both run distributed across Spark executors. + +## Syntax + +=== "SQL" + ```sql + ALTER TABLE lance.db.users REFRESH INDEX user_id_idx; + ``` + +## Options + +| Option | Type | Description | +|----------------|---------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `num_segments` | Integer | Number of parallel build tasks, and so of index segments added (clamped to the number of uncovered fragments when larger). Each task takes a contiguous run of fragments, sized to balance estimated workloads by row count. Defaults to `min(uncovered_fragments, spark.default.parallelism)`. | + +Option names are case-insensitive: `WITH (NUM_SEGMENTS = 8)` and `WITH (num_segments = 8)` are the same option. + +Index method options — such as `rows_per_zone` for `zonemap` or the tokenizer settings for `fts` — +are accepted as well, and are applied to the segments this command builds. See +[Index Method Options](#index-method-options) for when you need to pass them. + +## Examples + +### Refresh after appending data + +=== "SQL" + ```sql + INSERT INTO lance.db.users VALUES (100, 'new row'); + + ALTER TABLE lance.db.users REFRESH INDEX user_id_idx; + ``` + +### Refresh after compacting + +[OPTIMIZE](./optimize.md) replaces compacted fragments with new ones. A `zonemap` or `bloomfilter` +index records physical row addresses, which a rewrite invalidates, so its coverage drops and a +refresh restores it: + +=== "SQL" + ```sql + OPTIMIZE lance.db.users; + + ALTER TABLE lance.db.users REFRESH INDEX zone_idx; + ``` + +The other methods record row ids, which follow their rows through a rewrite, so compaction carries +their coverage over and a refresh afterwards has nothing to do. See +[Compaction](#notes-and-limitations) below. + +### Control build parallelism + +=== "SQL" + ```sql + ALTER TABLE lance.db.users REFRESH INDEX user_id_idx WITH (num_segments = 8); + ``` + +### Populate a deferred index + +An index created with `train = false` covers no fragments, so refreshing it builds it over the whole +table through the normal distributed path: + +=== "SQL" + ```sql + ALTER TABLE lance.db.users CREATE INDEX idx_id USING zonemap (id) WITH (train = false); + + ALTER TABLE lance.db.users REFRESH INDEX idx_id; + ``` + +The deferred index holds no built data to inherit method options from, so pass them to the refresh +that populates it — see [Index Method Options](#index-method-options). + +### Check coverage before and after + +[SHOW INDEXES](./show-indexes.md) reports how much of the table each index covers: + +=== "SQL" + ```sql + SHOW INDEXES FROM lance.db.users; + ``` + +A `num_unindexed_fragments` above zero means a refresh has work to do. + +## Output + +| Column | Type | Description | +|---------------------|--------|--------------------------------------------------------------------| +| `fragments_indexed` | Long | Number of previously uncovered fragments the commit covers. | +| `segments_added` | Long | Number of new physical index segments committed. | +| `index_name` | String | Name of the refreshed index, as stored in the table metadata. | + +A refresh with nothing to do returns zeros and commits no new table version. + +`fragments_indexed` counts the fragments actually covered, not the fragments planned: if a concurrent +operation retires one while the build runs, it is excluded and a warning names it. + +## How It Works + +1. **Planning**: the driver resolves the named index, subtracts its fragment coverage from the + table's fragments, and splits the remainder into batches, each a contiguous run of fragments + sized to balance row counts. Contiguity keeps [OPTIMIZE](./optimize.md) able to group the + fragments a segment covers. +2. **Distributed Build**: each batch becomes a Spark task that builds one uncommitted index + segment. Uncommitted segments are invisible to readers, so a failed build cannot affect query + results. +3. **Transactional Commit**: the driver commits the new segments as part of the same logical index. + Segments covering already-indexed fragments are retained, so existing coverage is preserved and + the change is atomic. + +Planning and building both run against a single pinned table version, so every task sees the +fragment set the driver planned over. + +## Index Method Options + +`REFRESH INDEX` builds its segments from the options in its own `WITH` clause; options you leave out +fall back to the index type's defaults rather than to whatever the index was built with. Lance +records a built index's parameters inside the index itself, not as table metadata a command can read +back, so pass the same options to the refresh that you passed to `CREATE INDEX`: + +=== "SQL" + ```sql + ALTER TABLE lance.db.users CREATE INDEX idx_id USING zonemap (id) WITH (rows_per_zone = 2048); + + ALTER TABLE lance.db.users REFRESH INDEX idx_id WITH (rows_per_zone = 2048); + ``` + +For most methods each segment is queried on its own and records its own configuration, so a mismatch +changes performance characteristics rather than results. + +`fts` (or `inverted`) is the exception: a full-text index is read with one configuration for all of +its segments, and a set that disagrees cannot be queried. `REFRESH INDEX` therefore compares what it +built against the segments it would join and **fails without committing** when they differ, leaving +the index exactly as it was: + +``` +Index 'idx_text' uses the fts method, whose segments must all share one configuration, and the +segments this build produced are configured differently from the ones they would join. Nothing was +committed and the index is unchanged. Re-run with the options the index was created with, or rebuild +it in full with ALTER TABLE ... CREATE INDEX. +``` + +Re-run with the original options, or rebuild with [CREATE INDEX](./create-index.md), which replaces +every segment and so has nothing to agree with. An `fts` index built by an older Lance version can +also record a configuration the current one does not reproduce even from the same options; a full +rebuild is the fix there. + +## Notes and Limitations + +- **Index Methods**: supported for the same methods as + [CREATE INDEX](./create-index.md#index-methods): `zonemap`, `btree`, `bitmap`, `label_list`, + `ngram`, `bloomfilter`, `rtree`, and `fts` (or `inverted`). A vector index cannot be created from + Spark SQL either, so refreshing one is rejected and pointed at the Lance SDK rather than at + `CREATE INDEX`, which could not build it. +- **Full-Build Options**: `train`, `build_mode`, and `rows_per_range` are rejected. Range-mode + `btree` redistributes and sorts the whole table, which an incremental refresh does not do — use + `CREATE INDEX` to rebuild that way. +- **Method Options Are Not Inherited**: new segments are built from the `WITH` clause and the index + type's defaults, not from the existing index's configuration. For `fts` a mismatch is rejected; + for the other methods it changes performance only. See + [Index Method Options](#index-method-options). +- **Indexes Created Before Coverage Tracking**: a segment that predates fragment coverage tracking + reports no coverage, so refreshing such an index rebuilds the whole table rather than a subset. + If the index mixes such a segment with a tracked one, the refresh fails rather than commit partial + coverage it cannot place; rebuild it with `CREATE INDEX`. +- **Multi-Field Indexes**: an index keyed on more than one field, or carrying columns beyond its key, + is rejected — a refreshed segment would declare only the key. Rebuild it with `CREATE INDEX`. +- **System Indexes**: Lance-maintained indexes, including fragment-reuse and MemWAL indexes, cannot + be refreshed. +- **Compaction**: [OPTIMIZE](./optimize.md) replaces the fragments it compacts with new ones. Whether + that costs the index its coverage depends on what the index stores. `zonemap` and `bloomfilter` + record physical row addresses, which the rewrite invalidates, so they lose coverage and need a + refresh afterwards. `btree`, `bitmap`, `label_list`, `ngram`, `rtree` and `fts` record row ids, + which follow their rows, so compaction carries their coverage over and a refresh reports + `fragments_indexed = 0`. +- **Stale `zonemap` Coverage Can Hide Rows**: a partially covered `zonemap` index prunes the + fragments it does not cover, so a predicate on the indexed column can return fewer rows than the + table holds (`COUNT(*)` over the whole table stays correct). Refresh a `zonemap` index after + appending or compacting before relying on filters over it. The other methods return complete + results while partially covered. +- **Concurrent Retirement**: if a concurrent operation retires a fragment while a refresh is + building it, that fragment is left out of the commit and named in a warning, and + `fragments_indexed` counts only what was covered. Re-run the refresh to pick up whatever replaced + it. A refresh whose fragments were *all* retired commits nothing and fails instead. +- **Concurrent `DROP INDEX`**: do not drop an index while a refresh of it is in flight. The refresh + re-resolves its target immediately before committing and fails if it is gone, but the check and the + commit are separate transactions and Lance does not currently treat a concurrent same-name drop as + a conflict. A `DROP INDEX` landing in that window is undone: the index reappears covering only the + fragments the refresh built. Tracked upstream in + [lance#6806](https://github.com/lance-format/lance/pull/6806). +- **Segment Growth**: each refresh adds segments to the index, and Lance can only compact fragments + that are covered by the identical set of index segments. So accumulated refreshes progressively + narrow what [OPTIMIZE](./optimize.md) can group: run `OPTIMIZE` before `REFRESH INDEX` rather than + after, and rebuild with [CREATE INDEX](./create-index.md) to consolidate the segments back into one + when the count grows. [SHOW INDEXES](./show-indexes.md) reports `num_segments`. + +## See Also + +- [CREATE INDEX](./create-index.md) +- [SHOW INDEXES](./show-indexes.md) +- [DROP INDEX](./drop-index.md) +- [OPTIMIZE](./optimize.md) diff --git a/docs/src/operations/ddl/show-indexes.md b/docs/src/operations/ddl/show-indexes.md index 7e47e7682..0ae4b87b0 100755 --- a/docs/src/operations/ddl/show-indexes.md +++ b/docs/src/operations/ddl/show-indexes.md @@ -54,12 +54,36 @@ The `SHOW INDEXES` command returns the following columns: | `num_indexed_rows` | long | Approximate number of rows covered by the index. | | `num_unindexed_fragments` | long | Number of fragments that are not yet indexed. | | `num_unindexed_rows` | long | Approximate number of rows that are not yet covered by the index. | +| `indexed_percent` | double | Share of rows the index covers, as a percentage truncated to two decimals, so it never overstates coverage. Null for an empty table. | +| `num_segments` | long | Number of physical index segments backing this logical index. | +| `size_bytes` | long | Total size of all index files across the segments. Null if any segment predates index file size tracking. | + +## Interpreting the Output + +An `indexed_percent` below 100 means part of the table is not covered, either because rows were +appended since the index was last built, or because [OPTIMIZE](./optimize.md) rewrote fragments a +`zonemap` or `bloomfilter` index had covered. For the index methods Spark SQL can build, use +[REFRESH INDEX](./refresh-index.md) to index just the uncovered fragments. A vector index is listed +here too, but Spark SQL can neither create nor refresh one; maintain it through the Lance SDK. + +Do not treat partial coverage as merely slower. For most methods the uncovered fragments are scanned +and results stay complete, but a partially covered `zonemap` index prunes them instead, so a +predicate on the indexed column can return fewer rows than the table holds. Refresh before relying on +filters over a `zonemap` index that reports less than 100. + +`num_segments` reflects how the index was built: a distributed build produces one segment per +parallel task, and each [REFRESH INDEX](./refresh-index.md) adds more, since a refresh appends +coverage rather than rewriting existing segments. Queries search every segment, and Lance can only +compact fragments covered by the identical set of segments, so a high count costs both query time and +`OPTIMIZE`'s ability to coalesce. Rebuilding with [CREATE INDEX](./create-index.md) consolidates them. ## Notes - The `fields` column returns the logical column names from the Lance schema, ordered according to the index definition. - Lance-maintained system indexes, including fragment-reuse and MemWAL indexes, are excluded from the output. +- Row counts are approximate, so `indexed_percent` is a guide rather than an exact figure. ## See Also - [CREATE INDEX](./create-index.md) +- [REFRESH INDEX](./refresh-index.md) diff --git a/integration-tests/test_lance_spark.py b/integration-tests/test_lance_spark.py index 9c799ef3b..2f66462ec 100644 --- a/integration-tests/test_lance_spark.py +++ b/integration-tests/test_lance_spark.py @@ -1138,6 +1138,203 @@ def test_drop_index_then_recreate(self, spark): assert len(query_result) == 1 assert query_result[0].id == 50 + def test_refresh_index_covers_appended_fragments(self, spark): + """Test REFRESH INDEX indexes only the fragments added after index creation.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING, + value DOUBLE + ) + """) + + def append(start, end): + data = [(i, f"Name{i}", float(i * 10)) for i in range(start, end)] + df = spark.createDataFrame(data, ["id", "name", "value"]) + df.coalesce(1).writeTo("default.test_table").append() + + append(0, 50) + spark.sql(""" + ALTER TABLE default.test_table + CREATE INDEX idx_id USING zonemap (id) + """) + + indexed = spark.sql("SHOW INDEXES IN default.test_table").collect()[0] + assert indexed["num_unindexed_fragments"] == 0 + assert indexed["indexed_percent"] == 100.0 + + # Two more fragments leave the index behind + append(50, 100) + append(100, 150) + + stale = spark.sql("SHOW INDEXES IN default.test_table").collect()[0] + assert stale["num_unindexed_fragments"] == 2 + assert stale["indexed_percent"] < 100.0 + + result = spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX idx_id + """).collect() + assert len(result) == 1 + assert result[0]["fragments_indexed"] == 2 + assert result[0]["segments_added"] >= 1 + assert result[0]["index_name"] == "idx_id" + + refreshed = spark.sql("SHOW INDEXES IN default.test_table").collect()[0] + assert refreshed["num_unindexed_fragments"] == 0 + assert refreshed["indexed_percent"] == 100.0 + + # Data in the newly indexed fragments stays queryable + assert spark.table("default.test_table").count() == 150 + query_result = spark.sql(""" + SELECT * FROM default.test_table WHERE id = 125 + """).collect() + assert len(query_result) == 1 + assert query_result[0].id == 125 + + def test_refresh_index_is_noop_when_fully_indexed(self, spark): + """Test REFRESH INDEX reports no work when the index already covers the table.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING, + value DOUBLE + ) + """) + + data = [(i, f"Name{i}", float(i * 10)) for i in range(100)] + df = spark.createDataFrame(data, ["id", "name", "value"]) + df.writeTo("default.test_table").append() + + spark.sql(""" + ALTER TABLE default.test_table + CREATE INDEX idx_id USING btree (id) + """) + + result = spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX idx_id + """).collect() + assert result[0]["fragments_indexed"] == 0 + assert result[0]["segments_added"] == 0 + + def test_refresh_index_populates_deferred_index(self, spark): + """Test REFRESH INDEX builds an index registered with train = false.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING, + value DOUBLE + ) + """) + + data = [(i, f"Name{i}", float(i * 10)) for i in range(100)] + df = spark.createDataFrame(data, ["id", "name", "value"]) + df.writeTo("default.test_table").append() + + spark.sql(""" + ALTER TABLE default.test_table + CREATE INDEX idx_id USING zonemap (id) WITH (train = false) + """) + + deferred = spark.sql("SHOW INDEXES IN default.test_table").collect()[0] + assert deferred["num_indexed_rows"] == 0 + + result = spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX idx_id + """).collect() + assert result[0]["fragments_indexed"] >= 1 + + populated = spark.sql("SHOW INDEXES IN default.test_table").collect()[0] + assert populated["num_unindexed_rows"] == 0 + assert populated["indexed_percent"] == 100.0 + + def test_refresh_index_named_after_its_column(self, spark): + """Test REFRESH INDEX works for an index named like Lance's own default, _idx.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING, + value DOUBLE + ) + """) + + def append(start, end): + data = [(i, f"Name{i}", float(i * 10)) for i in range(start, end)] + df = spark.createDataFrame(data, ["id", "name", "value"]) + df.coalesce(1).writeTo("default.test_table").append() + + append(0, 50) + spark.sql(""" + ALTER TABLE default.test_table + CREATE INDEX id_idx USING zonemap (id) + """) + append(50, 100) + + result = spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX id_idx + """).collect() + assert result[0]["fragments_indexed"] == 1 + assert result[0]["index_name"] == "id_idx" + + refreshed = spark.sql("SHOW INDEXES IN default.test_table").collect()[0] + assert refreshed["num_unindexed_fragments"] == 0 + query_result = spark.sql(""" + SELECT * FROM default.test_table WHERE id = 75 + """).collect() + assert len(query_result) == 1 + + def test_refresh_index_option_names_are_case_insensitive(self, spark): + """Test an upper-case option name takes effect instead of being silently ignored.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING, + value DOUBLE + ) + """) + + def append(start, end): + data = [(i, f"Name{i}", float(i * 10)) for i in range(start, end)] + df = spark.createDataFrame(data, ["id", "name", "value"]) + df.coalesce(1).writeTo("default.test_table").append() + + append(0, 10) + spark.sql(""" + ALTER TABLE default.test_table + CREATE INDEX idx_id USING zonemap (id) + """) + for start in range(10, 50, 10): + append(start, start + 10) + + result = spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX idx_id WITH (NUM_SEGMENTS = 2) + """).collect() + assert result[0]["fragments_indexed"] == 4 + assert result[0]["segments_added"] == 2 + + with pytest.raises(Exception, match="not supported for REFRESH INDEX"): + spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX idx_id WITH (TRAIN = false) + """).collect() + + def test_refresh_index_rejects_unknown_index(self, spark): + """Test REFRESH INDEX fails with an actionable message for a missing index.""" + spark.sql(""" + CREATE TABLE default.test_table ( + id INT, + name STRING, + value DOUBLE + ) + """) + + data = [(i, f"Name{i}", float(i * 10)) for i in range(10)] + df = spark.createDataFrame(data, ["id", "name", "value"]) + df.writeTo("default.test_table").append() + + with pytest.raises(Exception, match="does not exist"): + spark.sql(""" + ALTER TABLE default.test_table REFRESH INDEX missing_idx + """).collect() + class TestDDLOptimize: """Test DDL OPTIMIZE operations for compacting table fragments.""" @@ -3521,6 +3718,65 @@ def test_lance_multi_match_with_operator(self, spark): assert 4 in ids assert 6 in ids + # The fixture's index uses non-default tokenizer options, so a refresh that omits them builds + # segments the read path cannot combine with the existing ones. REFRESH INDEX must refuse that + # rather than commit an index that reports full coverage and answers nothing. + FTS_INDEX_OPTIONS = ( + "base_tokenizer = 'simple', language = 'English', " + "max_token_length = 40, lower_case = true, " + "stem = false, remove_stop_words = false, " + "ascii_folding = false, with_position = true" + ) + + def _append_fts_doc(self, spark): + df = spark.createDataFrame( + [(7, "Refreshed Doc", "Apache Spark refreshed segment body")], + ["id", "title", "body"], + ) + df.coalesce(1).writeTo("default.fts_docs").append() + + def test_refresh_fts_index_rejects_mismatched_options(self, spark): + """REFRESH INDEX refuses to mix FTS segments built with different options.""" + self._append_fts_doc(spark) + + with pytest.raises(Exception, match="share one configuration"): + spark.sql( + "ALTER TABLE default.fts_docs REFRESH INDEX fts_body" + ).collect() + + # Nothing was committed, so the index still answers exactly as before. + rows = spark.sql( + "SELECT id FROM default.fts_docs WHERE lance_match(body, 'spark')" + ).collect() + ids = sorted([r.id for r in rows]) + assert 1 in ids + assert 4 in ids + + def test_refresh_fts_index_with_matching_options(self, spark): + """Repeating the original options lets the refresh commit and keeps queries complete.""" + self._append_fts_doc(spark) + + result = spark.sql( + "ALTER TABLE default.fts_docs REFRESH INDEX fts_body " + f"WITH ({self.FTS_INDEX_OPTIONS})" + ).collect() + assert result[0]["fragments_indexed"] >= 1 + + stats = spark.sql("SHOW INDEXES IN default.fts_docs").collect() + body = [r for r in stats if r["name"] == "fts_body"][0] + assert body["num_unindexed_fragments"] == 0 + + rows = spark.sql( + "SELECT id FROM default.fts_docs WHERE lance_match(body, 'refreshed')" + ).collect() + assert [r.id for r in rows] == [7] + # Positions survive, so phrase queries keep working across the new segment. + rows = spark.sql( + "SELECT id FROM default.fts_docs " + "WHERE lance_match_phrase(body, 'refreshed segment')" + ).collect() + assert [r.id for r in rows] == [7] + def test_show_functions_lists_fts(self, spark): """SHOW FUNCTIONS returns all three FTS function names.""" functions = spark.sql("SHOW FUNCTIONS").collect() diff --git a/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 1ca2fb957..95573f322 100644 --- a/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -16,9 +16,11 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, RefreshIndex, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} +import java.util.Locale + import scala.collection.JavaConverters._ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) @@ -26,6 +28,16 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) private def cleanIdentifier(text: String): String = ParserUtils.cleanIdentifier(text) + /** + * A WITH-clause option name, normalized to lower case. + * + * ANTLR reports identifier text as written, and every command matches its option names against + * lower-case literals, so without normalizing here `WITH (NUM_SEGMENTS = 4)` parses into an option + * no command recognizes: it is silently ignored, and forwarded to Lance as an index parameter. + */ + private def normalizedOptionName(text: String): String = + cleanIdentifier(text).toLowerCase(Locale.ROOT) + override def visitSingleStatement(ctx: LanceSqlExtensionsParser.SingleStatementContext) : LogicalPlan = { visit(ctx.statement).asInstanceOf[LogicalPlan] @@ -73,7 +85,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -84,7 +96,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -98,7 +110,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val columns = visitFieldPathList(ctx.fieldPathList()) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -116,6 +128,19 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) LanceDropIndex(table, indexName) } + override def visitRefreshIndex(ctx: LanceSqlExtensionsParser.RefreshIndexContext) + : RefreshIndex = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + val indexName = cleanIdentifier(ctx.indexName.getText) + val args = ctx.namedArgument().asScala.map(a => + LanceNamedArgument( + normalizedOptionName(a.identifier().getText), + a.constant().accept(this))) + .toSeq + + RefreshIndex(table, indexName, args) + } + override def visitCreateBranchRefMain(ctx: LanceSqlExtensionsParser.CreateBranchRefMainContext) : LanceCreateBranch = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java index 691284eca..d090c6851 100644 --- a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java @@ -23,6 +23,7 @@ import org.apache.spark.sql.catalyst.plans.logical.AddColumnsBackfill; import org.apache.spark.sql.catalyst.plans.logical.AddIndex; import org.apache.spark.sql.catalyst.plans.logical.Optimize; +import org.apache.spark.sql.catalyst.plans.logical.RefreshIndex; import org.apache.spark.sql.catalyst.plans.logical.ShowIndexes; import org.apache.spark.sql.catalyst.plans.logical.UpdateColumnsBackfill; import org.apache.spark.sql.catalyst.plans.logical.Vacuum; @@ -142,6 +143,36 @@ public void testCreateIndexWithBacktickedIdentifiers() { assertEquals(List.of("`col-a`"), JavaConverters.seqAsJavaList(plan.columns())); } + @Test + public void testRefreshIndexWithBacktickedIdentifiers() { + // Keywords must be uppercase here: the grammar's LETTER fragment is [A-Z], and this test + // feeds the lexer a raw stream rather than the UpperCaseCharStream the session parser uses. + LanceSqlExtensionsParser parser = + createParser( + "ALTER TABLE `my-catalog`.`my-table` REFRESH INDEX `my-idx` WITH (NUM_SEGMENTS = 4)"); + RefreshIndex plan = (RefreshIndex) astBuilder.visitSingleStatement(parser.singleStatement()); + + UnresolvedIdentifier table = (UnresolvedIdentifier) plan.table(); + assertEquals( + List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts())); + assertEquals("my-idx", plan.indexName()); + assertEquals(1, plan.args().size()); + // Option names are normalized at parse time: commands match lower-case literals, so an + // upper-case spelling must reach them as the option they recognize rather than as an unknown + // one. + assertEquals("num_segments", plan.args().apply(0).name()); + assertEquals(4L, plan.args().apply(0).value()); + } + + @Test + public void testRefreshIndexWithoutOptions() { + LanceSqlExtensionsParser parser = createParser("ALTER TABLE DB.T REFRESH INDEX IDX"); + RefreshIndex plan = (RefreshIndex) astBuilder.visitSingleStatement(parser.singleStatement()); + + assertEquals("IDX", plan.indexName()); + assertEquals(0, plan.args().size()); + } + @Test public void testOptimizeWithBacktickedTableName() { LanceSqlExtensionsParser parser = createParser("OPTIMIZE `my-catalog`.`my-table`"); diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/RefreshIndexTest.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/RefreshIndexTest.java new file mode 100644 index 000000000..d9491f5d8 --- /dev/null +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/update/RefreshIndexTest.java @@ -0,0 +1,16 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.update; + +public class RefreshIndexTest extends BaseRefreshIndexTest {} diff --git a/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 40a228e97..858578d35 100644 --- a/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -16,9 +16,11 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, RefreshIndex, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} +import java.util.Locale + import scala.collection.JavaConverters._ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) @@ -26,6 +28,16 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) private def cleanIdentifier(text: String): String = ParserUtils.cleanIdentifier(text) + /** + * A WITH-clause option name, normalized to lower case. + * + * ANTLR reports identifier text as written, and every command matches its option names against + * lower-case literals, so without normalizing here `WITH (NUM_SEGMENTS = 4)` parses into an option + * no command recognizes: it is silently ignored, and forwarded to Lance as an index parameter. + */ + private def normalizedOptionName(text: String): String = + cleanIdentifier(text).toLowerCase(Locale.ROOT) + override def visitSingleStatement(ctx: LanceSqlExtensionsParser.SingleStatementContext) : LogicalPlan = { visit(ctx.statement).asInstanceOf[LogicalPlan] @@ -73,7 +85,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -84,7 +96,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -98,7 +110,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val columns = visitFieldPathList(ctx.fieldPathList()) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -116,6 +128,19 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) LanceDropIndex(table, indexName) } + override def visitRefreshIndex(ctx: LanceSqlExtensionsParser.RefreshIndexContext) + : RefreshIndex = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + val indexName = cleanIdentifier(ctx.indexName.getText) + val args = ctx.namedArgument().asScala.map(a => + LanceNamedArgument( + normalizedOptionName(a.identifier().getText), + a.constant().accept(this))) + .toSeq + + RefreshIndex(table, indexName, args) + } + override def visitCreateBranchRefMain(ctx: LanceSqlExtensionsParser.CreateBranchRefMainContext) : LanceCreateBranch = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java index d387087e9..47895a360 100644 --- a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/LanceSqlExtensionsAstBuilderTest.java @@ -26,6 +26,7 @@ import org.apache.spark.sql.catalyst.plans.logical.LanceDropBranch; import org.apache.spark.sql.catalyst.plans.logical.LanceShowBranches; import org.apache.spark.sql.catalyst.plans.logical.Optimize; +import org.apache.spark.sql.catalyst.plans.logical.RefreshIndex; import org.apache.spark.sql.catalyst.plans.logical.ShowIndexes; import org.apache.spark.sql.catalyst.plans.logical.UpdateColumnsBackfill; import org.apache.spark.sql.catalyst.plans.logical.Vacuum; @@ -166,6 +167,36 @@ public void testCreateIndexWithNestedFieldPath() { JavaConverters.seqAsJavaList(plan.columns())); } + @Test + public void testRefreshIndexWithBacktickedIdentifiers() { + // Keywords must be uppercase here: the grammar's LETTER fragment is [A-Z], and this test + // feeds the lexer a raw stream rather than the UpperCaseCharStream the session parser uses. + LanceSqlExtensionsParser parser = + createParser( + "ALTER TABLE `my-catalog`.`my-table` REFRESH INDEX `my-idx` WITH (NUM_SEGMENTS = 4)"); + RefreshIndex plan = (RefreshIndex) astBuilder.visitSingleStatement(parser.singleStatement()); + + UnresolvedIdentifier table = (UnresolvedIdentifier) plan.table(); + assertEquals( + List.of("my-catalog", "my-table"), JavaConverters.seqAsJavaList(table.nameParts())); + assertEquals("my-idx", plan.indexName()); + assertEquals(1, plan.args().size()); + // Option names are normalized at parse time: commands match lower-case literals, so an + // upper-case spelling must reach them as the option they recognize rather than as an unknown + // one. + assertEquals("num_segments", plan.args().apply(0).name()); + assertEquals(4L, plan.args().apply(0).value()); + } + + @Test + public void testRefreshIndexWithoutOptions() { + LanceSqlExtensionsParser parser = createParser("ALTER TABLE DB.T REFRESH INDEX IDX"); + RefreshIndex plan = (RefreshIndex) astBuilder.visitSingleStatement(parser.singleStatement()); + + assertEquals("IDX", plan.indexName()); + assertEquals(0, plan.args().size()); + } + @Test public void testOptimizeWithBacktickedTableName() { LanceSqlExtensionsParser parser = createParser("OPTIMIZE `my-catalog`.`my-table`"); diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/RefreshIndexTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/RefreshIndexTest.java new file mode 100644 index 000000000..d9491f5d8 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/update/RefreshIndexTest.java @@ -0,0 +1,16 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.update; + +public class RefreshIndexTest extends BaseRefreshIndexTest {} diff --git a/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 86068fd67..d7d00d1fe 100644 --- a/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -16,9 +16,11 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, RefreshIndex, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} +import java.util.Locale + import scala.jdk.CollectionConverters._ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) @@ -26,6 +28,16 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) private def cleanIdentifier(text: String): String = ParserUtils.cleanIdentifier(text) + /** + * A WITH-clause option name, normalized to lower case. + * + * ANTLR reports identifier text as written, and every command matches its option names against + * lower-case literals, so without normalizing here `WITH (NUM_SEGMENTS = 4)` parses into an option + * no command recognizes: it is silently ignored, and forwarded to Lance as an index parameter. + */ + private def normalizedOptionName(text: String): String = + cleanIdentifier(text).toLowerCase(Locale.ROOT) + override def visitSingleStatement(ctx: LanceSqlExtensionsParser.SingleStatementContext) : LogicalPlan = { visit(ctx.statement).asInstanceOf[LogicalPlan] @@ -73,7 +85,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -84,7 +96,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -98,7 +110,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val columns = visitFieldPathList(ctx.fieldPathList()) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -116,6 +128,19 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) LanceDropIndex(table, indexName) } + override def visitRefreshIndex(ctx: LanceSqlExtensionsParser.RefreshIndexContext) + : RefreshIndex = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + val indexName = cleanIdentifier(ctx.indexName.getText) + val args = ctx.namedArgument().asScala.map(a => + LanceNamedArgument( + normalizedOptionName(a.identifier().getText), + a.constant().accept(this))) + .toSeq + + RefreshIndex(table, indexName, args) + } + override def visitCreateBranchRefMain(ctx: LanceSqlExtensionsParser.CreateBranchRefMainContext) : LanceCreateBranch = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) diff --git a/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 86068fd67..d7d00d1fe 100644 --- a/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -16,9 +16,11 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, RefreshIndex, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} +import java.util.Locale + import scala.jdk.CollectionConverters._ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) @@ -26,6 +28,16 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) private def cleanIdentifier(text: String): String = ParserUtils.cleanIdentifier(text) + /** + * A WITH-clause option name, normalized to lower case. + * + * ANTLR reports identifier text as written, and every command matches its option names against + * lower-case literals, so without normalizing here `WITH (NUM_SEGMENTS = 4)` parses into an option + * no command recognizes: it is silently ignored, and forwarded to Lance as an index parameter. + */ + private def normalizedOptionName(text: String): String = + cleanIdentifier(text).toLowerCase(Locale.ROOT) + override def visitSingleStatement(ctx: LanceSqlExtensionsParser.SingleStatementContext) : LogicalPlan = { visit(ctx.statement).asInstanceOf[LogicalPlan] @@ -73,7 +85,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -84,7 +96,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -98,7 +110,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val columns = visitFieldPathList(ctx.fieldPathList()) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -116,6 +128,19 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) LanceDropIndex(table, indexName) } + override def visitRefreshIndex(ctx: LanceSqlExtensionsParser.RefreshIndexContext) + : RefreshIndex = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + val indexName = cleanIdentifier(ctx.indexName.getText) + val args = ctx.namedArgument().asScala.map(a => + LanceNamedArgument( + normalizedOptionName(a.identifier().getText), + a.constant().accept(this))) + .toSeq + + RefreshIndex(table, indexName, args) + } + override def visitCreateBranchRefMain(ctx: LanceSqlExtensionsParser.CreateBranchRefMainContext) : LanceCreateBranch = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) diff --git a/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 86068fd67..d7d00d1fe 100644 --- a/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -16,9 +16,11 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, RefreshIndex, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} +import java.util.Locale + import scala.jdk.CollectionConverters._ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) @@ -26,6 +28,16 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) private def cleanIdentifier(text: String): String = ParserUtils.cleanIdentifier(text) + /** + * A WITH-clause option name, normalized to lower case. + * + * ANTLR reports identifier text as written, and every command matches its option names against + * lower-case literals, so without normalizing here `WITH (NUM_SEGMENTS = 4)` parses into an option + * no command recognizes: it is silently ignored, and forwarded to Lance as an index parameter. + */ + private def normalizedOptionName(text: String): String = + cleanIdentifier(text).toLowerCase(Locale.ROOT) + override def visitSingleStatement(ctx: LanceSqlExtensionsParser.SingleStatementContext) : LogicalPlan = { visit(ctx.statement).asInstanceOf[LogicalPlan] @@ -73,7 +85,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -84,7 +96,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -98,7 +110,7 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) val columns = visitFieldPathList(ctx.fieldPathList()) val args = ctx.namedArgument().asScala.map(a => LanceNamedArgument( - cleanIdentifier(a.identifier().getText), + normalizedOptionName(a.identifier().getText), a.constant().accept(this))) .toSeq @@ -116,6 +128,19 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) LanceDropIndex(table, indexName) } + override def visitRefreshIndex(ctx: LanceSqlExtensionsParser.RefreshIndexContext) + : RefreshIndex = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + val indexName = cleanIdentifier(ctx.indexName.getText) + val args = ctx.namedArgument().asScala.map(a => + LanceNamedArgument( + normalizedOptionName(a.identifier().getText), + a.constant().accept(this))) + .toSeq + + RefreshIndex(table, indexName, args) + } + override def visitCreateBranchRefMain(ctx: LanceSqlExtensionsParser.CreateBranchRefMainContext) : LanceCreateBranch = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) diff --git a/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 b/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 index 788525249..416056c5a 100644 --- a/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 +++ b/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 @@ -23,6 +23,7 @@ statement | ALTER TABLE multipartIdentifier UPDATE COLUMNS columnList FROM identifier #updateColumnsBackfill | ALTER TABLE multipartIdentifier CREATE INDEX indexName=identifier USING method=identifier '(' fieldPathList ')' (WITH '(' (namedArgument (',' namedArgument)*)? ')')? #createIndex | ALTER TABLE multipartIdentifier DROP INDEX indexName=identifier #dropIndex + | ALTER TABLE multipartIdentifier REFRESH INDEX indexName=identifier (WITH '(' (namedArgument (',' namedArgument)*)? ')')? #refreshIndex | ALTER TABLE multipartIdentifier CREATE BRANCH (IF NOT EXISTS)? branchName=identifier (AS OF VERSION refMainVersion=versionNumber)? #createBranchRefMain | ALTER TABLE multipartIdentifier CREATE BRANCH (IF NOT EXISTS)? branchName=identifier @@ -113,6 +114,7 @@ NOT: 'NOT'; OF: 'OF'; OPTIMIZE: 'OPTIMIZE'; PRIMARY: 'PRIMARY'; +REFRESH: 'REFRESH'; SET: 'SET'; SHOW: 'SHOW'; TABLE: 'TABLE'; diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/RefreshIndex.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/RefreshIndex.scala new file mode 100644 index 000000000..58f37d5e9 --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/RefreshIndex.scala @@ -0,0 +1,51 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.catalyst.plans.logical + +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.types.{DataTypes, StructField, StructType} + +/** + * RefreshIndex logical plan representing incremental index maintenance on a Lance dataset. + * + * Unlike [[AddIndex]], which rebuilds every fragment, this command builds index segments only for + * the fragments an existing index does not yet cover, and adds them to that index. + */ +case class RefreshIndex( + table: LogicalPlan, + indexName: String, + args: Seq[LanceNamedArgument]) extends Command { + + override def children: Seq[LogicalPlan] = Seq(table) + + override def output: Seq[Attribute] = RefreshIndexOutputType.SCHEMA + + override def simpleString(maxFields: Int): String = { + s"RefreshIndex(${indexName})" + } + + override protected def withNewChildrenInternal(newChildren: IndexedSeq[LogicalPlan]) + : RefreshIndex = { + copy(newChildren(0), this.indexName, this.args) + } +} + +object RefreshIndexOutputType { + val SCHEMA: Seq[Attribute] = StructType( + Array( + StructField("fragments_indexed", DataTypes.LongType, nullable = true), + StructField("segments_added", DataTypes.LongType, nullable = true), + StructField("index_name", DataTypes.StringType, nullable = true))) + .map(field => AttributeReference(field.name, field.dataType, field.nullable, field.metadata)()) +} diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ShowIndexes.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ShowIndexes.scala index 938912bec..b4f28c29a 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ShowIndexes.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ShowIndexes.scala @@ -47,6 +47,9 @@ object ShowIndexesOutputType { StructField("num_indexed_fragments", DataTypes.LongType, nullable = true), StructField("num_indexed_rows", DataTypes.LongType, nullable = true), StructField("num_unindexed_fragments", DataTypes.LongType, nullable = true), - StructField("num_unindexed_rows", DataTypes.LongType, nullable = true))) + StructField("num_unindexed_rows", DataTypes.LongType, nullable = true), + StructField("indexed_percent", DataTypes.DoubleType, nullable = true), + StructField("num_segments", DataTypes.LongType, nullable = true), + StructField("size_bytes", DataTypes.LongType, nullable = true))) .map(field => AttributeReference(field.name, field.dataType, field.nullable, field.metadata)()) } diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index 085090996..85845ffb5 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -18,6 +18,7 @@ import com.fasterxml.jackson.databind.node.ObjectNode import org.apache.arrow.c.{ArrowArrayStream, Data} import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.ipc.ArrowReader +import org.apache.spark.SparkContext import org.apache.spark.internal.Logging import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow} @@ -27,7 +28,7 @@ import org.apache.spark.sql.types.StructType import org.apache.spark.sql.util.LanceArrowUtils import org.apache.spark.sql.util.LanceSerializeUtil.{decode, encode} import org.apache.spark.unsafe.types.UTF8String -import org.lance.{CommitBuilder, Dataset, Transaction} +import org.lance.Dataset import org.lance.index.{Index, IndexOptions, IndexParams, IndexType} import org.lance.index.scalar.{BTreeIndexParams, ScalarIndexParams} import org.lance.operation.{CreateIndex => AddIndexOperation} @@ -60,10 +61,10 @@ import scala.reflect.ClassTag *

Deferred training ({@code WITH (train=false)}): commits an empty index on the driver * with an empty fragment bitmap (all rows appear unindexed), skipping data processing. Supported * for all supported scalar index methods. Empty tables use the same path even when - * {@code train=true}, since there are no fragments to train. Populate the index later by re-running - * {@code CREATE INDEX} with the same name (a full distributed build that replaces the empty index) - * or, for incremental coverage of appended fragments, by {@code Dataset.optimizeIndices} (the SQL - * {@code OPTIMIZE} only compacts fragments). {@code num_segments} is rejected with + * {@code train=true}, since there are no fragments to train. Populate the index later with + * {@code REFRESH INDEX} (a distributed build of the fragments the index does not cover) or by + * re-running {@code CREATE INDEX} with the same name (a full distributed rebuild); the SQL + * {@code OPTIMIZE} only compacts fragments. {@code num_segments} is rejected with * {@code train=false}, since no segmented build occurs. * *

The following options are consumed at the Spark execution layer and are never forwarded @@ -88,7 +89,14 @@ case class AddIndexExec( val btreeBuildMode = IndexUtils.btreeBuildMode(indexType, args) val scalarSegmentIndexType = IndexUtils.scalarSegmentIndexType(method) - val (fragmentWorkloads, canonicalColumns) = { + // Plan and build against a single pinned version. Tasks open the dataset themselves, so without + // pinning each one resolves the latest version independently. The driver's batches fix which + // fragments a segment covers either way, but not the version each segment records, and that + // version is what core checks: it validates a segment's coverage only when the stamp predates + // the commit, so a task that opens after a concurrent rewrite of the indexed column stamps the + // current version over keys it read earlier and the stale keys are trusted. Coverage the commit + // cannot establish is accounted for at commit time; see IndexUtils.establishedCoverage. + val (fragmentWorkloads, canonicalColumns, buildReadOptions) = { val ds = Utils.openDatasetBuilder(readOptions).build() try { val canonical = columns.map { column => @@ -96,10 +104,9 @@ case class AddIndexExec( FieldPathUtils.pathByFieldId(ds.getLanceSchema, field.getId) } ( - ds.getFragments.asScala - .map(fragment => FragmentWorkload(fragment.getId, fragment.metadata().getNumRows)) - .toList, - canonical) + IndexUtils.fragmentWorkloads(ds), + canonical, + IndexUtils.pinVersion(readOptions, ds)) } finally { ds.close() } @@ -126,22 +133,7 @@ case class AddIndexExec( throw new IllegalArgumentException( "num_segments is not supported with train=false: a deferred index performs no segmented build") } - val validatedNumSegments: Option[Int] = numSegmentsOpt.map { arg => - arg.value match { - case null => - throw new IllegalArgumentException( - "num_segments must be a positive integer, got: null") - case n: Number => - val asLong = n.longValue() - if (asLong < 1L || asLong > Int.MaxValue) - throw new IllegalArgumentException( - s"num_segments must be a positive integer that fits in Int, got: $asLong") - asLong.toInt - case other => - throw new IllegalArgumentException( - s"num_segments must be a positive integer, got: $other") - } - } + val validatedNumSegments: Option[Int] = numSegmentsOpt.map(IndexUtils.parseNumSegments) // train=false, or an empty table: commit an empty index on the driver and skip data // processing. Index and option validation above still applies to empty tables. @@ -161,29 +153,37 @@ case class AddIndexExec( } val (nsImpl, nsProps, tableId, initialStorageOpts) = - extractNamespaceInfo(lanceDataset, readOptions) + IndexUtils.extractNamespaceInfo(catalog, lanceDataset, readOptions) - // Range-mode BTree uses preprocessed data from Spark and keeps its dedicated path. + // Range-mode BTree uses preprocessed data from Spark and keeps its dedicated path: its coverage + // follows the fragment ids in the scanned rows rather than the fragment list planned above. The + // build is pinned like the segmented one all the same. The scan resolves its own version through + // the catalog, so an unpinned executor can open a version newer than the rows it was handed and + // stamp the segment with it, and core only validates segments stamped older than the commit. if (btreeBuildMode.contains("range")) { val segments = new RangeBasedBTreeIndexJob( this.copy(columns = canonicalColumns), - readOptions, + buildReadOptions, fragmentIds.size, nsImpl, nsProps, tableId, initialStorageOpts).run() - commitIndexSegments(readOptions, canonicalColumns.head, segments) + val indexed = commitIndexSegments(readOptions, canonicalColumns.head, segments) return Seq(new GenericInternalRow(Array[Any]( - fragmentIds.size.toLong, + indexed.toLong, UTF8String.fromString(indexName)))) } // Scalar segment indexes use the logical segment commit path. if (scalarSegmentIndexType.isDefined) { val segmentJob = new ScalarSegmentIndexJob( - this.copy(columns = canonicalColumns), - readOptions, + session.sparkContext, + indexName, + method, + canonicalColumns.toList, + IndexUtils.toJson(args), + buildReadOptions, fragmentWorkloads, validatedNumSegments, nsImpl, @@ -192,9 +192,9 @@ case class AddIndexExec( initialStorageOpts) val segments = segmentJob.run() // Atomic add+remove via Lance core; see commitIndexSegments - commitIndexSegments(readOptions, canonicalColumns.head, segments) + val indexed = commitIndexSegments(readOptions, canonicalColumns.head, segments) return Seq(new GenericInternalRow(Array[Any]( - fragmentIds.size.toLong, + indexed.toLong, UTF8String.fromString(indexName)))) } @@ -231,39 +231,36 @@ case class AddIndexExec( // Lance core's commitExistingIndexSegments handles atomic replacement: // it finds existing segments whose fragments overlap with incoming ones // and removes them in the same CreateIndex transaction. + // + // Returns the number of fragments the commit actually covers, which is what the command reports. private def commitIndexSegments( readOptions: LanceSparkReadOptions, column: String, - segments: Seq[Index]): Unit = { + segments: Seq[Index]): Int = { val dataset = Utils.openDatasetBuilder(readOptions).build() try { - dataset.commitExistingIndexSegments( + IndexUtils.requireCommittableCoverage( + IndexUtils.liveFragmentIds(dataset), + segments, + indexName) + val committed = dataset.commitExistingIndexSegments( indexName, column, segments.toList.asJava) + // The commit advances this handle to the manifest it wrote, so both the returned metadata and + // the fragment list below describe the committed state rather than the one validated above. + IndexUtils + .establishedCoverage( + segments, + committed.asScala.toSeq, + IndexUtils.liveFragmentIds(dataset), + indexName) + .size } finally { dataset.close() } } - private def extractNamespaceInfo( - lanceDataset: LanceDataset, - readOptions: LanceSparkReadOptions): ( - Option[String], - Option[Map[String, String]], - Option[List[String]], - Option[Map[String, String]]) = { - catalog match { - case nsCatalog: BaseLanceNamespaceSparkCatalog => - ( - Option(nsCatalog.getNamespaceImpl), - Option(nsCatalog.getNamespaceProperties).map(_.asScala.toMap), - Option(readOptions.getTableId).map(_.asScala.toList), - Option(lanceDataset.getInitialStorageOptions).map(_.asScala.toMap)) - case _ => (None, None, None, None) - } - } - } /** @@ -276,6 +273,13 @@ case class AddIndexExec( * covering exactly those fragments, so the resulting segments have disjoint fragment coverage and * can be committed directly as a single logical index. * + * Unlike the segmented path, coverage here is derived from the fragment ids present in the scanned + * rows rather than from a fragment list fixed at planning time. The read options are pinned even so. + * The scan resolves its own version through the catalog, which is at or after the pinned one, so a + * segment can end up stamped older than the rows it holds and lose its coverage to core's staleness + * pruning; leaving the build unpinned instead lets an executor stamp the current version over rows + * read earlier, and core validates a segment only when its stamp predates the commit. + * * @param addIndexExec The AddIndexExec instance that initiated this job * @param readOptions Configuration options for reading the Lance dataset * @param numFragments Number of fragments in the dataset, used to bound shuffle partitions @@ -335,6 +339,7 @@ class RangeBasedBTreeIndexJob( val indexBuilder = RangeBTreeIndexBuilder( encode(readOptions), + addIndexExec.indexName, columns, zoneSize, nsImpl, @@ -364,6 +369,7 @@ class RangeBasedBTreeIndexJob( * This class is serialized and sent to executors to build the index for a specific range of data. * * @param encodedReadOptions Serialized configuration for Lance dataset access. + * @param indexName Name of the logical index the segment will belong to. * @param columns The names of the columns to be indexed. * @param zoneSize Optional size of zones within the B-tree index. * @param namespaceImpl Optional implementation class for namespace operations, used for credential vending. @@ -374,6 +380,7 @@ class RangeBasedBTreeIndexJob( */ case class RangeBTreeIndexBuilder( encodedReadOptions: String, + indexName: String, columns: List[String], zoneSize: Option[Long], namespaceImpl: Option[String], @@ -438,9 +445,10 @@ case class RangeBTreeIndexBuilder( Data.exportArrayStream(allocator, reader, stream) // Build an uncommitted BTree segment for this fragment group from the - // pre-sorted data. No index name or UUID is set: Lance generates the - // segment UUID, and the fragment ids declare the segment's coverage so - // the per-partition segments stay disjoint. + // pre-sorted data. No UUID is set: Lance generates the segment UUID, and + // the fragment ids declare the segment's coverage so the per-partition + // segments stay disjoint. See ScalarSegmentIndexTask for why the segment + // build names the index and sets replace. val btreeParamsBuilder = BTreeIndexParams.builder() if (zoneSize.isDefined) { btreeParamsBuilder.zoneSize(zoneSize.get) @@ -451,7 +459,8 @@ case class RangeBTreeIndexBuilder( val indexOptions = IndexOptions .builder(columns.asJava, IndexType.BTREE, indexParams) - .replace(false) + .withIndexName(indexName) + .replace(true) .withFragmentIds(fragmentIds.toList.asJava) .withPreprocessedData(stream) .build() @@ -471,9 +480,16 @@ case class RangeBTreeIndexBuilder( * A job implementation for creating scalar segment indexes using logical segment commit. * Fragments are batched into segments, each built in parallel, and committed * as a logical index on the driver. + * + * Shared by CREATE INDEX, which passes every fragment, and REFRESH INDEX, which passes only the + * fragments an existing index does not cover. */ class ScalarSegmentIndexJob( - addIndexExec: AddIndexExec, + sc: SparkContext, + indexName: String, + method: String, + columns: List[String], + argsJson: String, readOptions: LanceSparkReadOptions, fragmentWorkloads: List[FragmentWorkload], numSegments: Option[Int], @@ -483,23 +499,22 @@ class ScalarSegmentIndexJob( initialStorageOpts: Option[Map[String, String]]) { def run(): Seq[Index] = { - val indexType = IndexUtils.scalarSegmentIndexType(addIndexExec.method).getOrElse { + val indexType = IndexUtils.scalarSegmentIndexType(method).getOrElse { throw new UnsupportedOperationException( - s"Unsupported Lance index method: ${addIndexExec.method}") + s"Unsupported Lance index method: $method") } val encodedReadOptions = encode(readOptions) - val columns = addIndexExec.columns.toList - val argsJson = IndexUtils.toJson(addIndexExec.args) val fragmentBatches = IndexUtils.batchFragments( fragmentWorkloads, numSegments, - addIndexExec.session.sparkContext.defaultParallelism) + sc.defaultParallelism) val tasks = fragmentBatches.map { batch => ScalarSegmentIndexTask( encodedReadOptions, + indexName, columns, - addIndexExec.method, + method, argsJson, batch, nsImpl, @@ -509,7 +524,7 @@ class ScalarSegmentIndexJob( }.toSeq IndexUtils.runSegmentTasks( - addIndexExec.session.sparkContext, + sc, tasks, s"${indexType.name()} index build failed. Uncommitted segments are not " + "visible to readers and will not affect query correctness.")(_.execute()) @@ -520,9 +535,18 @@ final private[v2] case class FragmentWorkload(fragmentId: Integer, numRows: Long /** * A task to create a scalar index segment on a batch of fragments. + * + * The segment is named after the logical index it will join, and {@code replace} is set because that + * index may already exist: on the uncommitted build path Lance consults {@code replace} only to + * reject a name that is already taken, which is precisely the normal case for a REFRESH, and for a + * CREATE INDEX that replaces an index of the same name. Nothing is removed here — the driver's + * single {@code commitExistingIndexSegments} transaction decides which existing segments to keep. + * Leaving the name unset instead makes Lance derive its own default (`{column}_idx`) and reject the + * build whenever an index of that name exists on the column. */ case class ScalarSegmentIndexTask( encodedReadOptions: String, + indexName: String, columns: List[String], method: String, argsJson: String, @@ -545,8 +569,9 @@ case class ScalarSegmentIndexTask( val indexOptions = IndexOptions .builder(java.util.Arrays.asList(columns: _*), indexType, params) + .withIndexName(indexName) .withFragmentIds(fragmentIds.asJava) - .replace(false) + .replace(true) .build() val dataset = Utils.openDatasetBuilder(readOptions) @@ -614,11 +639,259 @@ object IndexUtils extends Logging { IndexType.RTREE -> "rtree", IndexType.INVERTED -> "inverted") + // Reverse of `methodToIndexTypes`, for commands that resolve an already-created index and need + // the method name back. INVERTED maps to the canonical "fts" spelling rather than its alias. + private val methodByIndexType: Map[IndexType, String] = Map( + IndexType.BTREE -> "btree", + IndexType.ZONEMAP -> "zonemap", + IndexType.BITMAP -> "bitmap", + IndexType.LABEL_LIST -> "label_list", + IndexType.NGRAM -> "ngram", + IndexType.BLOOM_FILTER -> "bloomfilter", + IndexType.RTREE -> "rtree", + IndexType.INVERTED -> "fts") + def scalarSegmentIndexType(method: String): Option[IndexType] = methodToIndexTypes .get(method.toLowerCase(Locale.ROOT)) .filter(scalarSegmentIndexTypes.contains) + /** The SQL method name for an index type that supports segmented builds, if any. */ + def methodForIndexType(indexType: IndexType): Option[String] = + Option(indexType).filter(scalarSegmentIndexTypes.contains).flatMap(methodByIndexType.get) + + private val SystemIndexNames: Set[String] = Set("__lance_frag_reuse", "__lance_mem_wal") + + /** True for indexes Lance maintains itself, which user commands must not target. */ + def isSystemIndex(indexName: String): Boolean = + indexName != null && SystemIndexNames.exists(_.equalsIgnoreCase(indexName)) + + /** Parses and range-checks the `num_segments` option shared by the segmented build paths. */ + def parseNumSegments(arg: LanceNamedArgument): Int = arg.value match { + case null => + throw new IllegalArgumentException("num_segments must be a positive integer, got: null") + case n: Number => + val asLong = n.longValue() + if (asLong < 1L || asLong > Int.MaxValue) { + throw new IllegalArgumentException( + s"num_segments must be a positive integer that fits in Int, got: $asLong") + } + asLong.toInt + case other => + throw new IllegalArgumentException(s"num_segments must be a positive integer, got: $other") + } + + /** + * Pins `readOptions` to the version `dataset` is open at. + * + * Distributed index builds hand read options to tasks that open the dataset themselves. Pinning + * makes every task observe the fragment set the driver planned over instead of resolving the + * latest version independently. + */ + def pinVersion( + readOptions: LanceSparkReadOptions, + dataset: Dataset): LanceSparkReadOptions = + readOptions.withRef(Utils.pinOpenedRef(dataset, readOptions.getRef)) + + /** + * Fragment ids and live row counts of `dataset`, in manifest order. + * + * Reads the primitive fragment-statistics view rather than [[Dataset#getFragments]], which + * materializes a Java object per fragment and per data file. Both commands enumerate fragments on + * the driver, once to plan and once to check the commit, so on a large table that difference is + * the bulk of planning cost. + */ + def fragmentWorkloads(dataset: Dataset): List[FragmentWorkload] = { + val stats = dataset.getFragmentStatistics + val ids = stats.getIds + val rowCounts = stats.getRowCounts + List.tabulate(ids.length)(index => + FragmentWorkload(Integer.valueOf(ids(index)), rowCounts(index))) + } + + /** Fragment ids live in `dataset`. See [[fragmentWorkloads]] for why this avoids getFragments. */ + def liveFragmentIds(dataset: Dataset): Set[Int] = + dataset.getFragmentStatistics.getIds.toSet + + /** + * Segments of `indexName` that a segment commit will keep, resolved against `dataset` as it is now. + * + * Fails when the index is gone. A refresh resolves its target before the distributed build and + * commits after it, and `commitExistingIndexSegments` against a name Lance no longer knows creates + * that index rather than extending it: a DROP INDEX during the build would otherwise be undone, + * leaving the index back in place with only the coverage this build happened to plan for. + */ + def resolveRetainedSegments( + dataset: Dataset, + indexName: String, + liveFragmentIds: Set[Int]): Seq[Index] = + retainedSegments(dataset.getIndexes.asScala.toSeq, indexName, liveFragmentIds) + + /** Dataset-free form of [[resolveRetainedSegments]]. */ + def retainedSegments( + allSegments: Seq[Index], + indexName: String, + liveFragmentIds: Set[Int]): Seq[Index] = { + val current = allSegments.filter(segment => indexName == segment.name()) + if (current.isEmpty) { + throw new IllegalStateException( + s"Index '$indexName' no longer exists: it was dropped or replaced while the build was " + + "running. Nothing was committed; re-create it with ALTER TABLE ... CREATE INDEX.") + } + // Commit keeps an existing segment only while it still covers a live fragment, and removes the + // rest in the same transaction, so only these have to fit the incoming segments. + current.filter(segment => declaredCoverage(Seq(segment)).exists(liveFragmentIds.contains)) + } + + /** Fragment ids the given segments declare coverage of. */ + def declaredCoverage(segments: Seq[Index]): Set[Int] = + segments.iterator + .flatMap(_.fragments().orElse(Collections.emptyList[Integer]()).asScala) + .map(_.intValue) + .toSet + + /** + * Refuses to publish a segment set that would establish no coverage at all. + * + * Commit intersects each segment's declared coverage with the dataset's live fragments, so a + * fragment retired while the build ran contributes nothing. Lance accepts such a set: an empty + * fragment bitmap is valid metadata, and an existing segment is trivially disjoint from it and so + * survives. Nothing is corrupted, but the transaction would publish segments that index no data and + * report a build that achieved nothing, and this is the last point at which it can still be + * declined. + * + * Partial loss is not refused. A fragment leaves the manifest either because its rows moved + * (compaction, an in-place rewrite) or because they were all deleted; only the first leaves data + * unindexed, and either way the segments covering what remains are correct. Discarding a finished + * distributed build over that would be the wrong response, and it would make a routine concurrent + * DELETE fatal on a long one. What the commit actually established is reported afterwards, by + * [[establishedCoverage]]. + */ + def requireCommittableCoverage( + liveFragmentIds: Set[Int], + segments: Seq[Index], + indexName: String): Unit = { + val declared = declaredCoverage(segments) + if (declared.nonEmpty && declared.intersect(liveFragmentIds).isEmpty) { + throw new IllegalStateException( + s"Index '$indexName' build raced a concurrent operation: every fragment it covers " + + s"(${describeFragmentIds(declared)}) was retired while the build ran, so the segments " + + "would cover nothing. No index change was committed; re-run the command.") + } + } + + /** + * The coverage a commit established, read back from the metadata the commit returned. + * + * A check taken before the commit cannot answer this, for two reasons. It reads the manifest its + * handle was opened at, while the commit lands on whatever version is current by then. And Lance + * prunes an incoming segment's coverage not only for a fragment that is gone but also for one whose + * indexed field was rewritten under the same id, which no comparison of fragment ids can see. The + * returned metadata is post-pruning, so it is the only truthful account of what was indexed. + * + * Only the segments this build produced are counted. Existing segments that were disjoint from them + * survive the commit, and their coverage is not this command's to report. + * + * @return the fragment ids the commit covered with these segments + */ + def establishedCoverage( + builtSegments: Seq[Index], + committedSegments: Seq[Index], + liveFragmentIds: Set[Int], + indexName: String): Set[Int] = { + val builtUuids = builtSegments.map(_.uuid).toSet + val established = + declaredCoverage(committedSegments.filter(segment => builtUuids.contains(segment.uuid))) + .intersect(liveFragmentIds) + val uncovered = declaredCoverage(builtSegments).diff(established) + if (uncovered.nonEmpty) { + logWarning( + s"Index '$indexName' build raced a concurrent operation: fragments " + + s"${describeFragmentIds(uncovered)} are not covered by this commit, because they were " + + "retired or because their indexed field was rewritten while the build ran. The segments " + + "for the remaining fragments are committed; re-run the command to cover them.") + } + established + } + + private def describeFragmentIds(ids: Set[Int]): String = { + val ordered = ids.toSeq.sorted + val shown = ordered.take(10).mkString(", ") + if (ordered.size > 10) s"$shown, ... (${ordered.size} total)" else shown + } + + // Index types whose segments must all describe the same build configuration. Lance builds each + // segment independently and queries most of them independently too, so for those, segments built + // with different options differ in performance only. An inverted (FTS) index is the exception: its + // read path loads one set of index details for the whole logical index and rejects a set whose + // segments disagree, which fails every full-text query on the column. Kept explicit rather than + // inferred; extend it if core grows another type with the same requirement. + private val uniformDetailsIndexTypes: Set[IndexType] = Set(IndexType.INVERTED) + + /** True when segments of this index type must all share one build configuration. */ + def requiresUniformSegmentDetails(indexType: IndexType): Boolean = + indexType != null && uniformDetailsIndexTypes.contains(indexType) + + /** + * Fails before commit when newly built segments describe a different build configuration than the + * segments they are about to join, for an index type that requires them to agree. + * + * Index details are the serialized build parameters, so this compares what was built rather than + * what was asked for, and it runs while the new segments are still uncommitted: a rejection leaves + * the index exactly as it was. Segments without index details predate the field, and an index + * being replaced wholesale has nothing to agree with; both cases defer to Lance core, which + * validates the segment set it is handed. + */ + def requireUniformSegmentDetails( + indexType: IndexType, + indexName: String, + method: String, + retainedSegments: Seq[Index], + builtSegments: Seq[Index]): Unit = { + if (!requiresUniformSegmentDetails(indexType)) { + return + } + val retained = distinctIndexDetails(retainedSegments) + val built = distinctIndexDetails(builtSegments) + if (retained.isEmpty || built.isEmpty) { + return + } + if (retained.size > 1 || built.size > 1 || retained != built) { + throw new IllegalArgumentException( + s"Index '$indexName' uses the $method method, whose segments must all share one " + + "configuration, and the segments this build produced are configured differently from the " + + "ones they would join. Nothing was committed and the index is unchanged. Re-run with the " + + "options the index was created with, or rebuild it in full with " + + "ALTER TABLE ... CREATE INDEX.") + } + } + + private def distinctIndexDetails(segments: Seq[Index]): Set[Seq[Byte]] = + segments.iterator + .flatMap(segment => Option(segment.indexDetails().orElse(null))) + .map(_.toSeq) + .toSet + + /** Namespace and storage context that tasks need to reopen the dataset on an executor. */ + def extractNamespaceInfo( + catalog: TableCatalog, + lanceDataset: LanceDataset, + readOptions: LanceSparkReadOptions): ( + Option[String], + Option[Map[String, String]], + Option[List[String]], + Option[Map[String, String]]) = { + catalog match { + case nsCatalog: BaseLanceNamespaceSparkCatalog => + ( + Option(nsCatalog.getNamespaceImpl), + Option(nsCatalog.getNamespaceProperties).map(_.asScala.toMap), + Option(readOptions.getTableId).map(_.asScala.toList), + Option(lanceDataset.getInitialStorageOptions).map(_.asScala.toMap)) + case _ => (None, None, None, None) + } + } + def resolveIndexField( schema: LanceSchema, indexType: IndexType, @@ -634,7 +907,7 @@ object IndexUtils extends Logging { * Extracts the `train` option from named arguments, defaulting to `true`. * * When `train=false`, index creation registers an empty index without processing any data. - * All existing rows will be unindexed and covered by a subsequent OPTIMIZE INDEX call. + * All existing rows are left unindexed until a subsequent REFRESH INDEX covers them. */ def extractTrain(args: Seq[LanceNamedArgument]): Boolean = args.find(_.name == "train") match { @@ -715,7 +988,7 @@ object IndexUtils extends Logging { } def runSegmentTasks[T <: Serializable: ClassTag]( - sc: org.apache.spark.SparkContext, + sc: SparkContext, tasks: Seq[T], failureMessage: String)(execute: T => String): Seq[Index] = { if (tasks.isEmpty) { @@ -733,6 +1006,23 @@ object IndexUtils extends Logging { } } + /** + * Splits `fragments` into `numSegments` batches, each a contiguous run of fragment ids, chosen so + * that the heaviest batch is as light as possible. + * + * Contiguity is not cosmetic. Lance's compaction planner only groups fragments that are covered by + * the identical set of index segments, so batches whose fragment ids interleave leave every + * adjacent pair of fragments in a different group and make OPTIMIZE a no-op for the whole table. + * Since an index accumulates one segment set per build, and REFRESH INDEX adds more over time, + * interleaved coverage would permanently block compaction on exactly the append-heavy tables this + * command exists for. Balance is not sacrificed to get it: the optimal contiguous partition is + * found exactly, so a workload the previous least-loaded-first assignment balanced perfectly + * still is. + * + * Assignment is deterministic: the same fragments and segment count always produce the same + * batches, whatever order `fragments` arrives in. Every batch holds at least one fragment, so the + * result always has exactly `segmentCount` entries. + */ def batchFragments( fragments: List[FragmentWorkload], numSegments: Option[Int], @@ -754,35 +1044,139 @@ object IndexUtils extends Logging { case None => math.max(1, math.min(fragmentCount, defaultParallelism)) } - final class SegmentBatch(val index: Int) { - val fragmentIds: ArrayBuffer[Integer] = ArrayBuffer.empty - var numRows: Long = 0L + val ordered = fragments.sortBy(_.fragmentId.intValue) + // Prefix sums drive the search. addExact rejects a workload that cannot be summed rather than + // balancing against a wrapped total. + val rowsUpTo = new Array[Long](fragmentCount + 1) + ordered.iterator.zipWithIndex.foreach { case (fragment, index) => + rowsUpTo(index + 1) = Math.addExact(rowsUpTo(index), fragment.numRows) + } + + var offset = 0 + balancedRunLengths(rowsUpTo, segmentCount).map { length => + val batch = ordered.slice(offset, offset + length).map(_.fragmentId) + offset += length + batch + } + } + + /** + * Lengths of exactly `segmentCount` contiguous runs over `rowsUpTo`, minimising the heaviest run. + * + * The smallest row budget a contiguous packing can respect is found by binary search, which is + * exact rather than approximate: for a fixed budget, extending each run as far as it will go uses + * the fewest runs, so the smallest feasible budget is the optimal maximum. The floor of the search + * is the widest single fragment, below which no packing exists. + * + * Packing at that budget can use fewer runs than were asked for, which would cost parallelism, so + * the remainder are split at their own balance points. A split only ever lowers the heaviest run, + * so optimality survives it. + */ + private def balancedRunLengths(rowsUpTo: Array[Long], segmentCount: Int): Seq[Int] = { + val fragmentCount = rowsUpTo.length - 1 + def rowsOf(start: Int, length: Int): Long = rowsUpTo(start + length) - rowsUpTo(start) + + var widestFragment = 0L + var index = 0 + while (index < fragmentCount) { + widestFragment = math.max(widestFragment, rowsOf(index, 1)) + index += 1 + } + + var low = widestFragment + var high = rowsUpTo(fragmentCount) + while (low < high) { + val budget = low + (high - low) / 2 + if (runsWithin(rowsUpTo, budget) <= segmentCount) high = budget else low = budget + 1 + } - def add(fragment: FragmentWorkload): Unit = { - numRows = Math.addExact(numRows, fragment.numRows) - fragmentIds += fragment.fragmentId + // runLengthAt(start) is the length of the run beginning at `start`, and 0 elsewhere. + val runLengthAt = new Array[Int](fragmentCount) + var runCount = 0 + var start = 0 + while (start < fragmentCount) { + var length = 1 + while (start + length < fragmentCount && rowsOf(start, length + 1) <= low) { + length += 1 } + runLengthAt(start) = length + runCount += 1 + start += length } - val segmentOrdering: Ordering[SegmentBatch] = - Ordering - .by[SegmentBatch, (Long, Int, Int)](segment => - (segment.numRows, segment.fragmentIds.size, segment.index)) - .reverse + if (runCount < segmentCount) { + // Heaviest splittable run first, so each extra batch is spent where it helps most. Ties break + // on length then on the earlier run, to keep the result independent of heap internals. + val splittable = PriorityQueue.empty[(Long, Int, Int)]( + Ordering.by[(Long, Int, Int), (Long, Int, Int)] { + case (rows, length, begin) => (rows, length, -begin) + }) + var scan = 0 + while (scan < fragmentCount) { + val length = runLengthAt(scan) + if (length > 1) { + splittable.enqueue((rowsOf(scan, length), length, scan)) + } + scan += length + } + while (runCount < segmentCount && splittable.nonEmpty) { + val (_, length, begin) = splittable.dequeue() + val cut = balancePoint(rowsUpTo, begin, length) + runLengthAt(begin) = cut + runLengthAt(begin + cut) = length - cut + runCount += 1 + if (cut > 1) { + splittable.enqueue((rowsOf(begin, cut), cut, begin)) + } + if (length - cut > 1) { + splittable.enqueue((rowsOf(begin + cut, length - cut), length - cut, begin + cut)) + } + } + } - val segments = PriorityQueue.empty[SegmentBatch](segmentOrdering) - (0 until segmentCount).foreach(index => segments.enqueue(new SegmentBatch(index))) + val runs = ArrayBuffer.empty[Int] + var cursor = 0 + while (cursor < fragmentCount) { + val length = runLengthAt(cursor) + runs += length + cursor += length + } + runs.toList + } - val sortedFragments = fragments.sortBy(fragment => (-fragment.numRows, fragment.fragmentId)) - sortedFragments.foreach { fragment => - val segment = segments.dequeue() - segment.add(fragment) - segments.enqueue(segment) + /** + * Fewest contiguous runs that keep every run's row count within `budget`. + * + * A fragment wider than `budget` still forms a run of its own, so callers must not search below + * the widest fragment or the count would understate what the budget can actually hold. + */ + private def runsWithin(rowsUpTo: Array[Long], budget: Long): Int = { + val fragmentCount = rowsUpTo.length - 1 + var runs = 0 + var start = 0 + while (start < fragmentCount) { + var length = 1 + while (start + length < fragmentCount && + rowsUpTo(start + length + 1) - rowsUpTo(start) <= budget) { + length += 1 + } + runs += 1 + start += length } + runs + } - segments.toSeq - .sortBy(_.index) - .map(segment => segment.fragmentIds.sortBy(_.intValue()).toList) + /** Where to cut a run of `length` fragments so the two halves are as even as possible. */ + private def balancePoint(rowsUpTo: Array[Long], start: Int, length: Int): Int = { + val total = rowsUpTo(start + length) - rowsUpTo(start) + def leftOf(at: Int): Long = rowsUpTo(start + at) - rowsUpTo(start) + def heavierHalf(at: Int): Long = math.max(leftOf(at), total - leftOf(at)) + + var cut = 1 + while (cut < length - 1 && leftOf(cut) < total - leftOf(cut)) { + cut += 1 + } + if (cut > 1 && heavierHalf(cut - 1) < heavierHalf(cut)) cut - 1 else cut } } diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala index f17e0d52c..76b09d519 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala @@ -51,6 +51,9 @@ case class LanceDataSourceV2Strategy(session: SparkSession) extends SparkStrateg case LanceDropIndex(ResolvedIdentifier(catalog, ident), indexName) => LanceDropIndexExec(asTableCatalog(catalog), ident, indexName.toLowerCase) :: Nil + case RefreshIndex(ResolvedIdentifier(catalog, ident), indexName, args) => + RefreshIndexExec(asTableCatalog(catalog), ident, indexName.toLowerCase, args) :: Nil + case LanceCreateBranch(ResolvedIdentifier(catalog, ident), branchName, ref, ifNotExists) => LanceCreateBranchExec(asTableCatalog(catalog), ident, branchName, ref, ifNotExists) :: Nil diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RefreshIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RefreshIndexExec.scala new file mode 100644 index 000000000..11b9e62ba --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/RefreshIndexExec.scala @@ -0,0 +1,273 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.spark.sql.execution.datasources.v2 + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow} +import org.apache.spark.sql.catalyst.plans.logical.{LanceNamedArgument, RefreshIndexOutputType} +import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog} +import org.apache.spark.unsafe.types.UTF8String +import org.lance.Dataset +import org.lance.index.IndexType +import org.lance.spark.{LanceDataset, LanceSparkReadOptions} +import org.lance.spark.utils.{FieldPathUtils, Utils} + +import java.util.Collections + +import scala.collection.JavaConverters._ + +/** + * Physical execution of distributed REFRESH INDEX (ALTER TABLE ... REFRESH INDEX ...). + * + *

Where {@code CREATE INDEX} rebuilds every fragment, this command indexes only the fragments + * the named index does not already cover, then adds the resulting segments to it. Cost is + * proportional to the new data rather than to the table, which makes it the incremental + * counterpart to a full rebuild for tables that are appended to continuously. + * + *

Execution mirrors {@code CREATE INDEX}: the driver resolves the index, diffs its fragment + * coverage against the dataset, and balances the remainder into batches; executors build one + * uncommitted segment per batch; the driver commits them as one logical index. Lance core keeps + * existing segments whose fragments are disjoint from the incoming ones, so prior coverage + * survives the commit. + * + *

Because a deferred index ({@code WITH (train=false)}) covers no fragments, refreshing one + * indexes the whole table through the same distributed path. + * + *

Index build parameters are taken from the {@code WITH} clause, falling back to the index type's + * defaults rather than to the existing index's configuration: Lance records a built index's + * parameters inside the index rather than as table metadata a command can read back. For index types + * whose segments are queried independently that only changes performance, so the caller is asked to + * pass the original options. For a type whose segments must agree — see + * {@code IndexUtils.requiresUniformSegmentDetails} — the mismatch is detected against the built + * segments and rejected before anything is committed. + */ +case class RefreshIndexExec( + catalog: TableCatalog, + ident: Identifier, + indexName: String, + args: Seq[LanceNamedArgument]) extends LeafV2CommandExec { + + override def output: Seq[Attribute] = RefreshIndexOutputType.SCHEMA + + override protected def run(): Seq[InternalRow] = { + val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "RefreshIndex") + val readOptions = lanceDataset.readOptions() + + val numSegments = validateArgs() + + val plan = { + val ds = Utils.openDatasetBuilder(readOptions).build() + try { + planRefresh(ds, readOptions) + } finally { + ds.close() + } + } + + if (plan.unindexedFragments.isEmpty) { + logInfo(s"Index '${plan.resolvedName}' already covers every fragment; nothing to refresh") + return Seq(noOpResult(plan.resolvedName)) + } + + val (nsImpl, nsProps, tableId, initialStorageOpts) = + IndexUtils.extractNamespaceInfo(catalog, lanceDataset, readOptions) + + val segments = new ScalarSegmentIndexJob( + session.sparkContext, + plan.resolvedName, + plan.method, + List(plan.column), + IndexUtils.toJson(args), + plan.buildReadOptions, + plan.unindexedFragments, + numSegments, + nsImpl, + nsProps, + tableId, + initialStorageOpts).run() + + // Lance core's commitExistingIndexSegments keeps existing segments that are disjoint from the + // incoming fragments and removes the ones they supersede, all in one CreateIndex transaction. + // + // The index is re-resolved here rather than carried over from planning: the whole distributed + // build sits between the two, so the state the commit has to fit is the state now, not the state + // the plan saw. + val dataset = Utils.openDatasetBuilder(readOptions).build() + val fragmentsIndexed = + try { + val liveFragmentIds = IndexUtils.liveFragmentIds(dataset) + val retainedSegments = + IndexUtils.resolveRetainedSegments(dataset, plan.resolvedName, liveFragmentIds) + IndexUtils.requireUniformSegmentDetails( + plan.indexType, + plan.resolvedName, + plan.method, + retainedSegments, + segments) + IndexUtils.requireCommittableCoverage(liveFragmentIds, segments, plan.resolvedName) + val committed = + dataset.commitExistingIndexSegments( + plan.resolvedName, + plan.column, + segments.toList.asJava) + // The commit advances this handle, so the returned metadata and the fragment list below both + // describe the committed state rather than the one the checks above validated. + IndexUtils + .establishedCoverage( + segments, + committed.asScala.toSeq, + IndexUtils.liveFragmentIds(dataset), + plan.resolvedName) + .size + } finally { + dataset.close() + } + + Seq(new GenericInternalRow(Array[Any]( + fragmentsIndexed.toLong, + segments.size.toLong, + UTF8String.fromString(plan.resolvedName)))) + } + + private def noOpResult(resolvedName: String): InternalRow = + new GenericInternalRow(Array[Any](0L, 0L, UTF8String.fromString(resolvedName))) + + /** + * Validates the WITH clause and returns the requested segment count. + * + * Options that only make sense for a full build are rejected rather than ignored, so a + * misapplied option fails instead of silently changing nothing. + */ + private def validateArgs(): Option[Int] = { + args.find(_.name == "train").foreach { _ => + throw new IllegalArgumentException( + "train is not supported for REFRESH INDEX: refreshing an index exists to populate it. " + + "Use CREATE INDEX WITH (train = false) to register an index without building it.") + } + Seq("build_mode", "rows_per_range").foreach { option => + args.find(_.name == option).foreach { _ => + throw new IllegalArgumentException( + s"$option is not supported for REFRESH INDEX: range mode redistributes and sorts the " + + "whole table, which an incremental refresh does not do. Use CREATE INDEX to rebuild " + + "with build_mode = 'range'.") + } + } + args.find(_.name == "num_segments").map(IndexUtils.parseNumSegments) + } + + private def planRefresh(ds: Dataset, readOptions: LanceSparkReadOptions): RefreshPlan = { + if (IndexUtils.isSystemIndex(indexName)) { + throw new IllegalArgumentException( + s"'$indexName' is a Lance-maintained system index and cannot be refreshed") + } + + // The parser lowercases the requested name, so matching has to ignore case to reach an index + // created elsewhere with mixed case. Lance keys indexes by exact name, so that match can span + // two distinct indexes; unioning their coverage would compute the wrong unindexed set and + // commit it under one of the two names. Refuse instead of guessing. + val matched = ds.getIndexes.asScala.toSeq + .filter(idx => indexName.equalsIgnoreCase(idx.name())) + .groupBy(_.name()) + if (matched.isEmpty) { + throw new IllegalArgumentException( + s"Index '$indexName' does not exist on table ${ident.toString}. " + + "Create it with ALTER TABLE ... CREATE INDEX first.") + } + if (matched.size > 1) { + throw new IllegalArgumentException( + s"'$indexName' matches ${matched.size} indexes differing only in case " + + s"(${matched.keys.toSeq.sorted.mkString(", ")}). Drop or rebuild them so one remains.") + } + val segments = matched.head._2 + + val indexType = segments.head.indexType() + // An index type this command cannot rebuild is also one CREATE INDEX cannot build, since both go + // through the same method mapping: a vector index, for instance, only exists because something + // outside Spark SQL created it. Pointing such a user at CREATE INDEX would be a dead end. + val method = Option(indexType).flatMap(IndexUtils.methodForIndexType).getOrElse { + val described = Option(indexType).map(_.name()).getOrElse("unknown") + throw new UnsupportedOperationException( + s"Spark SQL cannot build index type $described, so '$indexName' cannot be refreshed here. " + + "Maintain it through the Lance SDK, which is also where it was created.") + } + + val fieldIds = segments.head.fields() + if (fieldIds == null || fieldIds.isEmpty) { + throw new IllegalStateException( + s"Index '$indexName' declares no indexed field; rebuild it with CREATE INDEX") + } + // Every method this command supports is keyed on exactly one column, so more fields than that + // means the index carries extra columns or spans several. A rebuilt segment would declare only + // the keyed column, and one logical index needs one field declaration across its segments, so + // committing that mix would produce metadata `describe_indices` rejects. Lance core refuses to + // optimize a covering index for the same reason. + if (fieldIds.size() != 1) { + throw new UnsupportedOperationException( + s"REFRESH INDEX does not support index '$indexName': it declares ${fieldIds.size()} " + + "fields, and a refreshed segment would only cover the column it is keyed on. Rebuild it " + + "with ALTER TABLE ... CREATE INDEX instead.") + } + // Guaranteed by the check above to be the single field the index is keyed on. The path is + // resolved from the field id rather than remembered, so a renamed column still resolves; a + // field id absent from the current schema resolves to nothing and is reported as such. + val fieldId = fieldIds.get(0) + val column = Option(FieldPathUtils.pathByFieldId(ds.getLanceSchema, fieldId)).getOrElse { + throw new IllegalStateException( + s"Index '$indexName' is keyed on field id $fieldId, which is not in the table's current " + + "schema. Drop the index with ALTER TABLE ... DROP INDEX.") + } + + // A segment with no fragment bitmap predates coverage tracking, and one with an empty bitmap + // is a deferred index. Both read as covering nothing here, which plans a full build. Lance + // core rejects a partial-coverage commit against a segment it cannot place, so it stays the + // authority on whether that build is a legal replacement. + val covered = segments + .flatMap(_.fragments().orElse(Collections.emptyList[Integer]()).asScala) + .map(_.intValue) + .toSet + + val unindexed = IndexUtils + .fragmentWorkloads(ds) + .filterNot(fragment => covered.contains(fragment.fragmentId.intValue)) + + // Commit under the name the manifest actually stores. The parser lowercases the requested name, + // so an index created elsewhere with mixed case is matched case-insensitively above; committing + // under the lowercased spelling would fork a second logical index with overlapping coverage. + RefreshPlan( + segments.head.name(), + indexType, + method, + column, + unindexed, + IndexUtils.pinVersion(readOptions, ds)) + } +} + +/** + * The driver-side plan for one refresh. + * + * @param resolvedName index name as stored in the manifest, which the commit must reuse + * @param indexType type of the existing index, which the rebuilt segments must match + * @param method SQL index method resolved from the existing index's type + * @param column canonical path of the column the index is keyed on + * @param unindexedFragments fragments the index does not cover, with their row counts + * @param buildReadOptions read options pinned to the version the plan was computed against + */ +final private[v2] case class RefreshPlan( + resolvedName: String, + indexType: IndexType, + method: String, + column: String, + unindexedFragments: List[FragmentWorkload], + buildReadOptions: LanceSparkReadOptions) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ShowIndexesExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ShowIndexesExec.scala index 997778458..08cf55ac3 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ShowIndexesExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ShowIndexesExec.scala @@ -24,13 +24,6 @@ import org.lance.spark.utils.{FieldPathUtils, Utils} import scala.collection.JavaConverters._ -object ShowIndexesExec { - private val SystemIndexNames = Set("__lance_frag_reuse", "__lance_mem_wal") - - private def isSystemIndex(indexName: String): Boolean = - indexName != null && SystemIndexNames.exists(_.equalsIgnoreCase(indexName)) -} - /** * Physical execution of SHOW INDEXES for Lance datasets. * @@ -53,15 +46,17 @@ case class ShowIndexesExec( val dataset = Utils.openDatasetBuilder(readOptions).build() try { + // Group by logical index: one row per name, with every physical segment kept so segment-level + // metadata can be aggregated. val indexes = dataset.getIndexes.asScala.toSeq - .filterNot(idx => ShowIndexesExec.isSystemIndex(idx.name())) + .filterNot(idx => IndexUtils.isSystemIndex(idx.name())) .groupBy(_.name()) .toSeq .sortBy(_._1) - .map(_._2.head) val lanceSchema = dataset.getLanceSchema() - indexes.map { idx => + indexes.map { case (_, indexSegments) => + val idx = indexSegments.head val fieldIds = idx.fields() val fieldNamesArray = if (fieldIds == null) { @@ -98,6 +93,40 @@ case class ShowIndexesExec( val numUnindexedFragments = getLong("num_unindexed_fragments") val numUnindexedRows = getLong("num_unindexed_rows") + // Share of rows the index covers, truncated to two decimals. Truncating rather than + // rounding keeps the value from ever overstating coverage: a table one row short of being + // fully indexed reads as 99.99, not as 100. Null rather than 100 for an empty table, so + // "no rows" is not reported as fully indexed either. + val indexedPercent: java.lang.Double = + if (numIndexedRows == null || numUnindexedRows == null) { + null + } else { + val total = numIndexedRows.longValue() + numUnindexedRows.longValue() + if (total <= 0L) { + null + } else { + val percent = 100.0 * numIndexedRows.longValue() / total + java.lang.Double.valueOf(math.floor(percent * 100.0) / 100.0) + } + } + + // Physical segments backing this logical index. Older cores report only `num_indices`. + val numSegments = { + val reported = getLong("num_segments") + if (reported != null) reported else getLong("num_indices") + } + + // Total across segments, or null when any segment predates index file size tracking: a + // partial sum would understate the index rather than admit it is unknown. + val sizeBytes: java.lang.Long = { + val perSegment = indexSegments.map(segment => segment.getSizeBytes) + if (perSegment.exists(!_.isPresent)) { + null + } else { + java.lang.Long.valueOf(perSegment.map(_.get.longValue()).sum) + } + } + new GenericInternalRow(Array[Any]( UTF8String.fromString(name), fieldNamesArray, @@ -105,7 +134,10 @@ case class ShowIndexesExec( numIndexedFragments, numIndexedRows, numUnindexedFragments, - numUnindexedRows)) + numUnindexedRows, + indexedPercent, + numSegments, + sizeBytes)) } } finally { dataset.close() diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java index 711292373..b51f55986 100755 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java @@ -17,8 +17,11 @@ import org.lance.index.Index; import org.lance.index.IndexCriteria; import org.lance.index.IndexDescription; +import org.lance.index.IndexOptions; +import org.lance.index.IndexParams; import org.lance.index.IndexType; import org.lance.index.OptimizeOptions; +import org.lance.index.scalar.ScalarIndexParams; import org.lance.ipc.FullTextQuery; import org.lance.ipc.LanceScanner; import org.lance.ipc.ScanOptions; @@ -105,6 +108,75 @@ public void tearDown() throws IOException { } } + /** + * Pins the two Lance behaviours the coverage report rests on: a segment commit returns the + * metadata of the index as committed, and the handle it was made on advances to the manifest it + * wrote. + * + *

Together they are what make the reported count truthful. A check taken before the commit + * reads the manifest its handle was opened at, so a fragment retired in between is still counted; + * intersecting the returned metadata with the fragments live after the commit is what + * excludes it. This asserts the contract rather than the connector code consuming it, because a + * change on either side would silently make the count overstate again. + */ + @Test + public void testSegmentCommitReportsCoverageAsCommitted() { + spark.sql(String.format("create table %s (id int) using lance", fullTable)); + spark.sql(String.format("insert into %s values (0), (1), (2)", fullTable)); + spark.sql(String.format("insert into %s values (3), (4), (5)", fullTable)); + + try (org.lance.Dataset committer = + Utils.openDatasetBuilder(LanceSparkReadOptions.from(tableDir)).build()) { + List fragments = committer.getFragments(); + int coveredFragmentId = fragments.get(fragments.size() - 1).getId(); + + IndexParams indexParams = + IndexParams.builder() + .setScalarIndexParams(ScalarIndexParams.create("zonemap", "{}")) + .build(); + Index built = + committer.createIndex( + IndexOptions.builder(Collections.singletonList("id"), IndexType.ZONEMAP, indexParams) + .withIndexName("idx_committed_coverage") + .replace(true) + .withFragmentIds(Collections.singletonList(coveredFragmentId)) + .build()); + Assertions.assertEquals( + Collections.singletonList(coveredFragmentId), + built.fragments().orElse(Collections.emptyList()), + "the uncommitted segment should declare the fragment it was built for"); + + // Retire every fragment from another handle, leaving the committer on a stale manifest. + spark.sql(String.format("delete from %s where id >= 0", fullTable)); + + List committed = + committer.commitExistingIndexSegments( + "idx_committed_coverage", "id", Collections.singletonList(built)); + + Set ours = Collections.singleton(built.uuid()); + Assertions.assertTrue( + committed.stream().anyMatch(index -> ours.contains(index.uuid())), + "the commit must return the metadata of the segments it was handed"); + + Set liveAfter = + committer.getFragments().stream().map(Fragment::getId).collect(Collectors.toSet()); + Assertions.assertFalse( + liveAfter.contains(coveredFragmentId), + "the committing handle must advance to the manifest the commit wrote"); + + Set established = + committed.stream() + .filter(index -> ours.contains(index.uuid())) + .flatMap(index -> index.fragments().orElse(Collections.emptyList()).stream()) + .filter(liveAfter::contains) + .collect(Collectors.toSet()); + Assertions.assertEquals( + Collections.emptySet(), + established, + "coverage read from the committed state must not count a fragment retired in between"); + } + } + private void prepareDataset() { spark.sql(String.format("create table %s (id int, text string) using lance;", fullTable)); // First insert to create initial fragments diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseRefreshIndexTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseRefreshIndexTest.java new file mode 100644 index 000000000..10de4711b --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseRefreshIndexTest.java @@ -0,0 +1,766 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.spark.update; + +import org.lance.index.Index; +import org.lance.index.IndexOptions; +import org.lance.index.IndexParams; +import org.lance.index.IndexType; +import org.lance.index.scalar.ScalarIndexParams; +import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.utils.Utils; + +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.io.IOException; +import java.nio.file.FileSystems; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.UUID; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +/** Base test for distributed REFRESH INDEX. */ +public abstract class BaseRefreshIndexTest { + + private static final StructType SCHEMA = + new StructType( + new StructField[] { + DataTypes.createStructField("id", DataTypes.IntegerType, false), + DataTypes.createStructField("text", DataTypes.StringType, false) + }); + + protected String catalogName = "lance_test"; + protected String tableName; + protected String fullTable; + + protected SparkSession spark; + + @TempDir Path tempDir; + protected String tableDir; + + @BeforeEach + public void setup() throws IOException { + Path rootPath = tempDir.resolve(UUID.randomUUID().toString()); + Files.createDirectories(rootPath); + String testRoot = rootPath.toString(); + spark = + SparkSession.builder() + .appName("lance-refresh-index-test") + .master("local[3]") + .config("spark.default.parallelism", "10") + .config( + "spark.sql.catalog." + catalogName, "org.lance.spark.LanceNamespaceSparkCatalog") + .config( + "spark.sql.extensions", "org.lance.spark.extensions.LanceSparkSessionExtensions") + .config("spark.sql.catalog." + catalogName + ".impl", "dir") + .config("spark.sql.catalog." + catalogName + ".root", testRoot) + .config("spark.sql.catalog." + catalogName + ".single_level_ns", "true") + // The FTS predicate functions are resolved through the current catalog, so the + // full-text tests below need it to be the Lance one. + .config("spark.sql.defaultCatalog", catalogName) + .getOrCreate(); + this.tableName = "refresh_index_test_" + UUID.randomUUID().toString().replace("-", ""); + this.fullTable = this.catalogName + ".default." + this.tableName; + this.tableDir = + FileSystems.getDefault().getPath(testRoot, this.tableName + ".lance").toString(); + } + + @AfterEach + public void tearDown() throws IOException { + if (spark != null) { + // Spark 4 declares SparkSession.close() as throwing IOException; Spark 3 does not. + spark.close(); + } + } + + /** Refresh indexes the fragments an index does not cover, and only those. */ + @Test + public void testRefreshIndexesOnlyUncoveredFragments() { + createTable(); + appendFragment(0, 10); + appendFragment(10, 20); + createZonemapIndex(); + + Assertions.assertEquals(2, coverage().size(), "Create should cover both initial fragments"); + Set segmentsBefore = segmentUuids(); + + appendFragment(20, 30); + appendFragment(30, 40); + + Row result = refresh(""); + Assertions.assertEquals(2L, result.getLong(0), "Only the two new fragments should be indexed"); + Assertions.assertTrue(result.getLong(1) >= 1L, "At least one segment should be added"); + Assertions.assertEquals("idx_id", result.getString(2)); + + Assertions.assertEquals(4, coverage().size(), "All four fragments should now be covered"); + Assertions.assertTrue( + segmentUuids().containsAll(segmentsBefore), + "Refresh must preserve the segments covering already-indexed fragments"); + } + + /** A refresh with nothing to do must not commit a new version. */ + @Test + public void testRefreshIsNoOpWhenFullyCovered() { + createTable(); + appendFragment(0, 10); + appendFragment(10, 20); + createZonemapIndex(); + + long versionBefore = datasetVersion(); + Row result = refresh(""); + + Assertions.assertEquals(0L, result.getLong(0)); + Assertions.assertEquals(0L, result.getLong(1)); + Assertions.assertEquals( + versionBefore, datasetVersion(), "A no-op refresh must not commit a version"); + } + + /** A deferred index covers nothing, so refreshing one builds it over the whole table. */ + @Test + public void testRefreshPopulatesDeferredIndex() { + createTable(); + appendFragment(0, 10); + appendFragment(10, 20); + spark.sql( + String.format( + "alter table %s create index idx_id using zonemap (id) with (train = false)", + fullTable)); + + Assertions.assertTrue(coverage().isEmpty(), "A deferred index should cover no fragments"); + + // Statistics must be readable for an index that has no files yet. + Row deferred = + spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals(0L, deferred.getLong(4), "A deferred index covers no rows"); + Assertions.assertEquals( + 0.0d, deferred.getDouble(7), 1e-9, "A deferred index should report 0 percent coverage"); + + Row result = refresh(""); + Assertions.assertEquals(2L, result.getLong(0)); + Assertions.assertEquals(2, coverage().size()); + + Row stats = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals(0L, stats.getLong(5), "No fragment should remain unindexed"); + Assertions.assertEquals(0L, stats.getLong(6), "No row should remain unindexed"); + Assertions.assertEquals(100.0d, stats.getDouble(7), 1e-9, "Index should report full coverage"); + } + + /** Data stays queryable and complete through a refresh. */ + @Test + public void testRefreshedIndexAnswersQueries() { + createTable(); + appendFragment(0, 10); + appendFragment(10, 20); + createZonemapIndex(); + appendFragment(20, 30); + + refresh(""); + + Assertions.assertEquals(30L, spark.table(fullTable).count(), "No rows should be lost"); + Assertions.assertEquals( + 1L, + spark.sql(String.format("select * from %s where id = 25", fullTable)).count(), + "Point lookup in a refreshed fragment should find its row"); + Assertions.assertEquals( + 10L, + spark.sql(String.format("select * from %s where id >= 20 and id < 30", fullTable)).count(), + "Range scan over a refreshed fragment should return every row"); + Assertions.assertEquals( + 0L, + spark.sql(String.format("select * from %s where id = 99", fullTable)).count(), + "Absent value should match nothing"); + } + + /** + * Compaction retires the fragments an index covers and creates new ones the index does not, so a + * refresh is what restores coverage afterwards. + * + *

Asserts only the post-refresh state: query results are correct and every fragment is + * covered. It deliberately does not pin down the intermediate post-compaction state, which is a + * separate concern from this command. + */ + @Test + public void testRefreshRestoresCoverageAfterCompaction() { + createTable(); + // One fragment already at the compaction target so it is not a candidate, plus smaller ones + // that are: compaction then retires part of the index's coverage and leaves the rest. + appendFragment(0, 100); + appendFragment(100, 110); + appendFragment(110, 120); + appendFragment(120, 130); + spark.sql( + String.format( + "alter table %s create index idx_id using zonemap (id) with (num_segments = 1)", + fullTable)); + + spark.sql(String.format("optimize %s with (target_rows_per_fragment = 50)", fullTable)); + + refresh(""); + + Row stats = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals( + 0L, + stats.getLong(5), + "Every fragment should be covered after refreshing a compacted table"); + Assertions.assertEquals( + 100.0d, stats.getDouble(7), 1e-9, "Coverage should be complete after refresh"); + + Assertions.assertEquals(130L, spark.table(fullTable).count(), "No rows should be lost"); + Assertions.assertEquals( + 1L, + spark.sql(String.format("select * from %s where id = 105", fullTable)).count(), + "A row moved by compaction must be found through the refreshed index"); + Assertions.assertEquals( + 30L, + spark.sql(String.format("select * from %s where id >= 95 and id < 125", fullTable)).count(), + "A range spanning compacted and untouched fragments must return every row exactly once"); + } + + /** + * A table one row short of full coverage must not report 100 percent, or an operator polling SHOW + * INDEXES would conclude the index is complete while a fragment is still unindexed. + */ + @Test + public void testIndexedPercentNeverOverstatesCoverage() { + createTable(); + appendFragment(0, 20000); + createZonemapIndex(); + appendFragment(20000, 20001); + + Row stats = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals(1L, stats.getLong(6), "Exactly one row should be unindexed"); + Assertions.assertTrue( + stats.getDouble(7) < 100.0d, + "A partially indexed table must not report 100 percent, got " + stats.getDouble(7)); + + refresh(""); + + Row after = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals( + 100.0d, after.getDouble(7), 1e-9, "Full coverage should report exactly 100 percent"); + } + + /** num_segments bounds the number of parallel build tasks, and so the segments committed. */ + @Test + public void testRefreshRespectsNumSegments() { + createTable(); + appendFragment(0, 10); + createZonemapIndex(); + appendFragment(10, 20); + appendFragment(20, 30); + appendFragment(30, 40); + appendFragment(40, 50); + + Row result = refresh("with (num_segments = 2)"); + Assertions.assertEquals(4L, result.getLong(0)); + Assertions.assertEquals(2L, result.getLong(1), "Four fragments should build as two segments"); + Assertions.assertEquals(5, coverage().size()); + } + + /** + * The column to rebuild is resolved from the index's field id, not from remembered text, so a + * nested field must round-trip back to a path the index builder accepts. + */ + @Test + public void testRefreshIndexOnNestedField() { + spark.sql( + String.format( + "create table %s (id int, payload struct) using lance;", fullTable)); + spark.sql(String.format("insert into %s values (1, named_struct('value', 10))", fullTable)); + spark.sql( + String.format( + "alter table %s create index idx_nested using btree (payload.value)", fullTable)); + spark.sql(String.format("insert into %s values (2, named_struct('value', 20))", fullTable)); + + Row result = + spark + .sql(String.format("alter table %s refresh index idx_nested", fullTable)) + .collectAsList() + .get(0); + + Assertions.assertEquals(1L, result.getLong(0), "The appended fragment should be indexed"); + Assertions.assertEquals("idx_nested", result.getString(2)); + Assertions.assertEquals( + 1L, + spark.sql(String.format("select * from %s where payload.value = 20", fullTable)).count(), + "The refreshed nested index must answer a lookup on the new fragment"); + } + + /** Index options are forwarded to the segment build. */ + @Test + public void testRefreshForwardsIndexOptions() { + createTable(); + appendFragment(0, 10); + createZonemapIndex(); + appendFragment(10, 20); + + Row result = refresh("with (rows_per_zone = 2048)"); + Assertions.assertEquals(1L, result.getLong(0)); + Assertions.assertEquals(2, coverage().size()); + } + + @Test + public void testRefreshUnknownIndexFails() { + createTable(); + appendFragment(0, 10); + + Exception error = + Assertions.assertThrows( + Exception.class, + () -> spark.sql(String.format("alter table %s refresh index missing_idx", fullTable))); + Assertions.assertTrue( + causeChain(error).contains("does not exist"), + "Expected an actionable message, got: " + causeChain(error)); + } + + @Test + public void testRefreshSystemIndexFails() { + createTable(); + appendFragment(0, 10); + + Exception error = + Assertions.assertThrows( + Exception.class, + () -> + spark.sql( + String.format("alter table %s refresh index __lance_frag_reuse", fullTable))); + Assertions.assertTrue( + causeChain(error).contains("system index"), + "Expected a system-index rejection, got: " + causeChain(error)); + } + + /** Options that only apply to a full rebuild are rejected rather than silently ignored. */ + @ParameterizedTest + @ValueSource( + strings = { + "train = false", + "build_mode = 'range'", + "rows_per_range = 100", + // Option names are normalized before validation, so an upper-case spelling cannot slip past + // the rejection and take the silent default instead. + "TRAIN = false", + "BUILD_MODE = 'range'", + "ROWS_PER_RANGE = 100" + }) + public void testRefreshRejectsFullBuildOnlyOptions(String option) { + createTable(); + appendFragment(0, 10); + createZonemapIndex(); + appendFragment(10, 20); + + Exception error = + Assertions.assertThrows(Exception.class, () -> refresh(String.format("with (%s)", option))); + Assertions.assertTrue( + causeChain(error).contains("not supported for REFRESH INDEX"), + "Expected a rejection naming the option, got: " + causeChain(error)); + } + + /** + * An inverted index loads one configuration for the whole logical index and rejects a segment set + * that disagrees, so a refresh that repeats the original options must keep queries working. + */ + @Test + public void testRefreshFtsIndexWithMatchingOptions() { + createTable(); + appendFragment(0, 10); + createFtsIndex(); + Assertions.assertEquals(10L, ftsMatchCount(), "The index should answer before the refresh"); + + appendFragment(10, 20); + Row result = + spark + .sql( + String.format( + "alter table %s refresh index idx_text with (%s)", fullTable, FTS_OPTIONS)) + .collectAsList() + .get(0); + + Assertions.assertEquals(1L, result.getLong(0)); + Assertions.assertEquals(20L, ftsMatchCount(), "Every row should match after the refresh"); + Assertions.assertEquals( + 1L, ftsPhraseCount("text 15"), "Positions must survive so phrase queries still work"); + } + + /** + * Lance does not report a built index's options, so a refresh that omits them builds segments + * configured differently. For an inverted index that combination fails every full-text query, so + * the command must refuse it rather than commit an index that looks healthy and answers nothing. + */ + @Test + public void testRefreshFtsIndexRejectsMismatchedOptions() { + createTable(); + appendFragment(0, 10); + createFtsIndex(); + appendFragment(10, 20); + + Exception error = + Assertions.assertThrows( + Exception.class, + () -> + spark + .sql(String.format("alter table %s refresh index idx_text", fullTable)) + .collectAsList()); + Assertions.assertTrue( + causeChain(error).contains("share one configuration"), + "Expected a configuration-mismatch rejection, got: " + causeChain(error)); + + Assertions.assertEquals( + 1, segments("idx_text").size(), "A rejected refresh must not commit segments"); + Assertions.assertEquals( + 20L, ftsMatchCount(), "Full-text queries must keep working and stay complete"); + Assertions.assertEquals(1L, ftsPhraseCount("text 15"), "Phrase queries must keep working"); + } + + /** An fts index built with defaults on both sides needs no options to refresh. */ + @Test + public void testRefreshFtsIndexWithDefaultOptions() { + createTable(); + appendFragment(0, 10); + spark.sql(String.format("alter table %s create index idx_text using fts (text)", fullTable)); + appendFragment(10, 20); + + Row result = + spark + .sql(String.format("alter table %s refresh index idx_text", fullTable)) + .collectAsList() + .get(0); + + Assertions.assertEquals(1L, result.getLong(0)); + Assertions.assertEquals(20L, ftsMatchCount(), "Every row should match after the refresh"); + } + + /** + * {@code _idx} is the name Lance derives for an unnamed index, so a segment build that + * lets it derive its own name collides with the very index being refreshed. + */ + @Test + public void testRefreshIndexNamedAfterItsColumn() { + createTable(); + appendFragment(0, 10); + spark.sql(String.format("alter table %s create index id_idx using zonemap (id)", fullTable)); + appendFragment(10, 20); + + Row result = + spark + .sql(String.format("alter table %s refresh index id_idx", fullTable)) + .collectAsList() + .get(0); + + Assertions.assertEquals(1L, result.getLong(0)); + Assertions.assertEquals("id_idx", result.getString(2)); + Assertions.assertEquals( + 1L, + spark.sql(String.format("select * from %s where id = 15", fullTable)).count(), + "The refreshed index must answer a lookup on the new fragment"); + } + + /** Option names are matched case-insensitively, so an upper-case spelling still takes effect. */ + @Test + public void testRefreshHonorsUpperCaseOptionNames() { + createTable(); + appendFragment(0, 10); + createZonemapIndex(); + appendFragment(10, 20); + appendFragment(20, 30); + appendFragment(30, 40); + appendFragment(40, 50); + + Row result = refresh("WITH (NUM_SEGMENTS = 2)"); + Assertions.assertEquals(4L, result.getLong(0)); + Assertions.assertEquals( + 2L, result.getLong(1), "An upper-case num_segments must bound the segment count"); + } + + /** SHOW INDEXES aggregates across every physical segment of one logical index. */ + @Test + public void testShowIndexesAggregatesAcrossSegments() { + createTable(); + appendFragment(0, 10); + spark.sql( + String.format( + "alter table %s create index idx_id using zonemap (id) with (num_segments = 1)", + fullTable)); + appendFragment(10, 20); + appendFragment(20, 30); + + Row added = refresh("with (num_segments = 2)"); + Assertions.assertEquals(2L, added.getLong(1)); + + List committed = segments("idx_id"); + Assertions.assertEquals(3, committed.size(), "One created segment plus two refreshed ones"); + + Row stats = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals(3L, stats.getLong(8), "num_segments must count every segment"); + long expectedSize = 0L; + for (Index segment : committed) { + expectedSize += segment.getSizeBytes().orElse(0L); + } + Assertions.assertEquals( + expectedSize, stats.getLong(9), "size_bytes must sum every segment's files"); + } + + /** + * Compaction rewrites the manifest bitmap of an index whose matches are row ids, so its coverage + * survives and a refresh has nothing to do. Only address-domain indexes (zonemap, bloomfilter) + * need one afterwards. + */ + @Test + public void testRefreshAfterCompactionIsNoOpForRowIdIndex() { + createTable(); + appendFragment(0, 100); + appendFragment(100, 110); + appendFragment(110, 120); + appendFragment(120, 130); + spark.sql( + String.format( + "alter table %s create index idx_btree using btree (id) with (num_segments = 1)", + fullTable)); + + spark.sql(String.format("optimize %s with (target_rows_per_fragment = 50)", fullTable)); + + Row stats = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + Assertions.assertEquals( + 0L, stats.getLong(5), "Compaction must not strand fragments for a row-id index"); + + Row result = + spark + .sql(String.format("alter table %s refresh index idx_btree", fullTable)) + .collectAsList() + .get(0); + Assertions.assertEquals(0L, result.getLong(0), "Nothing should be left to index"); + Assertions.assertEquals(0L, result.getLong(1)); + Assertions.assertEquals( + 1L, + spark.sql(String.format("select * from %s where id = 105", fullTable)).count(), + "A row moved by compaction must still be found"); + } + + /** + * Lance only compacts fragments covered by the identical set of index segments, so a segment set + * whose coverage interleaves fragment ids leaves OPTIMIZE with no two fragments it can group. + * Batching each segment over a contiguous run keeps compaction possible within that run. + * + *

Uses {@code btree} so compaction remaps the index's coverage and the assertions can stay on + * query results; an address-domain index would need a refresh first. + */ + @Test + public void testOptimizeStillCompactsWithMultipleIndexSegments() { + createTable(); + appendFragment(0, 10); + appendFragment(10, 20); + appendFragment(20, 30); + appendFragment(30, 40); + spark.sql( + String.format( + "alter table %s create index idx_btree using btree (id) with (num_segments = 2)", + fullTable)); + + Assertions.assertEquals(4, liveFragmentCount()); + spark.sql(String.format("optimize %s with (target_rows_per_fragment = 1000)", fullTable)); + + Assertions.assertTrue( + liveFragmentCount() < 4, + "OPTIMIZE must still coalesce fragments a multi-segment index covers, got " + + liveFragmentCount() + + " fragments"); + Assertions.assertEquals(40L, spark.table(fullTable).count(), "No rows should be lost"); + Assertions.assertEquals( + 1L, + spark.sql(String.format("select * from %s where id = 25", fullTable)).count(), + "Rows must stay reachable through the compacted fragments"); + } + + /** + * A segment commit must not resurrect an index that was dropped after the commit handle was + * opened. + * + *

{@code RefreshIndexExec} re-resolves the index on its commit handle, which closes the wide + * window — the whole distributed build — but the check and the commit are still two operations. + * Lance's rebase treats a concurrent {@code CreateIndex} as conflicting only when the *other* + * transaction also carries {@code new_indices}, and a drop carries only {@code removed_indices}, + * so a drop landing inside that window is rebased over rather than rejected. The resurrected + * index then holds only the segments the refresh built, i.e. partial coverage. + * + *

This exercises the core primitive directly, because that is where the guarantee has to come + * from: any connector-side check has the same race, and {@code commitExistingIndexSegments} takes + * no expected predecessor. + */ + @Test + @Disabled( + "blocked on Lance core: rebase does not treat a concurrent same-name index drop as a" + + " conflict. Fixed upstream by lance-format/lance#6806; enable once that is released and" + + " pinned.") + public void testStaleSegmentCommitDoesNotResurrectDroppedIndex() { + createTable(); + appendFragment(0, 10); + createZonemapIndex(); + appendFragment(10, 20); + + LanceSparkReadOptions readOptions = LanceSparkReadOptions.from(tableDir); + try (org.lance.Dataset commitDataset = Utils.openDatasetBuilder(readOptions).build()) { + int appendedFragmentId = + commitDataset.getFragments().get(commitDataset.getFragments().size() - 1).getId(); + IndexParams indexParams = + IndexParams.builder() + .setScalarIndexParams(ScalarIndexParams.create("zonemap", "{}")) + .build(); + Index segment = + commitDataset.createIndex( + IndexOptions.builder(Collections.singletonList("id"), IndexType.ZONEMAP, indexParams) + .withIndexName("idx_id") + .replace(true) + .withFragmentIds(Collections.singletonList(appendedFragmentId)) + .build()); + + // The drop commits on its own handle, leaving the one above holding a stale manifest. + try (org.lance.Dataset dropDataset = Utils.openDatasetBuilder(readOptions).build()) { + dropDataset.dropIndex("idx_id"); + } + Assertions.assertTrue(segments().isEmpty(), "The drop should have taken effect"); + + commitDataset.commitExistingIndexSegments("idx_id", "id", Collections.singletonList(segment)); + } + + Assertions.assertTrue( + segments().isEmpty(), + "A stale segment commit must not resurrect a concurrently dropped index"); + } + + private void createTable() { + spark.sql(String.format("create table %s (id int, text string) using lance;", fullTable)); + } + + /** + * A full, non-default tokenizer configuration. Exercising every option rather than a single flag + * keeps the mismatch guard covering the whole configuration a refresh has to reproduce. + */ + private static final String FTS_OPTIONS = + "base_tokenizer = 'simple', language = 'English', max_token_length = 40, " + + "lower_case = true, stem = false, remove_stop_words = false, " + + "ascii_folding = false, with_position = true"; + + private void createFtsIndex() { + spark.sql( + String.format( + "alter table %s create index idx_text using fts (text) with (%s)", + fullTable, FTS_OPTIONS)); + } + + private long ftsMatchCount() { + return spark + .sql(String.format("select id from %s where lance_match(text, 'text')", fullTable)) + .count(); + } + + private long ftsPhraseCount(String phrase) { + return spark + .sql( + String.format( + "select id from %s where lance_match_phrase(text, '%s')", fullTable, phrase)) + .count(); + } + + /** Appends one fragment holding ids in [startInclusive, endExclusive). */ + private void appendFragment(int startInclusive, int endExclusive) { + List rows = + IntStream.range(startInclusive, endExclusive) + .boxed() + .map(i -> RowFactory.create(i, String.format("text_%d", i))) + .collect(Collectors.toList()); + try { + // coalesce(1) keeps each append to a single fragment, so fragment counts stay predictable. + spark.createDataFrame(rows, SCHEMA).coalesce(1).writeTo(fullTable).append(); + } catch (NoSuchTableException e) { + throw new IllegalStateException("Test table was not created: " + fullTable, e); + } + } + + private void createZonemapIndex() { + spark.sql(String.format("alter table %s create index idx_id using zonemap (id)", fullTable)); + } + + private Row refresh(String withClause) { + return spark + .sql(String.format("alter table %s refresh index idx_id %s", fullTable, withClause)) + .collectAsList() + .get(0); + } + + /** Fragment ids covered by every segment of the test index. */ + private Set coverage() { + Set covered = new HashSet<>(); + for (Index segment : segments()) { + covered.addAll(segment.fragments().orElse(Collections.emptyList())); + } + return covered; + } + + private Set segmentUuids() { + return segments().stream().map(Index::uuid).collect(Collectors.toSet()); + } + + private List segments() { + return segments("idx_id"); + } + + private List segments(String indexName) { + try (org.lance.Dataset dataset = + Utils.openDatasetBuilder(LanceSparkReadOptions.from(tableDir)).build()) { + return dataset.getIndexes().stream() + .filter(index -> indexName.equals(index.name())) + .collect(Collectors.toList()); + } + } + + private int liveFragmentCount() { + try (org.lance.Dataset dataset = + Utils.openDatasetBuilder(LanceSparkReadOptions.from(tableDir)).build()) { + return dataset.getFragmentStatistics().size(); + } + } + + private long datasetVersion() { + try (org.lance.Dataset dataset = + Utils.openDatasetBuilder(LanceSparkReadOptions.from(tableDir)).build()) { + return dataset.version(); + } + } + + private static String causeChain(Throwable error) { + StringBuilder messages = new StringBuilder(); + for (Throwable current = error; current != null; current = current.getCause()) { + messages.append(current.getMessage()).append(" | "); + } + return messages.toString(); + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseShowIndexesTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseShowIndexesTest.java index 0c2ff0bfd..b123be4da 100755 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseShowIndexesTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseShowIndexesTest.java @@ -112,7 +112,7 @@ public void testShowIndexes() { Dataset result = spark.sql(String.format("show indexes from %s", fullTable)); Assertions.assertEquals( - "StructType(StructField(name,StringType,true),StructField(fields,ArrayType(StringType,true),true),StructField(index_type,StringType,true),StructField(num_indexed_fragments,LongType,true),StructField(num_indexed_rows,LongType,true),StructField(num_unindexed_fragments,LongType,true),StructField(num_unindexed_rows,LongType,true))", + "StructType(StructField(name,StringType,true),StructField(fields,ArrayType(StringType,true),true),StructField(index_type,StringType,true),StructField(num_indexed_fragments,LongType,true),StructField(num_indexed_rows,LongType,true),StructField(num_unindexed_fragments,LongType,true),StructField(num_unindexed_rows,LongType,true),StructField(indexed_percent,DoubleType,true),StructField(num_segments,LongType,true),StructField(size_bytes,LongType,true))", result.schema().toString()); List rows = result.collectAsList(); @@ -138,6 +138,31 @@ public void testShowIndexes() { // num_indexed_rows should be at least 1 long numIndexedRows = row.getLong(4); Assertions.assertTrue(numIndexedRows >= 1L, "num_indexed_rows should be at least 1"); + + // a freshly created index covers every row + Assertions.assertEquals(100.0d, row.getDouble(7), 1e-9, "indexed_percent should be 100"); + + // one logical index backed by at least one physical segment + Assertions.assertTrue(row.getLong(8) >= 1L, "num_segments should be at least 1"); + + // a built index occupies storage + Assertions.assertTrue(row.getLong(9) > 0L, "size_bytes should be positive"); + } + + /** + * With no rows to divide, coverage is undefined rather than complete: reporting 100 would tell an + * operator polling for staleness that an index covering nothing is up to date. + */ + @Test + public void testShowIndexesReportsNullPercentForEmptyTable() { + spark.sql(String.format("create table %s (id int, text string) using lance;", fullTable)); + spark.sql(String.format("alter table %s create index test_index using btree (id)", fullTable)); + + Row row = spark.sql(String.format("show indexes from %s", fullTable)).collectAsList().get(0); + + Assertions.assertTrue(row.isNullAt(7), "indexed_percent should be null for an empty table"); + Assertions.assertEquals(0L, row.getLong(4), "no rows can be indexed"); + Assertions.assertEquals(0L, row.getLong(6), "no rows can be unindexed"); } @Test diff --git a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala index 9644e6f1c..8efcba66b 100644 --- a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala +++ b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala @@ -16,7 +16,9 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.sql.catalyst.plans.logical.LanceNamedArgument import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.Test -import org.lance.index.IndexType +import org.lance.index.{Index, IndexType} + +import scala.collection.JavaConverters._ /** * Unit tests for [[IndexUtils]] helper methods. @@ -31,6 +33,38 @@ class IndexUtilsTest { FragmentWorkload(java.lang.Integer.valueOf(fragmentId), rowCount) }.toList + /** An index segment carrying only the metadata these helpers read. */ + private def segment( + fragmentIds: Option[Seq[Int]], + indexDetails: Option[Array[Byte]] = None): Index = { + val builder = Index + .builder() + .uuid(java.util.UUID.randomUUID()) + .name("idx_id") + .indexType(IndexType.INVERTED) + fragmentIds.foreach(ids => + builder.fragments(ids.map(java.lang.Integer.valueOf).asJava)) + indexDetails.foreach(builder.indexDetails) + builder.build() + } + + private def coveringSegment(fragmentIds: Int*): Index = segment(Some(fragmentIds)) + + /** + * Asserts the batches partition the fragments into contiguous runs of the id order. + * + * Concatenating the batches in order and getting an ascending id sequence back is exactly that + * property: any partition whose concatenation is sorted consists of consecutive slices. + */ + private def assertContiguousBatches(batches: Seq[List[Integer]]): Unit = { + val flattened = batches.flatten.map(_.intValue) + assertEquals( + flattened.sorted, + flattened, + s"batches must partition the fragments in id order, got $batches") + batches.foreach(batch => assertFalse(batch.isEmpty, s"no batch may be empty, got $batches")) + } + // ── extractTrain ────────────────────────────────────────────────────────── @Test @@ -228,33 +262,519 @@ class IndexUtilsTest { assertEquals(Seq.empty, IndexUtils.batchFragments(Nil, None, 4)) } + /** + * Interleaved coverage makes Lance's compaction planner treat every adjacent fragment pair as + * ungroupable, so OPTIMIZE stops coalescing the table entirely. Batches must be contiguous runs. + */ + @Test + def batchFragments_producesContiguousRuns(): Unit = { + Seq( + fragmentWorkloads(1, 1, 1, 1, 1, 1), + fragmentWorkloads(100, 1, 1, 1), + fragmentWorkloads(1, 1, 1, 100), + fragmentWorkloads(5, 9, 2, 7, 3, 8, 1, 6), + fragmentWorkloads(0, 0, 0, 0, 0)).foreach { fragments => + (1 to fragments.size).foreach { segments => + val batches = IndexUtils.batchFragments(fragments, Some(segments), 4) + assertEquals( + segments, + batches.size, + s"expected $segments batches for ${fragments.size} fragments, got $batches") + assertContiguousBatches(batches) + assertEquals( + fragments.map(_.fragmentId), + batches.flatten, + "every fragment must be assigned exactly once") + } + } + } + @Test def batchFragments_balancesRowsDeterministically(): Unit = { val fragments = fragmentWorkloads(80, 50, 30, 20) + // Splitting after fragment 0 gives 80 / 100; the contiguous alternative gives 130 / 50. val expected = Seq( - List(java.lang.Integer.valueOf(0), java.lang.Integer.valueOf(3)), - List(java.lang.Integer.valueOf(1), java.lang.Integer.valueOf(2))) + List(java.lang.Integer.valueOf(0)), + List( + java.lang.Integer.valueOf(1), + java.lang.Integer.valueOf(2), + java.lang.Integer.valueOf(3))) assertEquals(expected, IndexUtils.batchFragments(fragments, Some(2), 4)) assertEquals(expected, IndexUtils.batchFragments(fragments.reverse, Some(2), 4)) } + /** A single dominant fragment must not collapse the requested parallelism. */ + @Test + def batchFragments_keepsRequestedParallelismUnderSkew(): Unit = { + val batches = IndexUtils.batchFragments(fragmentWorkloads(100, 1, 1, 1), Some(4), 4) + + assertEquals( + Seq( + List(java.lang.Integer.valueOf(0)), + List(java.lang.Integer.valueOf(1)), + List(java.lang.Integer.valueOf(2)), + List(java.lang.Integer.valueOf(3))), + batches) + } + @Test def batchFragments_distributesZeroRowFragmentsAcrossSegments(): Unit = { val fragments = fragmentWorkloads(0, 0, 0, 0).reverse assertEquals( Seq( - List(java.lang.Integer.valueOf(0), java.lang.Integer.valueOf(3)), + List(java.lang.Integer.valueOf(0)), List(java.lang.Integer.valueOf(1)), - List(java.lang.Integer.valueOf(2))), + List(java.lang.Integer.valueOf(2), java.lang.Integer.valueOf(3))), IndexUtils.batchFragments(fragments, Some(3), 4)) } + /** Fragment ids are not necessarily dense or zero-based once a table has been compacted. */ + @Test + def batchFragments_keepsSparseFragmentIdsContiguousByPosition(): Unit = { + val fragments = List( + FragmentWorkload(java.lang.Integer.valueOf(17), 10L), + FragmentWorkload(java.lang.Integer.valueOf(4), 10L), + FragmentWorkload(java.lang.Integer.valueOf(9), 10L), + FragmentWorkload(java.lang.Integer.valueOf(31), 10L)) + + assertEquals( + Seq( + List(java.lang.Integer.valueOf(4), java.lang.Integer.valueOf(9)), + List(java.lang.Integer.valueOf(17), java.lang.Integer.valueOf(31))), + IndexUtils.batchFragments(fragments, Some(2), 4)) + } + + private def workloads(batches: Seq[List[Integer]], rows: Seq[Long]): Seq[Long] = + batches.map(_.map(id => rows(id.intValue)).sum) + + /** Smallest achievable heaviest batch over all contiguous partitions into `segmentCount` runs. */ + private def optimalHeaviestBatch(rows: Seq[Long], segmentCount: Int): Long = { + def best(from: Int, runs: Int): Long = + if (runs == 1) { + rows.drop(from).sum + } else { + (from until rows.size - runs + 1).map { cut => + math.max(rows.slice(from, cut + 1).sum, best(cut + 1, runs - 1)) + }.min + } + best(0, segmentCount) + } + + /** + * The heaviest batch has to be as light as a contiguous partition allows. This workload is the + * counter-example that sank an earlier prefix-crossing heuristic: it cut after the three + * indivisible leading fragments had already overshot their even shares, leaving one batch of 162 + * where 95 is forced by fragment 0 alone. + */ + @Test + def batchFragments_minimisesTheHeaviestBatch(): Unit = { + val rows = Seq(95L, 93L, 89L, 8L, 1L, 4L, 74L, 88L, 38L) + val fragments = rows.zipWithIndex.map { case (count, fragmentId) => + FragmentWorkload(java.lang.Integer.valueOf(fragmentId), count) + }.toList + + val batches = IndexUtils.batchFragments(fragments, Some(6), 1) + + assertEquals(6, batches.size) + assertContiguousBatches(batches) + assertEquals( + 95L, + workloads(batches, rows).max, + s"expected the optimal heaviest batch, got ${workloads(batches, rows)}") + } + + /** + * Optimality is checked against every contiguous partition rather than against a fixed expected + * split, so the property is pinned instead of one of its consequences. + */ + @Test + def batchFragments_matchesTheOptimalContiguousPartition(): Unit = { + val workloadShapes = Seq( + Seq(95L, 93L, 89L, 8L, 1L, 4L, 74L, 88L, 38L), + Seq(100L, 1L, 1L, 1L, 1L, 1L), + Seq(1L, 1L, 1L, 1L, 1L, 100L), + Seq(5L, 9L, 2L, 7L, 3L, 8L, 1L, 6L), + Seq(7L, 7L, 7L, 7L, 7L, 7L, 7L), + Seq(50L, 1L, 50L, 1L, 50L, 1L, 50L), + Seq(0L, 0L, 5L, 0L, 0L)) + + workloadShapes.foreach { rows => + val fragments = rows.zipWithIndex.map { case (count, fragmentId) => + FragmentWorkload(java.lang.Integer.valueOf(fragmentId), count) + }.toList + (1 to rows.size).foreach { segmentCount => + val batches = IndexUtils.batchFragments(fragments, Some(segmentCount), 1) + assertEquals(segmentCount, batches.size, s"$rows into $segmentCount") + assertContiguousBatches(batches) + assertEquals( + optimalHeaviestBatch(rows, segmentCount), + workloads(batches, rows).max, + s"$rows into $segmentCount batches: got ${workloads(batches, rows)}") + } + } + } + @Test def batchFragments_rejectsWorkloadOverflow(): Unit = { assertThrows( classOf[ArithmeticException], () => IndexUtils.batchFragments(fragmentWorkloads(Long.MaxValue, 1), Some(1), 1)) } + + // ── methodForIndexType ──────────────────────────────────────────────────── + + @Test + def methodForIndexType_roundTripsEverySegmentBuildableType(): Unit = { + val types = Seq( + IndexType.BTREE, + IndexType.ZONEMAP, + IndexType.BITMAP, + IndexType.LABEL_LIST, + IndexType.NGRAM, + IndexType.BLOOM_FILTER, + IndexType.RTREE, + IndexType.INVERTED) + + types.foreach { indexType => + val method = IndexUtils.methodForIndexType(indexType) + assertTrue(method.isDefined, s"expected a method name for $indexType") + assertEquals( + Some(indexType), + IndexUtils.scalarSegmentIndexType(method.get), + s"method '${method.get}' should resolve back to $indexType") + } + } + + @Test + def methodForIndexType_mapsInvertedToCanonicalFtsSpelling(): Unit = { + assertEquals(Some("fts"), IndexUtils.methodForIndexType(IndexType.INVERTED)) + } + + @Test + def methodForIndexType_returnsNoneForVectorType(): Unit = { + assertEquals(None, IndexUtils.methodForIndexType(IndexType.VECTOR)) + } + + @Test + def methodForIndexType_returnsNoneForNull(): Unit = { + assertEquals(None, IndexUtils.methodForIndexType(null)) + } + + // ── isSystemIndex ───────────────────────────────────────────────────────── + + @Test + def isSystemIndex_matchesLanceMaintainedIndexesCaseInsensitively(): Unit = { + assertTrue(IndexUtils.isSystemIndex("__lance_frag_reuse")) + assertTrue(IndexUtils.isSystemIndex("__LANCE_MEM_WAL")) + } + + @Test + def isSystemIndex_rejectsUserIndexesAndNull(): Unit = { + assertFalse(IndexUtils.isSystemIndex("idx_id")) + assertFalse(IndexUtils.isSystemIndex(null)) + } + + // ── declaredCoverage / committedCoverage ────────────────────────────────── + + @Test + def declaredCoverage_unionsSegmentBitmaps(): Unit = { + assertEquals( + Set(0, 1, 4), + IndexUtils.declaredCoverage(Seq(coveringSegment(0, 1), coveringSegment(4)))) + } + + @Test + def declaredCoverage_treatsAbsentBitmapAsNoCoverage(): Unit = { + assertEquals(Set(2), IndexUtils.declaredCoverage(Seq(segment(None), coveringSegment(2)))) + assertEquals(Set.empty[Int], IndexUtils.declaredCoverage(Seq.empty)) + } + + /** A segment keeping its identity but reporting narrower coverage, as a pruned commit returns. */ + private def prunedTo(built: Index, fragmentIds: Seq[Int]): Index = + Index + .builder() + .uuid(built.uuid) + .name(built.name) + .indexType(built.indexType) + .fragments(fragmentIds.map(java.lang.Integer.valueOf).asJava) + .build() + + @Test + def requireCommittableCoverage_acceptsCoverageThatSurvives(): Unit = { + IndexUtils.requireCommittableCoverage(Set(1, 5), Seq(coveringSegment(0, 1)), "idx_id") + } + + @Test + def requireCommittableCoverage_acceptsSegmentsThatDeclareNothing(): Unit = { + IndexUtils.requireCommittableCoverage(Set(1), Seq(segment(None)), "idx_id") + IndexUtils.requireCommittableCoverage(Set.empty[Int], Seq.empty, "idx_id") + } + + @Test + def requireCommittableCoverage_refusesASetThatWouldCoverNothing(): Unit = { + val error = assertThrows( + classOf[IllegalStateException], + () => IndexUtils.requireCommittableCoverage(Set(5), Seq(coveringSegment(0, 7)), "idx_id")) + + assertTrue(error.getMessage.contains("idx_id"), error.getMessage) + assertTrue(error.getMessage.contains("0, 7"), error.getMessage) + assertTrue(error.getMessage.contains("re-run"), error.getMessage) + } + + @Test + def requireCommittableCoverage_summarizesLargeRetiredSets(): Unit = { + val error = assertThrows( + classOf[IllegalStateException], + () => + IndexUtils.requireCommittableCoverage( + Set.empty[Int], + Seq(coveringSegment(0 to 20: _*)), + "idx_id")) + + assertTrue(error.getMessage.contains("21 total"), error.getMessage) + } + + // ── establishedCoverage ─────────────────────────────────────────────────── + + @Test + def establishedCoverage_reportsWhatTheCommitReturned(): Unit = { + val built = Seq(coveringSegment(0, 1), coveringSegment(2)) + + assertEquals( + Set(0, 1, 2), + IndexUtils.establishedCoverage(built, built, Set(0, 1, 2), "idx_id")) + } + + /** + * Lance prunes a fragment whose indexed field was rewritten under the same id, so the committed + * bitmap can be narrower than the one handed in while every fragment is still live. No comparison + * of fragment ids can see that, which is why the report has to come from what the commit returned. + */ + @Test + def establishedCoverage_followsAPrunedCommitEvenWhileEveryFragmentIsLive(): Unit = { + val built = coveringSegment(0, 1) + + assertEquals( + Set(0), + IndexUtils.establishedCoverage( + Seq(built), + Seq(prunedTo(built, Seq(0))), + Set(0, 1), + "idx_id")) + } + + @Test + def establishedCoverage_excludesFragmentsRetiredByTheCommit(): Unit = { + val built = coveringSegment(0, 1) + + assertEquals( + Set(1), + IndexUtils.establishedCoverage(Seq(built), Seq(built), Set(1, 5), "idx_id")) + } + + /** Existing segments survive a commit they are disjoint from; their coverage is not ours. */ + @Test + def establishedCoverage_countsOnlyTheSegmentsThisBuildProduced(): Unit = { + val built = coveringSegment(2) + val survivor = coveringSegment(0, 1) + + assertEquals( + Set(2), + IndexUtils.establishedCoverage( + Seq(built), + Seq(survivor, built), + Set(0, 1, 2), + "idx_id")) + } + + @Test + def establishedCoverage_isEmptyWhenNothingOfThisBuildSurvived(): Unit = { + val built = coveringSegment(0) + + assertEquals( + Set.empty[Int], + IndexUtils.establishedCoverage( + Seq(built), + Seq(prunedTo(built, Seq.empty)), + Set(0), + "idx_id")) + } + + // ── retainedSegments ────────────────────────────────────────────────────── + + private def namedSegment(name: String, fragmentIds: Int*): Index = + Index + .builder() + .uuid(java.util.UUID.randomUUID()) + .name(name) + .indexType(IndexType.ZONEMAP) + .fragments(fragmentIds.map(java.lang.Integer.valueOf).asJava) + .build() + + @Test + def retainedSegments_keepsOnlySegmentsOfTheNamedIndexThatStillCoverLiveFragments(): Unit = { + val kept = namedSegment("idx_id", 0, 1) + val allRetired = namedSegment("idx_id", 7) + val otherIndex = namedSegment("idx_other", 0) + + assertEquals( + Seq(kept), + IndexUtils.retainedSegments(Seq(kept, allRetired, otherIndex), "idx_id", Set(0, 1))) + } + + /** + * A commit against a name Lance no longer knows creates that index instead of extending it, so a + * DROP INDEX during the build would otherwise be undone with only the planned coverage. + */ + @Test + def retainedSegments_failsWhenTheIndexIsGone(): Unit = { + val error = assertThrows( + classOf[IllegalStateException], + () => IndexUtils.retainedSegments(Seq(namedSegment("idx_other", 0)), "idx_id", Set(0))) + + assertTrue(error.getMessage.contains("idx_id"), error.getMessage) + assertTrue(error.getMessage.contains("no longer exists"), error.getMessage) + assertTrue(error.getMessage.contains("CREATE INDEX"), error.getMessage) + } + + @Test + def retainedSegments_matchesIndexNameExactly(): Unit = { + assertThrows( + classOf[IllegalStateException], + () => IndexUtils.retainedSegments(Seq(namedSegment("IDX_ID", 0)), "idx_id", Set(0))) + } + + // ── requireUniformSegmentDetails ────────────────────────────────────────── + + @Test + def requiresUniformSegmentDetails_onlyForInverted(): Unit = { + assertTrue(IndexUtils.requiresUniformSegmentDetails(IndexType.INVERTED)) + Seq( + IndexType.BTREE, + IndexType.ZONEMAP, + IndexType.BITMAP, + IndexType.LABEL_LIST, + IndexType.NGRAM, + IndexType.BLOOM_FILTER, + IndexType.RTREE, + IndexType.VECTOR).foreach { indexType => + assertFalse(IndexUtils.requiresUniformSegmentDetails(indexType), indexType.name()) + } + assertFalse(IndexUtils.requiresUniformSegmentDetails(null)) + } + + @Test + def requireUniformSegmentDetails_acceptsMatchingDetails(): Unit = { + val details = Array[Byte](1, 2, 3) + IndexUtils.requireUniformSegmentDetails( + IndexType.INVERTED, + "idx_id", + "fts", + Seq(segment(Some(Seq(0)), Some(details.clone()))), + Seq(segment(Some(Seq(1)), Some(details.clone())))) + } + + /** + * An inverted index loads one set of details for the whole logical index, so segments built with + * different options fail every full-text query. Rejecting before the commit leaves it untouched. + */ + @Test + def requireUniformSegmentDetails_rejectsDifferingDetails(): Unit = { + val error = assertThrows( + classOf[IllegalArgumentException], + () => + IndexUtils.requireUniformSegmentDetails( + IndexType.INVERTED, + "idx_id", + "fts", + Seq(segment(Some(Seq(0)), Some(Array[Byte](1, 2, 3)))), + Seq(segment(Some(Seq(1)), Some(Array[Byte](1, 2, 4)))))) + + assertTrue(error.getMessage.contains("idx_id"), error.getMessage) + assertTrue(error.getMessage.contains("fts"), error.getMessage) + assertTrue(error.getMessage.contains("CREATE INDEX"), error.getMessage) + } + + @Test + def requireUniformSegmentDetails_rejectsDisagreementWithinEitherSide(): Unit = { + val a = Array[Byte](1) + val b = Array[Byte](2) + assertThrows( + classOf[IllegalArgumentException], + () => + IndexUtils.requireUniformSegmentDetails( + IndexType.INVERTED, + "idx_id", + "fts", + Seq(segment(Some(Seq(0)), Some(a)), segment(Some(Seq(1)), Some(b))), + Seq(segment(Some(Seq(2)), Some(a))))) + } + + @Test + def requireUniformSegmentDetails_ignoresIndexTypesWithoutTheRequirement(): Unit = { + IndexUtils.requireUniformSegmentDetails( + IndexType.ZONEMAP, + "idx_id", + "zonemap", + Seq(segment(Some(Seq(0)), Some(Array[Byte](1)))), + Seq(segment(Some(Seq(1)), Some(Array[Byte](2))))) + } + + /** Nothing to compare against: a segment predating index details, or a wholesale replacement. */ + @Test + def requireUniformSegmentDetails_defersWhenEitherSideHasNoDetails(): Unit = { + IndexUtils.requireUniformSegmentDetails( + IndexType.INVERTED, + "idx_id", + "fts", + Seq(segment(Some(Seq(0)))), + Seq(segment(Some(Seq(1)), Some(Array[Byte](1))))) + IndexUtils.requireUniformSegmentDetails( + IndexType.INVERTED, + "idx_id", + "fts", + Seq.empty, + Seq(segment(Some(Seq(1)), Some(Array[Byte](1))))) + } + + // ── parseNumSegments ────────────────────────────────────────────────────── + + @Test + def parseNumSegments_acceptsPositiveIntegers(): Unit = { + assertEquals( + 8, + IndexUtils.parseNumSegments(LanceNamedArgument("num_segments", java.lang.Long.valueOf(8)))) + } + + @Test + def parseNumSegments_rejectsNonPositiveValues(): Unit = { + Seq(0L, -1L).foreach { value => + assertThrows( + classOf[IllegalArgumentException], + () => + IndexUtils.parseNumSegments( + LanceNamedArgument("num_segments", java.lang.Long.valueOf(value)))) + } + } + + @Test + def parseNumSegments_rejectsValuesWiderThanInt(): Unit = { + assertThrows( + classOf[IllegalArgumentException], + () => + IndexUtils.parseNumSegments( + LanceNamedArgument("num_segments", java.lang.Long.valueOf(Int.MaxValue.toLong + 1)))) + } + + @Test + def parseNumSegments_rejectsNonNumericAndNullValues(): Unit = { + assertThrows( + classOf[IllegalArgumentException], + () => IndexUtils.parseNumSegments(LanceNamedArgument("num_segments", "eight"))) + assertThrows( + classOf[IllegalArgumentException], + () => IndexUtils.parseNumSegments(LanceNamedArgument("num_segments", null))) + } }