Merge remote-tracking branch 'origin/main' into gatekeeper/fix-1525-1

# Conflicts: # nodejs/__test__/arrow.test.ts # nodejs/lancedb/sanitize.ts
2026-08-18 12:08:35 +00:00 · 2026-08-08 19:39:27 +00:00
parent 62d4a345ee 36054be576
commit 73f3fbd0b7
2 changed files with 290 additions and 33 deletions
@@ -6,7 +6,9 @@ import * as arrow17 from "apache-arrow-17";
 import * as arrow18 from "apache-arrow-18";

 import {
+  Vector as CurrentVector,
  convertToTable,
+  tableFromIPC as currentTableFromIPC,
  fromBufferToRecordBatch,
  fromDataToBuffer,
  fromRecordBatchToBuffer,
@@ -19,6 +21,7 @@ import {
  FunctionOptions,
 } from "../lancedb/embedding/embedding_function";
 import { EmbeddingFunctionConfig } from "../lancedb/embedding/registry";
+import { sanitizeTable } from "../lancedb/sanitize";

 // biome-ignore lint/suspicious/noExplicitAny: skip
 function sampleRecords(): Array<Record<string, any>> {
@@ -65,7 +68,11 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
      tableFromIPC,
      DataType,
      Dictionary,
+      RecordBatch: ArrowRecordBatch,
+      Table: ArrowTable,
      Uint8: ArrowUint8,
+      makeData: arrowMakeData,
+      vectorFromArray,
      // biome-ignore lint/suspicious/noExplicitAny: <explanation>
    } = <any>arrow;
    type Schema = ApacheArrow["Schema"];
@@ -1070,6 +1077,114 @@ describe.each([arrow15, arrow16, arrow17, arrow18])(
        expect(actual.getChild("text")?.toJSON()).toEqual(["foo", "bar"]);
      });

