Skip to content

feat(python): expose LABEL_LIST segment index build #7884

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
Xuanwo merged 3 commits into from
Jul 28, 2026
Merged

feat(python): expose LABEL_LIST segment index build #7884

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7dd7dda
feat(python): expose LABEL_LIST segment index build
ddupg Jul 21, 2026
cee695f
fix(index): reject LABEL_LIST distributed index UUIDs
ddupg Jul 27, 2026
4f0345f
fix(index): reject nested LABEL_LIST item types
ddupg Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 16 additions & 10 deletions python/python/lance/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3252,8 +3252,12 @@ def _prepare_scalar_index_request(
", float, bool, str, large_str, fixed-size-binary, or temporal",
)
elif index_type == "LABEL_LIST":
if not pa.types.is_list(field_type):
raise TypeError(f"LABEL_LIST index column {column} must be a list")
if not (
pa.types.is_list(field_type) or pa.types.is_large_list(field_type)

@Xuanwo Xuanwo Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Accepting every LargeList here expands the supported schema beyond what LabelList bitmap keys can order, so index creation can succeed and a later query can panic. LargeList<List<Int64>> passes this check; the first indexed query reaches the (List, List) => todo!() comparator in btree.rs and surfaces ArrowInvalid: Task was aborted. The flat query returns the correct row.

The following ran on this head and fails at the final call:

with tempfile.TemporaryDirectory() as uri:
    ds = lance.write_dataset(
        pa.table({
            "id": pa.array([0, 1], type=pa.int32()),
            "labels": pa.array(
                [[[1, 2]], [[3]]],
                type=pa.large_list(pa.list_(pa.int64())),
            ),
        }),
        uri,
    )
    predicate = "array_has_any(labels, [[1, 2]])"
    assert ds.scanner(
        filter=predicate, use_scalar_index=False
    ).to_table()["id"].to_pylist() == [0]
    ds.create_scalar_index("labels", "LABEL_LIST")
    assert "ScalarIndexQuery" in ds.scanner(filter=predicate).explain_plan()
    ds.scanner(filter=predicate).to_table()  # Rust panic

The comparator limitation predates this PR for nested List, but this change newly accepts the equivalent LargeList schema and makes its any/all path reachable. A bounded fix could validate supported element types in Rust before writing the artifact, mirror that check in Python, and retain scan fallback for existing unsupported indexes.

@ddupg ddupg Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch. LABEL_LIST now rejects nested item types for both List and LargeList in Rust Core. The query parser applies the same check so legacy invalid indexes fall back to a scan instead of reaching the unsupported bitmap comparison. Both paths are covered by parameterized tests.

):
raise TypeError(
f"LABEL_LIST index column {column} must be a list or large list"
)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
elif index_type == "NGRAM":
if not pa.types.is_string(field_type) and not pa.types.is_large_string(
field_type
Expand Down Expand Up @@ -3311,6 +3315,7 @@ def _is_segment_native_scalar_index_type(
"RTREE",
"ZONEMAP",
"BLOOMFILTER",
"LABEL_LIST",

@Xuanwo Xuanwo Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Making LABEL_LIST segment-native exposes caller-supplied index_uuid without the collision guard used by BTree/RTree. LabelList writes _indices/<uuid>/bitmap_page_lookup.lance with replace semantics, so reusing a committed UUID overwrites a live index before any manifest commit.

On this head I published an index for labels_a, then staged labels_b with the live UUID. The dataset version stayed unchanged, but the live file hash changed. After reopening:

aaaa: flat=[0, 2, 4, 6], indexed=[]
cccc: flat=[], indexed=[0, 2, 4, 6]

This creates both false negatives and false positives without calling commit_existing_index_segments. LABEL_LIST could reject caller-supplied UUIDs before any I/O, with a regression asserting that the original artifact hash and indexed results remain unchanged after a rejected collision.

@ddupg ddupg Jul 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added a Core guard to reject caller-supplied UUIDs for distributed LABEL_LIST builds before any I/O.

}

@classmethod
Expand All @@ -3325,6 +3330,7 @@ def _requires_uncommitted_scalar_index(
"RTREE",
"ZONEMAP",
"BLOOMFILTER",
"LABEL_LIST",
}

def create_scalar_index(
Expand Down Expand Up @@ -4281,10 +4287,9 @@ def create_index_uncommitted(

This is the public distributed-build API for vector, BTREE scalar,
canonical bitmap scalar, INVERTED scalar, NGRAM scalar, RTREE scalar,
ZONEMAP scalar,
and BLOOMFILTER scalar index construction. Unlike
:meth:`create_index`, this method does not publish the index into the
dataset manifest. Instead, it writes one segment under
ZONEMAP scalar, BLOOMFILTER scalar, and LABEL_LIST scalar index construction.
Unlike :meth:`create_index`, this method does not publish the index into
the dataset manifest. Instead, it writes one segment under
``_indices/<segment_uuid>/`` and returns the resulting
:class:`Index` metadata.

Expand All @@ -4298,10 +4303,11 @@ def create_index_uncommitted(
4. commit the final segment list with
:meth:`commit_existing_index_segments`

BTREE, BITMAP, INVERTED, NGRAM, RTREE, ZONEMAP, and BLOOMFILTER segments may
be merged with :meth:`merge_existing_index_segments` before commit.
NGRAM segments built before a deferred compaction must be merged before
commit so their postings can be rebuilt against current row addresses.
BTREE, BITMAP, INVERTED, NGRAM, RTREE, ZONEMAP, BLOOMFILTER, and
LABEL_LIST segments may be merged with
:meth:`merge_existing_index_segments` before commit. NGRAM segments
built before a deferred compaction must be merged before commit so
their postings can be rebuilt against current row addresses.
Parameters are the same as :meth:`create_index`, with one additional
requirement:

Expand Down
67 changes: 67 additions & 0 deletions python/python/tests/test_scalar_index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4882,6 +4882,73 @@ def test_ngram_segment_merge_and_commit_from_python(tmp_path):
assert ds.count_rows("text IS NULL") == 1


@pytest.mark.parametrize(
"label_type",
[pa.list_(pa.string()), pa.large_list(pa.string())],
ids=["list", "large_list"],
)
def test_label_list_segment_index(tmp_path, label_type):
rows_per_fragment = 8
ds = lance.write_dataset(
pa.table(
{
"id": pa.array(range(rows_per_fragment * 4), type=pa.int32()),
"labels": pa.array(
[
["distributed"] if row_id % 2 == 0 else ["other"]
for row_id in range(rows_per_fragment * 4)
],
type=label_type,
),
}
),
tmp_path,
max_rows_per_file=rows_per_fragment,
)

fragment_ids = [fragment.fragment_id for fragment in ds.get_fragments()]
assert len(fragment_ids) == 4

with pytest.raises(ValueError, match="create_index_uncommitted"):
ds.create_scalar_index(
column="labels",
index_type="LABEL_LIST",
fragment_ids=[fragment_ids[0]],
)

index_name = "labels_segment_idx"
segments = [
ds.create_index_uncommitted(
column="labels",
index_type="LABEL_LIST",
name=index_name,
fragment_ids=[fragment_id],
)
for fragment_id in fragment_ids
]

merged_segment = ds.merge_existing_index_segments(segments)
ds = ds.commit_existing_index_segments(index_name, "labels", [merged_segment])

filter_expr = "array_has_any(labels, ['distributed'])"
without_index = ds.scanner(
filter=filter_expr,
columns=["id", "labels"],
use_scalar_index=False,
).to_table()
with_index = ds.scanner(
filter=filter_expr,
columns=["id", "labels"],
use_scalar_index=True,
).to_table()

assert with_index.equals(without_index)
assert (
"ScalarIndexQuery"
in ds.scanner(filter=filter_expr, use_scalar_index=True).explain_plan()
)


def test_zonemap_fragment_ids_parameter_validation(tmp_path):
ds = generate_multi_fragment_dataset(
tmp_path, num_fragments=2, rows_per_fragment=100
Expand Down
140 changes: 111 additions & 29 deletions rust/lance-index/src/scalar/expression.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use tokio::try_join;

use super::{
AnyQuery, BloomFilterQuery, LabelListQuery, MetricsCollector, SargableQuery, ScalarIndex,
SearchResult, TextQuery, TokenQuery,
SearchResult, TextQuery, TokenQuery, label_list::validate_label_list_data_type,
};
#[cfg(feature = "geo")]
use super::{GeoQuery, RelationQuery};
Expand Down Expand Up @@ -774,6 +774,11 @@ impl ScalarQueryParser for LabelListQueryParser {
if args.len() != 2 {
return None;
}
// LABEL_LIST stores unnested items as bitmap keys. Nested items cannot be
// ordered by the bitmap lookup, so legacy indexes must fall back to a scan.
if validate_label_list_data_type(data_type).is_err() {
return None;
}
// DataFusion normalizes array_contains to array_has
if func.name() == "array_has" {
let inner_type = match data_type {
Expand All @@ -796,34 +801,34 @@ impl ScalarQueryParser for LabelListQueryParser {
}

let label_list = maybe_scalar(&args[1], data_type)?;
if let ScalarValue::List(list_arr) = label_list {
let list_values = list_arr.values();
if list_values.is_empty() {
return None;
}
let mut scalars = Vec::with_capacity(list_values.len());
for idx in 0..list_values.len() {
scalars.push(ScalarValue::try_from_array(list_values.as_ref(), idx).ok()?);
}
if func.name() == "array_has_all" {
let query = LabelListQuery::HasAllLabels(scalars);
Some(IndexedExpression::index_query(
column.to_string(),
self.index_name.clone(),
self.index_type.clone(),
Arc::new(query),
))
} else if func.name() == "array_has_any" {
let query = LabelListQuery::HasAnyLabel(scalars);
Some(IndexedExpression::index_query(
column.to_string(),
self.index_name.clone(),
self.index_type.clone(),
Arc::new(query),
))
} else {
None
}
let list_values = match label_list {
ScalarValue::List(list_arr) => list_arr.values().clone(),
ScalarValue::LargeList(list_arr) => list_arr.values().clone(),

@Xuanwo Xuanwo Jul 24, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This new exact LargeList arm makes released nullable v0 LabelList indexes return wrong rows under negation. Lance v3.0.1 could write LargeList LabelList indexes with index_version = 0 and LabelListIndexDetails, but without list-level NULL metadata. The current compatibility path only warns, and read_list_nulls treats the missing metadata as an empty NULL set, so complementing NOT array_has_any/all includes NULL-list rows.

On this head, a v0 artifact over rows [[1], NULL, [2]] produced flat = [2] but indexed = [1, 2] for NOT array_has_any(labels, [1]); the plan confirmed ScalarIndexQuery@legacy_labels(LabelList).

A bounded compatibility fix could keep nullable v0 segments on scan/recheck, or reject them, before enabling exact LargeList planning. A historical-format regression covering both negated predicates would lock this down.

@ddupg ddupg Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the concrete repro. I agree that nullable v0 LabelList indices lack list-level NULL metadata, so NOT queries can be incorrect.

However, this limitation predates this PR: array_has already uses the index for both List and LargeList, while array_has_any/all already use it for List.
#6044 intentionally kept v0 indices readable with a rebuild warning. Newly created LabelList indices use v1 and persist the required NULL metadata.

This PR only extends the same behavior to LargeList array_has_any/all. Handling v0 differently would require a broader change across index loading and all LabelList predicates, rather than a local fix for this LargeList arm.

Could we keep the broader v0 handling out of scope for this PR and rely on the existing rebuild warning for legacy indices?

_ => return None,
};
if list_values.is_empty() {
return None;
}
let mut scalars = Vec::with_capacity(list_values.len());
for idx in 0..list_values.len() {
scalars.push(ScalarValue::try_from_array(list_values.as_ref(), idx).ok()?);
}
if func.name() == "array_has_all" {
let query = LabelListQuery::HasAllLabels(scalars);
Some(IndexedExpression::index_query(
column.to_string(),
self.index_name.clone(),
self.index_type.clone(),
Arc::new(query),
))
} else if func.name() == "array_has_any" {
let query = LabelListQuery::HasAnyLabel(scalars);
Some(IndexedExpression::index_query(
column.to_string(),
self.index_name.clone(),
self.index_type.clone(),
Arc::new(query),
))
} else {
None
}
Expand Down Expand Up @@ -2606,6 +2611,7 @@ mod tests {
use lance_datafusion::exec::{LanceExecutionOptions, get_session_context};
use lance_select::result::IndexExprResultWireFormat;
use roaring::RoaringBitmap;
use rstest::rstest;

use crate::scalar::json::{JsonQuery, JsonQueryParser};

Expand Down Expand Up @@ -2667,6 +2673,16 @@ mod tests {
Field::new("price", DataType::Float32, false),
Field::new("json", DataType::LargeBinary, false),
]);
check_with_schema(index_info, expr, expected, optimize, schema);
}

fn check_with_schema(
index_info: &dyn IndexInformationProvider,
expr: &str,
expected: Option<IndexedExpression>,
optimize: bool,
schema: Schema,
) {
let df_schema: DFSchema = schema.try_into().unwrap();

let ctx = get_session_context(&LanceExecutionOptions::default());
Expand Down Expand Up @@ -2756,6 +2772,72 @@ mod tests {
)
}

#[rstest]
#[case::list(DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))))]
#[case::large_list(DataType::LargeList(Arc::new(Field::new("item", DataType::Utf8, true))))]
fn test_label_list_query_parser(#[case] label_type: DataType) {
let index_info = MockIndexInfoProvider::new(vec![(
"labels",
ColInfo::new(
label_type.clone(),
Box::new(LabelListQueryParser::new(
"labels_idx".to_string(),
"LabelList".to_string(),
)),
),
)]);
let schema = Schema::new(vec![Field::new("labels", label_type, true)]);

check_with_schema(
&index_info,
"array_has_any(labels, ['distributed'])",
Some(IndexedExpression::index_query(
"labels".to_string(),
"labels_idx".to_string(),
"LabelList".to_string(),
Arc::new(LabelListQuery::HasAnyLabel(vec![ScalarValue::Utf8(Some(
"distributed".to_string(),
))])),
)),
true,
schema,
);
}

#[rstest]
#[case::list(DataType::List(Arc::new(Field::new(
"item",
DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
true,
))))]
#[case::large_list(DataType::LargeList(Arc::new(Field::new(
"item",
DataType::List(Arc::new(Field::new("item", DataType::Int64, true))),
true,
))))]
fn test_label_list_nested_item(#[case] label_type: DataType) {
let index_info = MockIndexInfoProvider::new(vec![(
"labels",
ColInfo::new(
label_type.clone(),
Box::new(LabelListQueryParser::new(
"labels_idx".to_string(),
"LabelList".to_string(),
)),
),
)]);
let schema = Schema::new(vec![Field::new("labels", label_type, true)]);

// Nested item types must not produce a scalar index query.
check_with_schema(
&index_info,
"array_has_any(labels, [[1]])",
None,
true,
schema,
);
}

#[test]
fn test_expressions() {
let index_info = MockIndexInfoProvider::new(vec![
Expand Down
Loading
Loading