feat(index): support flattened JSON sub-doc indexing

java/src/main/java/org/lance/index/scalar/InvertedIndexParams.java

@@ -57,6 +57,9 @@ public static final class Builder {
     private Integer maxNgramLength;
     private Boolean prefixOnly;
     private Integer blockSize = 128;
+    private Boolean disableCrossArrayUnnest;
+    private Long maxSubDocsPerRow;
+    private MaxSubDocsPerRowExceedAction maxSubDocsPerRowExceedAction;
     private Boolean splitIdentifiers;
     private Boolean splitOnNumerics;
     private Boolean preserveOriginal;
@@ -304,6 +307,53 @@ public Builder blockSize(int blockSize) {
       return this;
     }
 
+    /**
+     * Configure whether flattened JSON tokenization avoids cross-array unnesting.
+     *
+     * <p>When true, sibling arrays are indexed independently instead of producing their Cartesian
+     * product. This can reduce index build memory for JSON records with multiple arrays but can
+     * sacrifice result accuracy for queries that constrain values across those arrays. The default
+     * is false.
+     *
+     * @param disableCrossArrayUnnest whether to avoid cross-array unnesting
+     * @return this builder
+     */
+    public Builder disableCrossArrayUnnest(boolean disableCrossArrayUnnest) {
+      this.disableCrossArrayUnnest = disableCrossArrayUnnest;
+      return this;
+    }
+
+    /**
+     * Limit the number of flattened sub-documents emitted for one JSON row.
+     *
+     * <p>If unset, the number of sub-documents is unlimited.
+     *
+     * @param maxSubDocsPerRow maximum sub-documents per row, must be positive
+     * @return this builder
+     * @throws IllegalArgumentException if {@code maxSubDocsPerRow} is not positive
+     */
+    public Builder maxSubDocsPerRow(long maxSubDocsPerRow) {
+      if (maxSubDocsPerRow <= 0) {
+        throw new IllegalArgumentException("maxSubDocsPerRow must be positive");
+      }
+      this.maxSubDocsPerRow = maxSubDocsPerRow;
+      return this;
+    }
+
+    /**
+     * Configure the action taken when {@link #maxSubDocsPerRow(long)} is exceeded.
+     *
+     * <p>The default is {@link MaxSubDocsPerRowExceedAction#FAIL}.
+     *
+     * @param action action to take when the limit is exceeded
+     * @return this builder
+     */
+    public Builder maxSubDocsPerRowExceedAction(MaxSubDocsPerRowExceedAction action) {
+      this.maxSubDocsPerRowExceedAction =
+          Objects.requireNonNull(action, "maxSubDocsPerRowExceedAction must not be null");
+      return this;
+    }
+
     /**
      * Configure whether code identifiers are split into subwords.
      *
@@ -501,6 +551,16 @@ public ScalarIndexParams build() {
       if (blockSize != null) {
         params.put("block_size", blockSize);
       }
+      if (disableCrossArrayUnnest != null) {
+        params.put("disable_cross_array_unnest", disableCrossArrayUnnest);
+      }
+      if (maxSubDocsPerRow != null) {
+        params.put("max_sub_docs_per_row", maxSubDocsPerRow);
+      }
+      if (maxSubDocsPerRowExceedAction != null) {
+        params.put(
+            "max_sub_docs_per_row_exceed_action", maxSubDocsPerRowExceedAction.toRustString());
+      }
       if (splitIdentifiers != null) {
         params.put("split_identifiers", splitIdentifiers);
       }

java/src/main/java/org/lance/index/scalar/MaxSubDocsPerRowExceedAction.java

@@ -0,0 +1,32 @@
+/*
+ * 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.index.scalar;
+
+/** Action taken when a JSON row exceeds the configured flattened sub-document limit. */
+public enum MaxSubDocsPerRowExceedAction {
+  /** Abort index ingestion with an error. */
+  FAIL("fail"),
+  /** Omit the source row from the index and continue ingestion. */
+  SKIP_ROW("skip_row");
+
+  private final String rustValue;
+
+  MaxSubDocsPerRowExceedAction(String rustValue) {
+    this.rustValue = rustValue;
+  }
+
+  String toRustString() {
+    return rustValue;
+  }
+}

java/src/test/java/org/lance/index/scalar/InvertedIndexParamsTest.java

@@ -55,6 +55,29 @@ void documentGranularityIsSerialized() {
     assertEquals("list_element", json.get("document_granularity"));
   }
 
+  @Test
+  void disableCrossArrayUnnestIsSerialized() {
+    ScalarIndexParams params = InvertedIndexParams.builder().disableCrossArrayUnnest(true).build();
+
+    Map<String, Object> json = JsonUtils.fromJson(params.getJsonParams().orElseThrow());
+    assertEquals(true, json.get("disable_cross_array_unnest"));
+  }
+
+  @Test
+  void maxSubDocsPerRowOptionsAreSerialized() {
+    for (MaxSubDocsPerRowExceedAction action : MaxSubDocsPerRowExceedAction.values()) {
+      ScalarIndexParams params =
+          InvertedIndexParams.builder()
+              .maxSubDocsPerRow(128)
+              .maxSubDocsPerRowExceedAction(action)
+              .build();
+
+      Map<String, Object> json = JsonUtils.fromJson(params.getJsonParams().orElseThrow());
+      assertEquals(128, ((Number) json.get("max_sub_docs_per_row")).longValue());
+      assertEquals(action.toRustString(), json.get("max_sub_docs_per_row_exceed_action"));
+    }
+  }
+
   @Test
   void blockSizeIsSerialized() {
     ScalarIndexParams params = InvertedIndexParams.builder().blockSize(128).build();

protos/index_old.proto

@@ -57,6 +57,11 @@ message InvertedIndexDetails {
     LIST_ELEMENT = 1;
   }
 
+  enum MaxSubDocsPerRowExceedAction {
+    FAIL = 0;
+    SKIP_ROW = 1;
+  }
+
   message CodeTokenizerConfig {
     /* Split one lexical identifier into subwords, e.g. getUserName ->
      * get/user/name.
@@ -115,4 +120,18 @@ message InvertedIndexDetails {
    * which identifies the overall inverted-index layout.
    */
   optional uint32 posting_format_version = 15;
+  // JSON document tokenization mode. Absent means SingleDocument JSON tokenization,
+  // which is how indexes written before flattened JSON sub-docs are interpreted.
+  optional string json_tokenizer_mode = 16;
+  // If true, avoid cross-array unnesting during flattened JSON tokenization.
+  // The default false value preserves exact Cartesian-product semantics.
+  bool disable_cross_array_unnest = 17;
+  /* Maximum flattened sub-documents emitted for one JSON row. Absence means
+   * unlimited.
+   */
+  optional uint64 max_sub_docs_per_row = 18;
+  /* Action taken when max_sub_docs_per_row is exceeded. FAIL aborts index
+   * ingestion; SKIP_ROW omits the source row from the index.
+   */
+  MaxSubDocsPerRowExceedAction max_sub_docs_per_row_exceed_action = 19;
 }

python/python/lance/dataset.py

@@ -3720,6 +3720,20 @@ def create_scalar_index(
             ``[1, num_compute_cpus]``. If unset, Lance uses ``num_compute_cpus``
             workers unless ``LANCE_FTS_NUM_SHARDS`` is set. This parameter is
             only used for the current build and is not persisted with the index.
+        disable_cross_array_unnest: bool, default False
+            This is for the ``INVERTED`` index on JSON columns. If True, flattened
+            JSON tokenization indexes sibling arrays independently instead of
+            producing their Cartesian product. This reduces index build memory for
+            records with multiple arrays but can sacrifice result accuracy for
+            queries that constrain values across those arrays.
+        max_sub_docs_per_row: int, optional
+            This is for the ``INVERTED`` index on JSON columns. Maximum number of
+            flattened sub-documents one source row may produce. If unset, the
+            number is unlimited.
+        max_sub_docs_per_row_exceed_action: str, default "fail"
+            Action when ``max_sub_docs_per_row`` is exceeded. ``"fail"`` aborts
+            index ingestion with tuning guidance. ``"skip_row"`` omits the source
+            row from the index and continues.
         base_tokenizer: str, default "simple"
             This is for the ``INVERTED`` index. The base tokenizer to use. The
             value can be:

python/python/tests/test_scalar_index.py

@@ -6032,7 +6032,14 @@ def test_json_inverted_match_query(tmp_path):
         stem=True,
         lower_case=True,
         remove_stop_words=True,
+        disable_cross_array_unnest=True,
+        max_sub_docs_per_row=128,
+        max_sub_docs_per_row_exceed_action="skip_row",
     )
+    details = dataset.describe_indices()[0].details
+    assert details["disable_cross_array_unnest"] is True
+    assert details["max_sub_docs_per_row"] == 128
+    assert details["max_sub_docs_per_row_exceed_action"] == "skip_row"
 
     # Test match query with token exceeding max_token_length
     results = dataset.to_table(
@@ -6048,7 +6055,7 @@ def test_json_inverted_match_query(tmp_path):
 
     # Test language match
     results = dataset.to_table(
-        full_text_query=MatchQuery("Language,str,english", "json_col")
+        full_text_query=MatchQuery("Language[*],str,english", "json_col")
     )
     assert results.num_rows == 3
 

python/src/dataset.rs

@@ -2649,6 +2649,9 @@ impl Dataset {
                         "split_on_numerics",
                         "preserve_original",
                         "index_operators",
+                        "disable_cross_array_unnest",
+                        "max_sub_docs_per_row",
+                        "max_sub_docs_per_row_exceed_action",
                         "memory_limit",
                         "num_workers",
                         "format_version",
@@ -2682,6 +2685,24 @@ impl Dataset {
                             .block_size(block_size.extract()?)
                             .map_err(|e| PyValueError::new_err(e.to_string()))?;
                     }
+                    if let Some(disable_cross_array_unnest) =
+                        kwargs.get_item("disable_cross_array_unnest")?
+                    {
+                        params = params
+                            .disable_cross_array_unnest(disable_cross_array_unnest.extract()?);
+                    }
+                    if let Some(max_sub_docs_per_row) = kwargs.get_item("max_sub_docs_per_row")? {
+                        params = params
+                            .max_sub_docs_per_row(max_sub_docs_per_row.extract()?)
+                            .map_err(|err| PyValueError::new_err(err.to_string()))?;
+                    }
+                    if let Some(action) = kwargs.get_item("max_sub_docs_per_row_exceed_action")? {
+                        let action: String = action.extract()?;
+                        params =
+                            params.max_sub_docs_per_row_exceed_action(action.parse().map_err(
+                                |err: lance_core::Error| PyValueError::new_err(err.to_string()),
+                            )?);
+                    }
                     if let Some(memory_limit) = kwargs.get_item("memory_limit")? {
                         params = params.memory_limit_mb(memory_limit.extract()?);
                     }

rust/lance-index/src/scalar.rs

@@ -47,7 +47,7 @@ pub mod seed;
 pub mod zoned;
 pub mod zonemap;
 
-pub use inverted::tokenizer::InvertedIndexParams;
+pub use inverted::tokenizer::{InvertedIndexParams, MaxSubDocsPerRowExceedAction};
 
 /// Convert a `Vec<`[`lance_index_core::scalar::IndexFile`]`>` to a
 /// `Vec<`[`lance_table::format::IndexFile`]`>`.

rust/lance-index/src/scalar/inverted.rs

@@ -40,6 +40,32 @@ pub use tokenizer::*;
 
 use crate::scalar::inverted::query::{FtsSearchParams, Tokens, uses_fuzzy_expansion};
 
+pub(crate) fn collapse_scored_rows(
+    rows: impl IntoIterator<Item = (u64, f32)>,
+    limit: usize,
+) -> Vec<(u64, f32)> {
+    let mut scores_by_row_id = HashMap::new();
+    for (row_id, score) in rows {
+        scores_by_row_id
+            .entry(row_id)
+            .and_modify(|existing| {
+                if score > *existing {
+                    *existing = score;
+                }
+            })
+            .or_insert(score);
+    }
+
+    let mut rows = scores_by_row_id.into_iter().collect::<Vec<_>>();
+    rows.sort_unstable_by(|(left_id, left_score), (right_id, right_score)| {
+        right_score
+            .total_cmp(left_score)
+            .then_with(|| left_id.cmp(right_id))
+    });
+    rows.truncate(limit);
+    rows
+}
+
 /// Canonical token vocabulary and BM25 statistics for one indexed query leaf.
 ///
 /// Keeping these values together prevents a search path from expanding one
@@ -428,11 +454,11 @@ impl InvertedIndexPlugin {
         params.validate_format_version()?;
         let format_version = params.resolved_format_version();
         let is_element_document = params.get_document_granularity().is_list_element();
-        let details = pbold::InvertedIndexDetails::try_from(&params)?;
         let mut inverted_index =
             InvertedIndexBuilder::new_with_fragment_mask(params, fragment_mask)
                 .with_progress(progress);
         let files = inverted_index.update(data, index_store, None).await?;
+        let details = pbold::InvertedIndexDetails::try_from(inverted_index.params())?;
         Ok(CreatedIndex {
             index_details: prost_types::Any::from_msg(&details).unwrap(),
             index_version: if is_element_document {