diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index a45fb9e0202..8587ed3badf 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -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) + ): + raise TypeError( + f"LABEL_LIST index column {column} must be a list or large list" + ) elif index_type == "NGRAM": if not pa.types.is_string(field_type) and not pa.types.is_large_string( field_type @@ -3311,6 +3315,7 @@ def _is_segment_native_scalar_index_type( "RTREE", "ZONEMAP", "BLOOMFILTER", + "LABEL_LIST", } @classmethod @@ -3325,6 +3330,7 @@ def _requires_uncommitted_scalar_index( "RTREE", "ZONEMAP", "BLOOMFILTER", + "LABEL_LIST", } def create_scalar_index( @@ -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//`` and returns the resulting :class:`Index` metadata. @@ -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: diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index bcb5cece18b..c85be5dd99c 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -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 diff --git a/rust/lance-index/src/scalar/expression.rs b/rust/lance-index/src/scalar/expression.rs index 6fe635948cb..63d2922a0ff 100644 --- a/rust/lance-index/src/scalar/expression.rs +++ b/rust/lance-index/src/scalar/expression.rs @@ -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}; @@ -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 { @@ -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(), + _ => 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 } @@ -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}; @@ -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, + optimize: bool, + schema: Schema, + ) { let df_schema: DFSchema = schema.try_into().unwrap(); let ctx = get_session_context(&LanceExecutionOptions::default()); @@ -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![ diff --git a/rust/lance-index/src/scalar/label_list.rs b/rust/lance-index/src/scalar/label_list.rs index b96a3ba0346..bc250dd27a8 100644 --- a/rust/lance-index/src/scalar/label_list.rs +++ b/rust/lance-index/src/scalar/label_list.rs @@ -672,6 +672,33 @@ impl CacheKey for LabelListIndexStateKey { #[derive(Debug, Default)] pub struct LabelListIndexPlugin; +pub(super) fn validate_label_list_data_type(data_type: &DataType) -> Result<()> { + let item_type = match data_type { + DataType::List(item_field) | DataType::LargeList(item_field) => item_field.data_type(), + _ => { + return Err(Error::invalid_input_source( + format!( + "LabelList index can only be created on List or LargeList type columns. Column has type {:?}", + data_type + ) + .into(), + )); + } + }; + + if item_type.is_nested() { + return Err(Error::invalid_input_source( + format!( + "LabelList index item type must be non-nested. Column has type {:?}", + data_type + ) + .into(), + )); + } + + Ok(()) +} + #[async_trait] impl BasicTrainer for LabelListIndexPlugin { fn new_training_request( @@ -679,16 +706,7 @@ impl BasicTrainer for LabelListIndexPlugin { _params: &str, field: &Field, ) -> Result> { - if !matches!( - field.data_type(), - DataType::List(_) | DataType::LargeList(_) - ) { - return Err(Error::invalid_input_source(format!( - "LabelList index can only be created on List or LargeList type columns. Column has type {:?}", - field.data_type() - ) - .into())); - } + validate_label_list_data_type(field.data_type())?; Ok(Box::new(DefaultTrainingRequest::new( TrainingCriteria::new(TrainingOrdering::None).with_row_id(), @@ -723,16 +741,7 @@ impl BasicTrainer for LabelListIndexPlugin { })? .1; - if !matches!( - field.data_type(), - DataType::List(_) | DataType::LargeList(_) - ) { - return Err(Error::invalid_input_source(format!( - "LabelList index can only be created on List or LargeList type columns. Column has type {:?}", - field.data_type() - ) - .into())); - } + validate_label_list_data_type(field.data_type())?; let list_nulls = Arc::new(Mutex::new(RowAddrTreeMap::new())); let data = track_list_nulls(data, list_nulls.clone()); @@ -846,11 +855,42 @@ mod tests { use datafusion_common::ScalarValue; use lance_core::cache::CacheCodec; use lance_core::utils::address::RowAddress; + use rstest::rstest; use super::super::bitmap::BitmapIndexState; use super::super::btree::OrderableScalarValue; use super::*; + #[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_rejects_nested_item_type(#[case] data_type: DataType) { + let field = Field::new(VALUE_COLUMN_NAME, data_type, true); + let error = LabelListIndexPlugin + .new_training_request("", &field) + .err() + .expect("nested item type should be rejected"); + + assert!( + matches!(error, Error::InvalidInput { .. }), + "expected invalid input error, got: {error}" + ); + assert!( + error + .to_string() + .contains("LabelList index item type must be non-nested"), + "unexpected error: {error}" + ); + } + fn sample_state() -> LabelListIndexState { let mut index_map = BTreeMap::new(); for k in 0..32i32 { diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index b12a7198e5a..19ef499f1e4 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -19,11 +19,11 @@ use crate::index::DatasetIndexExt; use arrow::array::{AsArray, GenericListBuilder, GenericStringBuilder}; use arrow::datatypes::UInt64Type; use arrow_array::RecordBatch; -use arrow_array::{Array, GenericStringArray, StructArray, UInt64Array}; +use arrow_array::{Array, GenericStringArray, LargeListArray, ListArray, StructArray, UInt64Array}; use arrow_array::{ ArrayRef, Float32Array, Int32Array, RecordBatchIterator, StringArray, builder::StringDictionaryBuilder, - types::{Float32Type, Int32Type}, + types::{Float32Type, Int32Type, Int64Type}, }; use arrow_schema::{ DataType, Field as ArrowField, Field, Fields as ArrowFields, Schema as ArrowSchema, @@ -2904,6 +2904,95 @@ async fn test_bitmap_prewarm_with_serializing_backend_serves_query_with_no_io() ); } +#[rstest] +#[case::list(false)] +#[case::large_list(true)] +#[tokio::test] +async fn test_label_list_index_types(#[case] large_list: bool) { + let test_uri = TempStrDir::default(); + let label_values = vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(2)]), + Some(vec![Some(1)]), + Some(vec![Some(3)]), + ]; + let labels: ArrayRef = if large_list { + Arc::new(LargeListArray::from_iter_primitive::( + label_values, + )) + } else { + Arc::new(ListArray::from_iter_primitive::( + label_values, + )) + }; + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("id", DataType::Int32, false), + ArrowField::new("labels", labels.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from(vec![0, 1, 2, 3])), labels], + ) + .unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let mut dataset = Dataset::write(reader, &test_uri, None).await.unwrap(); + + let expected = dataset + .scan() + .project(&["id"]) + .unwrap() + .filter("array_has_any(labels, [1])") + .unwrap() + .try_into_batch() + .await + .unwrap(); + let expected_ids = expected + .column(0) + .as_primitive::() + .values() + .to_vec(); + assert_eq!(expected_ids, vec![0, 2]); + + dataset + .create_index( + &["labels"], + IndexType::LabelList, + Some("labels_idx".to_owned()), + &ScalarIndexParams::default(), + true, + ) + .await + .unwrap(); + + let result = dataset + .scan() + .project(&["id"]) + .unwrap() + .filter("array_has_any(labels, [1])") + .unwrap() + .try_into_batch() + .await + .unwrap(); + let result_ids = result + .column(0) + .as_primitive::() + .values() + .to_vec(); + assert_eq!(result_ids, expected_ids); + + let plan = dataset + .scan() + .filter("array_has_any(labels, [1])") + .unwrap() + .explain_plan(false) + .await + .unwrap(); + assert!( + plan.contains("ScalarIndexQuery") && plan.contains("LabelList"), + "Expected LabelList scalar index query in plan: {plan}" + ); +} + /// LabelList analogue: after prewarming, an `array_has_any` query against a /// `LabelList` index serves results without any further IO. Exercises the /// `LabelListIndexState` codec (which embeds the inner bitmap state and the diff --git a/rust/lance/src/index/create.rs b/rust/lance/src/index/create.rs index 6a55ffb2f6a..505cd6f3c07 100644 --- a/rust/lance/src/index/create.rs +++ b/rust/lance/src/index/create.rs @@ -878,6 +878,10 @@ fn ensure_index_uuid_allowed( || scalar_index_type.is_some_and(|index_type| index_type.eq_ignore_ascii_case("ngram")) { Some("NGram") + } else if index_type == IndexType::LabelList + || scalar_index_type.is_some_and(|index_type| index_type.eq_ignore_ascii_case("labellist")) + { + Some("LabelList") } else { None }; @@ -1693,6 +1697,110 @@ mod tests { } } + #[tokio::test] + async fn test_label_list_distributed_index_uuid_collision_rejected() { + use lance_datagen::{Dimension, array}; + use lance_index::scalar::label_list::BITMAP_LOOKUP_NAME; + + let test_dir = TempStrDir::default(); + let labels = |label| { + array::cycle_vec_var( + array::cycle::(vec![label]), + Dimension::from(1), + Dimension::from(2), + ) + }; + let mut dataset = gen_batch() + .col("labels_a", labels(1)) + .col("labels_b", labels(3)) + .into_dataset( + test_dir.as_str(), + FragmentCount::from(2), + FragmentRowCount::from(4), + ) + .await + .unwrap(); + let params = + ScalarIndexParams::for_builtin(lance_index::scalar::BuiltinIndexType::LabelList); + let live_index = dataset + .create_index( + &["labels_a"], + IndexType::LabelList, + Some("labels_a_idx".to_string()), + ¶ms, + false, + ) + .await + .unwrap(); + let artifact_path = dataset + .indices_dir() + .join(live_index.uuid.to_string()) + .join(BITMAP_LOOKUP_NAME); + let artifact_before = dataset + .object_store + .read_one_all(&artifact_path) + .await + .unwrap(); + let version_before = dataset.version().version; + let fragment_id = dataset.get_fragments()[0].id() as u32; + + for index_type in [IndexType::LabelList, IndexType::Scalar] { + let err = dataset + .create_index_builder(&["labels_b"], index_type, ¶ms) + .name("labels_b_idx".to_string()) + .fragments(vec![fragment_id]) + .index_uuid(live_index.uuid) + .execute_uncommitted() + .await + .unwrap_err(); + + assert!( + matches!(err, Error::InvalidInput { .. }), + "expected invalid input error, got: {err}" + ); + assert!( + err.to_string().contains( + "index_uuid is no longer accepted for LabelList distributed index builds" + ), + "unexpected error: {err}" + ); + } + assert_eq!(dataset.version().version, version_before); + let artifact_after = dataset + .object_store + .read_one_all(&artifact_path) + .await + .unwrap(); + assert_eq!(artifact_after, artifact_before); + + let dataset = Dataset::open(test_dir.as_str()).await.unwrap(); + assert_eq!( + dataset + .count_rows(Some("array_has_any(labels_a, [1])".to_string())) + .await + .unwrap(), + 8 + ); + assert_eq!( + dataset + .count_rows(Some("array_has_any(labels_a, [3])".to_string())) + .await + .unwrap(), + 0 + ); + let plan = dataset + .scan() + .filter("array_has_any(labels_a, [1])") + .unwrap() + .explain_plan(false) + .await + .unwrap(); + assert!( + plan.contains("ScalarIndexQuery") && plan.contains("LabelList"), + "expected LabelList scalar index query in plan: {plan}" + ); + } + #[tokio::test] async fn test_range_based_btree_index_create() { use crate::dataset::scanner::ColumnOrdering;