+      it("preserves a dictionary shared by multiple fields", async function () {
+        const values = ["alpha", "beta", "alpha"];
+        const dictionaryVector = vectorFromArray(values);
+        const batch = new ArrowRecordBatch({
+          first: dictionaryVector.data[0],
+          second: dictionaryVector.data[0],
+        });
+        const table = new ArrowTable([batch]);
+
+        const sanitized = sanitizeTable(table);
+        expect([...sanitized.getChild("first")!]).toEqual(values);
+        expect([...sanitized.getChild("second")!]).toEqual(values);
+        const firstType = sanitized.schema.fields[0].type as {
+          dictionary: unknown;
+        };
+        const secondType = sanitized.schema.fields[1].type as {
+          dictionary: unknown;
+        };
+        expect(secondType.dictionary).toBe(firstType.dictionary);
+        expect(sanitized.batches[0].data.children[1].dictionary).toBe(
+          sanitized.batches[0].data.children[0].dictionary,
+        );
+
+        const buf = await fromDataToBuffer(table);
+        const actual = currentTableFromIPC(buf);
+        expect([...actual.getChild("first")!]).toEqual(values);
+        expect([...actual.getChild("second")!]).toEqual(values);
+      });
+
+      it("preserves shared dictionary data from another Arrow version", async function () {
+        const values = ["alpha", "beta", "alpha"];
+        const dictionaryVector = vectorFromArray(values);
+        const firstBatch = new ArrowRecordBatch({
+          label: dictionaryVector.slice(0, 2).data[0],
+        });
+        const secondBatch = new ArrowRecordBatch({
+          label: dictionaryVector.slice(2).data[0],
+        });
+        const table = new ArrowTable([firstBatch, secondBatch]);
+
+        const sanitized = sanitizeTable(table);
+        expect([...sanitized.getChild("label")!]).toEqual(values);
+
+        const dictionaries = sanitized.batches.map(
+          (batch) => batch.data.children[0].dictionary,
+        );
+        expect(dictionaries[0]).toBeInstanceOf(CurrentVector);
+        expect(dictionaries[1]).toBe(dictionaries[0]);
+
+        const buf = await fromDataToBuffer(table);
+        const actual = currentTableFromIPC(buf);
+        expect([...actual.getChild("label")!]).toEqual(values);
+      });
+
+      it("preserves shared chunks in growing dictionaries", async function () {
+        const type = new Dictionary(new Utf8(), new Int32(), 42, false);
+        const firstDictionary = vectorFromArray(["alpha", "beta"], new Utf8());
+        const secondDictionary = firstDictionary.concat(
+          vectorFromArray(["gamma"], new Utf8()),
+        );
+        const firstData = arrowMakeData({
+          type,
+          data: Int32Array.from([0, 1]),
+          dictionary: firstDictionary,
+        });
+        const secondData = arrowMakeData({
+          type,
+          data: Int32Array.from([2]),
+          dictionary: secondDictionary,
+        });
+        const table = new ArrowTable([
+          new ArrowRecordBatch({ label: firstData }),
+          new ArrowRecordBatch({ label: secondData }),
+        ]);
+
+        const sanitized = sanitizeTable(table);
+        const expected = ["alpha", "beta", "gamma"];
+        expect([...sanitized.getChild("label")!]).toEqual(expected);
+        const firstLocalDictionary =
+          sanitized.batches[0].data.children[0].dictionary!;
+        const secondLocalDictionary =
+          sanitized.batches[1].data.children[0].dictionary!;
+        expect(secondLocalDictionary.data[0]).toBe(
+          firstLocalDictionary.data[0],
+        );
+
+        const buf = await fromTableToBuffer(sanitized);
+        const actual = currentTableFromIPC(buf);
+        expect([...actual.getChild("label")!]).toEqual(expected);
+      });
+
+      it("can serialize list data from another Arrow version", async function () {
+        const values = [["anime", "action"], [], null];
+        const vector = vectorFromArray(
+          values,
+          new List(new Field("item", new Utf8(), true)),
+        );
+        const table = new ArrowTable({ tags: vector });
+
+        const buf = await fromDataToBuffer(table);
+        const actual = currentTableFromIPC(buf);
+        const actualTags = actual.getChild("tags");
+
+        expect(actualTags?.get(0)?.toJSON()).toEqual(values[0]);
+        expect(actualTags?.get(1)?.toJSON()).toEqual(values[1]);
+        expect(actualTags?.get(2)).toBeNull();
+      });
+
      it("can still import data", async function () {
        const schema = new arrow15.Schema([
          new arrow15.Field("id", new arrow15.Int32()),
@@ -9,7 +9,7 @@
 // comes from the exact same library instance.  This is not always the case
 // and so we must sanitize the input to ensure that it is compatible.

-import { BufferType, Data } from "apache-arrow";
+import { BufferType, Data, Vector } from "apache-arrow";
 import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
 import {
  Binary,
@@ -72,9 +72,22 @@ import {
  Uint64,
  Union,
  Utf8,
-  Vector,
 } from "./arrow";

+type SanitizationContext = {
+  types: WeakMap<object, DataType>;
+  vectors: WeakMap<object, Vector>;
+  data: WeakMap<object, Data<DataType>>;
+};
+
+function createSanitizationContext(): SanitizationContext {
+  return {
+    types: new WeakMap(),
+    vectors: new WeakMap(),
+    data: new WeakMap(),
+  };
+}
+
 export function sanitizeMetadata(
  metadataLike?: unknown,
 ): Map<string, string> | undefined {
@@ -187,6 +200,13 @@ export function sanitizeInterval(typeLike: object) {
 }

 export function sanitizeList(typeLike: object) {
+  return sanitizeListWithContext(typeLike, createSanitizationContext());
+}
+
+function sanitizeListWithContext(
+  typeLike: object,
+  context: SanitizationContext,
+) {
  if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
    throw Error(
      "Expected a List type to have an array-like `children` property",
@@ -195,19 +215,35 @@ export function sanitizeList(typeLike: object) {
  if (typeLike.children.length !== 1) {
    throw Error("Expected a List type to have exactly one child");
  }
-  return new List(sanitizeField(typeLike.children[0]));
+  return new List(sanitizeFieldWithContext(typeLike.children[0], context));
 }

 export function sanitizeStruct(typeLike: object) {
+  return sanitizeStructWithContext(typeLike, createSanitizationContext());
+}
+
+function sanitizeStructWithContext(
+  typeLike: object,
+  context: SanitizationContext,
+) {
  if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
    throw Error(
      "Expected a Struct type to have an array-like `children` property",
    );
  }
-  return new Struct(typeLike.children.map((child) => sanitizeField(child)));
+  return new Struct(
+    typeLike.children.map((child) => sanitizeFieldWithContext(child, context)),
+  );
 }

 export function sanitizeUnion(typeLike: object) {
+  return sanitizeUnionWithContext(typeLike, createSanitizationContext());
+}
+
+function sanitizeUnionWithContext(
+  typeLike: object,
+  context: SanitizationContext,
+) {
  if (
    !("typeIds" in typeLike) ||
    !("mode" in typeLike) ||
@@ -227,7 +263,7 @@ export function sanitizeUnion(typeLike: object) {
    typeLike.mode,
    // biome-ignore lint/suspicious/noExplicitAny: skip
    typeLike.typeIds as any,
-    typeLike.children.map((child) => sanitizeField(child)),
+    typeLike.children.map((child) => sanitizeFieldWithContext(child, context)),
  );
 }

@@ -235,6 +271,19 @@ export function sanitizeTypedUnion(
  typeLike: object,
  // eslint-disable-next-line @typescript-eslint/naming-convention
  UnionType: typeof DenseUnion | typeof SparseUnion,
+) {
+  return sanitizeTypedUnionWithContext(
+    typeLike,
+    UnionType,
+    createSanitizationContext(),
+  );
+}
+
+function sanitizeTypedUnionWithContext(
+  typeLike: object,
+  // eslint-disable-next-line @typescript-eslint/naming-convention
+  UnionType: typeof DenseUnion | typeof SparseUnion,
+  context: SanitizationContext,
 ) {
  if (!("typeIds" in typeLike)) {
    throw Error(
@@ -249,7 +298,7 @@ export function sanitizeTypedUnion(

  return new UnionType(
    typeLike.typeIds as Int32Array | number[],
-    typeLike.children.map((child) => sanitizeField(child)),
+    typeLike.children.map((child) => sanitizeFieldWithContext(child, context)),
  );
 }

@@ -263,6 +312,16 @@ export function sanitizeFixedSizeBinary(typeLike: object) {
 }

 export function sanitizeFixedSizeList(typeLike: object) {
+  return sanitizeFixedSizeListWithContext(
+    typeLike,
+    createSanitizationContext(),
+  );
+}
+
+function sanitizeFixedSizeListWithContext(
+  typeLike: object,
+  context: SanitizationContext,
+) {
  if (!("listSize" in typeLike) || typeof typeLike.listSize !== "number") {
    throw Error("Expected a FixedSizeList type to have a `listSize` property");
  }
@@ -276,11 +335,18 @@ export function sanitizeFixedSizeList(typeLike: object) {
  }
  return new FixedSizeList(
    typeLike.listSize,
-    sanitizeField(typeLike.children[0]),
+    sanitizeFieldWithContext(typeLike.children[0], context),
  );
 }

 export function sanitizeMap(typeLike: object) {
+  return sanitizeMapWithContext(typeLike, createSanitizationContext());
+}
+
+function sanitizeMapWithContext(
+  typeLike: object,
+  context: SanitizationContext,
+) {
  if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
    throw Error(
      "Expected a Map type to have an array-like `children` property",
@@ -293,7 +359,10 @@ export function sanitizeMap(typeLike: object) {
    throw Error("Expected a Map type to have exactly one child");
  }

-  return new Map_(sanitizeField(typeLike.children[0]), typeLike.keysSorted);
+  return new Map_(
+    sanitizeFieldWithContext(typeLike.children[0], context),
+    typeLike.keysSorted,
+  );
 }

 export function sanitizeDuration(typeLike: object) {
@@ -304,6 +373,13 @@ export function sanitizeDuration(typeLike: object) {
 }

 export function sanitizeDictionary(typeLike: object) {
+  return sanitizeDictionaryWithContext(typeLike, createSanitizationContext());
+}
+
+function sanitizeDictionaryWithContext(
+  typeLike: object,
+  context: SanitizationContext,
+) {
  if (!("id" in typeLike) || typeof typeLike.id !== "number") {
    throw Error("Expected a Dictionary type to have an `id` property");
  }
@@ -317,8 +393,8 @@ export function sanitizeDictionary(typeLike: object) {
    throw Error("Expected a Dictionary type to have an `isOrdered` property");
  }
  return new Dictionary(
-    sanitizeType(typeLike.dictionary),
-    sanitizeType(typeLike.indices) as TKeys,
+    sanitizeTypeWithContext(typeLike.dictionary, context),
+    sanitizeTypeWithContext(typeLike.indices, context) as TKeys,
    typeLike.id,
    typeLike.isOrdered,
  );
@@ -326,12 +402,23 @@ export function sanitizeDictionary(typeLike: object) {

 // biome-ignore lint/suspicious/noExplicitAny: skip
 export function sanitizeType(typeLike: unknown): DataType<any> {
+  return sanitizeTypeWithContext(typeLike, createSanitizationContext());
+}
+
+function sanitizeTypeWithContext(
+  typeLike: unknown,
+  context: SanitizationContext,
+): DataType {
  if (typeof typeLike === "string") {
    return dataTypeFromName(typeLike);
  }
  if (typeof typeLike !== "object" || typeLike === null) {
    throw Error("Expected a Type but object was null/undefined");
  }
+  const cached = context.types.get(typeLike);
+  if (cached !== undefined) {
+    return cached;
+  }
  if (
    !("typeId" in typeLike) ||
    !(
@@ -350,6 +437,16 @@ export function sanitizeType(typeLike: unknown): DataType<any> {
    throw Error("Type's typeId property was not a function or number");
  }

+  const type = sanitizeTypeById(typeLike, typeId, context);
+  context.types.set(typeLike, type);
+  return type;
+}
+
+function sanitizeTypeById(
+  typeLike: object,
+  typeId: Type,
+  context: SanitizationContext,
+): DataType {
  switch (typeId) {
    case Type.NONE:
      throw Error("Received a Type with a typeId of NONE");
@@ -376,21 +473,21 @@ export function sanitizeType(typeLike: unknown): DataType<any> {
    case Type.Interval:
      return sanitizeInterval(typeLike);
    case Type.List:
-      return sanitizeList(typeLike);
+      return sanitizeListWithContext(typeLike, context);
    case Type.Struct:
-      return sanitizeStruct(typeLike);
+      return sanitizeStructWithContext(typeLike, context);
    case Type.Union:
-      return sanitizeUnion(typeLike);
+      return sanitizeUnionWithContext(typeLike, context);
    case Type.FixedSizeBinary:
      return sanitizeFixedSizeBinary(typeLike);
    case Type.FixedSizeList:
-      return sanitizeFixedSizeList(typeLike);
+      return sanitizeFixedSizeListWithContext(typeLike, context);
    case Type.Map:
-      return sanitizeMap(typeLike);
+      return sanitizeMapWithContext(typeLike, context);
    case Type.Duration:
      return sanitizeDuration(typeLike);
    case Type.Dictionary:
-      return sanitizeDictionary(typeLike);
+      return sanitizeDictionaryWithContext(typeLike, context);
    case Type.Int8:
      return new Int8();
    case Type.Int16:
@@ -434,9 +531,9 @@ export function sanitizeType(typeLike: unknown): DataType<any> {
    case Type.TimestampSecond:
      return sanitizeTypedTimestamp(typeLike, TimestampSecond);
    case Type.DenseUnion:
-      return sanitizeTypedUnion(typeLike, DenseUnion);
+      return sanitizeTypedUnionWithContext(typeLike, DenseUnion, context);
    case Type.SparseUnion:
-      return sanitizeTypedUnion(typeLike, SparseUnion);
+      return sanitizeTypedUnionWithContext(typeLike, SparseUnion, context);
    case Type.IntervalDayTime:
      return new IntervalDayTime();
    case Type.IntervalYearMonth:
@@ -455,6 +552,13 @@ export function sanitizeType(typeLike: unknown): DataType<any> {
 }

 export function sanitizeField(fieldLike: unknown): Field {
+  return sanitizeFieldWithContext(fieldLike, createSanitizationContext());
+}
+
+function sanitizeFieldWithContext(
+  fieldLike: unknown,
+  context: SanitizationContext,
+): Field {
  if (fieldLike instanceof Field) {
    return fieldLike;
  }
@@ -472,7 +576,7 @@ export function sanitizeField(fieldLike: unknown): Field {
  }
  let type: DataType;
  try {
-    type = sanitizeType(fieldLike.type);
+    type = sanitizeTypeWithContext(fieldLike.type, context);
  } catch (error: unknown) {
    throw Error(
      `Unable to sanitize type for field: ${fieldLike.name} due to error: ${error}`,
@@ -502,6 +606,13 @@ export function sanitizeField(fieldLike: unknown): Field {
 * than lancedb is using.
 */
 export function sanitizeSchema(schemaLike: SchemaLike): Schema {
+  return sanitizeSchemaWithContext(schemaLike, createSanitizationContext());
+}
+
+function sanitizeSchemaWithContext(
+  schemaLike: SchemaLike,
+  context: SanitizationContext,
+): Schema {
  if (schemaLike instanceof Schema) {
    return schemaLike;
  }
@@ -523,7 +634,7 @@ export function sanitizeSchema(schemaLike: SchemaLike): Schema {
    );
  }
  const sanitizedFields = schemaLike.fields.map((field) =>
-    sanitizeField(field),
+    sanitizeFieldWithContext(field, context),
  );
  return new Schema(sanitizedFields, metadata);
 }
@@ -545,13 +656,18 @@ export function sanitizeTable(tableLike: TableLike): Table {
      "The table passed in does not appear to be a table (no 'columns' property)",
    );
  }
-  const schema = sanitizeSchema(tableLike.schema);
-
-  const batches = tableLike.batches.map(sanitizeRecordBatch);
+  const context = createSanitizationContext();
+  const schema = sanitizeSchemaWithContext(tableLike.schema, context);
+  const batches = tableLike.batches.map((batch) =>
+    sanitizeRecordBatch(batch, context),
+  );
  return new Table(schema, batches);
 }

-function sanitizeRecordBatch(batchLike: RecordBatchLike): RecordBatch {
+function sanitizeRecordBatch(
+  batchLike: RecordBatchLike,
+  context: SanitizationContext,
+): RecordBatch {
  if (batchLike instanceof RecordBatch) {
    return batchLike;
  }
@@ -568,19 +684,43 @@ function sanitizeRecordBatch(batchLike: RecordBatchLike): RecordBatch {
      "The record batch passed in does not appear to be a record batch (no 'data' property)",
    );
  }
-  const schema = sanitizeSchema(batchLike.schema);
-  const data = sanitizeData(batchLike.data) as Data<Struct>;
+  const schema = sanitizeSchemaWithContext(batchLike.schema, context);
+  const data = sanitizeData(batchLike.data, context) as Data<Struct>;
  return new RecordBatch(schema, data);
 }
-function sanitizeData(dataLike: DataLike): Data {
+
+type DictionaryVectorLike = {
+  data: readonly DataLike[];
+};
+
+type DictionaryDataLike = DataLike & {
+  dictionary?: DictionaryVectorLike;
+};
+
+function sanitizeData(
+  dataLike: DataLike,
+  context: SanitizationContext,
+): Data<DataType> {
  if (dataLike instanceof Data) {
    return dataLike;
  }
-  const dictionary = dataLike.dictionary
-    ? new Vector(dataLike.dictionary.data.map(sanitizeData))
-    : undefined;
-  return new Data(
-    sanitizeType(dataLike.type),
+  const cachedData = context.data.get(dataLike);
+  if (cachedData !== undefined) {
+    return cachedData;
+  }
+  const dictionaryLike = (dataLike as DictionaryDataLike).dictionary;
+  let dictionary: Vector | undefined;
+  if (dictionaryLike !== undefined) {
+    dictionary = context.vectors.get(dictionaryLike);
+    if (dictionary === undefined) {
+      dictionary = new Vector(
+        dictionaryLike.data.map((data) => sanitizeData(data, context)),
+      );
+      context.vectors.set(dictionaryLike, dictionary);
+    }
+  }
+  const data = new Data(
+    sanitizeTypeWithContext(dataLike.type, context),
    dataLike.offset,
    dataLike.length,
    dataLike.nullCount,
@@ -590,9 +730,11 @@ function sanitizeData(dataLike: DataLike): Data {
      [BufferType.VALIDITY]: dataLike.nullBitmap,
      [BufferType.TYPE]: dataLike.typeIds,
    },
-    dataLike.children.map(sanitizeData),
+    dataLike.children.map((child) => sanitizeData(child, context)),
    dictionary,
  );
+  context.data.set(dataLike, data);
+  return data;
 }

 const constructorsByTypeName = {