lancedb/nodejs/lancedb/sanitize.ts

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors

// The utilities in this file help sanitize data from the user's arrow
// library into the types expected by vectordb's arrow library.  Node
// generally allows for mulitple versions of the same library (and sometimes
// even multiple copies of the same version) to be installed at the same
// time.  However, arrow-js uses instanceof which expected that the input
// 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, Vector } from "apache-arrow";
import type { IntBitWidth, TKeys, TimeBitWidth } from "apache-arrow/type";
import {
  Binary,
  Bool,
  DataLike,
  DataType,
  DateDay,
  DateMillisecond,
  type DateUnit,
  Date_,
  Decimal,
  DenseUnion,
  Dictionary,
  Duration,
  DurationMicrosecond,
  DurationMillisecond,
  DurationNanosecond,
  DurationSecond,
  Field,
  FixedSizeBinary,
  FixedSizeList,
  Float,
  Float16,
  Float32,
  Float64,
  Int,
  Int8,
  Int16,
  Int32,
  Int64,
  Interval,
  IntervalDayTime,
  IntervalYearMonth,
  List,
  Map_,
  Null,
  type Precision,
  RecordBatch,
  RecordBatchLike,
  Schema,
  SchemaLike,
  SparseUnion,
  Struct,
  Table,
  TableLike,
  Time,
  TimeMicrosecond,
  TimeMillisecond,
  TimeNanosecond,
  TimeSecond,
  Timestamp,
  TimestampMicrosecond,
  TimestampMillisecond,
  TimestampNanosecond,
  TimestampSecond,
  Type,
  Uint8,
  Uint16,
  Uint32,
  Uint64,
  Union,
  Utf8,
} 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 {
  if (metadataLike === undefined || metadataLike === null) {
    return undefined;
  }

  let entries: IterableIterator<[unknown, unknown]>;
  try {
    entries = Map.prototype.entries.call(metadataLike);
  } catch {
    throw Error("Expected metadata, if present, to be a Map<string, string>");
  }

  const metadata = new Map<string, string>();
  for (const [key, value] of entries) {
    if (typeof key !== "string" || typeof value !== "string") {
      throw Error(
        "Expected metadata, if present, to be a Map<string, string> but it had non-string keys or values",
      );
    }
    metadata.set(key, value);
  }
  return metadata;
}

export function sanitizeInt(typeLike: object) {
  if (
    !("bitWidth" in typeLike) ||
    typeof typeLike.bitWidth !== "number" ||
    !("isSigned" in typeLike) ||
    typeof typeLike.isSigned !== "boolean"
  ) {
    throw Error(
      "Expected an Int Type to have a `bitWidth` and `isSigned` property",
    );
  }
  return new Int(typeLike.isSigned, typeLike.bitWidth as IntBitWidth);
}

export function sanitizeFloat(typeLike: object) {
  if (!("precision" in typeLike) || typeof typeLike.precision !== "number") {
    throw Error("Expected a Float Type to have a `precision` property");
  }
  return new Float(typeLike.precision as Precision);
}

export function sanitizeDecimal(typeLike: object) {
  if (
    !("scale" in typeLike) ||
    typeof typeLike.scale !== "number" ||
    !("precision" in typeLike) ||
    typeof typeLike.precision !== "number" ||
    !("bitWidth" in typeLike) ||
    typeof typeLike.bitWidth !== "number"
  ) {
    throw Error(
      "Expected a Decimal Type to have `scale`, `precision`, and `bitWidth` properties",
    );
  }
  return new Decimal(typeLike.scale, typeLike.precision, typeLike.bitWidth);
}

export function sanitizeDate(typeLike: object) {
  if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
    throw Error("Expected a Date type to have a `unit` property");
  }
  return new Date_(typeLike.unit as DateUnit);
}

export function sanitizeTime(typeLike: object) {
  if (
    !("unit" in typeLike) ||
    typeof typeLike.unit !== "number" ||
    !("bitWidth" in typeLike) ||
    typeof typeLike.bitWidth !== "number"
  ) {
    throw Error(
      "Expected a Time type to have `unit` and `bitWidth` properties",
    );
  }
  return new Time(typeLike.unit, typeLike.bitWidth as TimeBitWidth);
}

export function sanitizeTimestamp(typeLike: object) {
  if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
    throw Error("Expected a Timestamp type to have a `unit` property");
  }
  let timezone = null;
  if ("timezone" in typeLike && typeof typeLike.timezone === "string") {
    timezone = typeLike.timezone;
  }
  return new Timestamp(typeLike.unit, timezone);
}

export function sanitizeTypedTimestamp(
  typeLike: object,
  // eslint-disable-next-line @typescript-eslint/naming-convention
  Datatype:
    | typeof TimestampNanosecond
    | typeof TimestampMicrosecond
    | typeof TimestampMillisecond
    | typeof TimestampSecond,
) {
  let timezone = null;
  if ("timezone" in typeLike && typeof typeLike.timezone === "string") {
    timezone = typeLike.timezone;
  }
  return new Datatype(timezone);
}

export function sanitizeInterval(typeLike: object) {
  if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
    throw Error("Expected an Interval type to have a `unit` property");
  }
  return new Interval(typeLike.unit);
}

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",
    );
  }
  if (typeLike.children.length !== 1) {
    throw Error("Expected a List type to have exactly one child");
  }
  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) => 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) ||
    typeof typeLike.mode !== "number"
  ) {
    throw Error(
      "Expected a Union type to have `typeIds` and `mode` properties",
    );
  }
  if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
    throw Error(
      "Expected a Union type to have an array-like `children` property",
    );
  }

  return new Union(
    typeLike.mode,
    // biome-ignore lint/suspicious/noExplicitAny: skip
    typeLike.typeIds as any,
    typeLike.children.map((child) => sanitizeFieldWithContext(child, context)),
  );
}

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(
      "Expected a DenseUnion/SparseUnion type to have a `typeIds` property",
    );
  }
  if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
    throw Error(
      "Expected a DenseUnion/SparseUnion type to have an array-like `children` property",
    );
  }

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

export function sanitizeFixedSizeBinary(typeLike: object) {
  if (!("byteWidth" in typeLike) || typeof typeLike.byteWidth !== "number") {
    throw Error(
      "Expected a FixedSizeBinary type to have a `byteWidth` property",
    );
  }
  return new FixedSizeBinary(typeLike.byteWidth);
}

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");
  }
  if (!("children" in typeLike) || !Array.isArray(typeLike.children)) {
    throw Error(
      "Expected a FixedSizeList type to have an array-like `children` property",
    );
  }
  if (typeLike.children.length !== 1) {
    throw Error("Expected a FixedSizeList type to have exactly one child");
  }
  return new FixedSizeList(
    typeLike.listSize,
    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",
    );
  }
  if (!("keysSorted" in typeLike) || typeof typeLike.keysSorted !== "boolean") {
    throw Error("Expected a Map type to have a `keysSorted` property");
  }
  if (typeLike.children.length !== 1) {
    throw Error("Expected a Map type to have exactly one child");
  }

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

export function sanitizeDuration(typeLike: object) {
  if (!("unit" in typeLike) || typeof typeLike.unit !== "number") {
    throw Error("Expected a Duration type to have a `unit` property");
  }
  return new Duration(typeLike.unit);
}

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");
  }
  if (!("indices" in typeLike) || typeof typeLike.indices !== "object") {
    throw Error("Expected a Dictionary type to have an `indices` property");
  }
  if (!("dictionary" in typeLike) || typeof typeLike.dictionary !== "object") {
    throw Error("Expected a Dictionary type to have an `dictionary` property");
  }
  if (!("isOrdered" in typeLike) || typeof typeLike.isOrdered !== "boolean") {
    throw Error("Expected a Dictionary type to have an `isOrdered` property");
  }
  return new Dictionary(
    sanitizeTypeWithContext(typeLike.dictionary, context),
    sanitizeTypeWithContext(typeLike.indices, context) as TKeys,
    typeLike.id,
    typeLike.isOrdered,
  );
}

// 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) ||
    !(
      typeof typeLike.typeId !== "function" ||
      typeof typeLike.typeId !== "number"
    )
  ) {
    throw Error("Expected a Type to have a typeId property");
  }
  let typeId: Type;
  if (typeof typeLike.typeId === "function") {
    typeId = (typeLike.typeId as () => unknown)() as Type;
  } else if (typeof typeLike.typeId === "number") {
    typeId = typeLike.typeId as Type;
  } else {
    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");
    case Type.Null:
      return new Null();
    case Type.Int:
      return sanitizeInt(typeLike);
    case Type.Float:
      return sanitizeFloat(typeLike);
    case Type.Binary:
      return new Binary();
    case Type.Utf8:
      return new Utf8();
    case Type.Bool:
      return new Bool();
    case Type.Decimal:
      return sanitizeDecimal(typeLike);
    case Type.Date:
      return sanitizeDate(typeLike);
    case Type.Time:
      return sanitizeTime(typeLike);
    case Type.Timestamp:
      return sanitizeTimestamp(typeLike);
    case Type.Interval:
      return sanitizeInterval(typeLike);
    case Type.List:
      return sanitizeListWithContext(typeLike, context);
    case Type.Struct:
      return sanitizeStructWithContext(typeLike, context);
    case Type.Union:
      return sanitizeUnionWithContext(typeLike, context);
    case Type.FixedSizeBinary:
      return sanitizeFixedSizeBinary(typeLike);
    case Type.FixedSizeList:
      return sanitizeFixedSizeListWithContext(typeLike, context);
    case Type.Map:
      return sanitizeMapWithContext(typeLike, context);
    case Type.Duration:
      return sanitizeDuration(typeLike);
    case Type.Dictionary:
      return sanitizeDictionaryWithContext(typeLike, context);
    case Type.Int8:
      return new Int8();
    case Type.Int16:
      return new Int16();
    case Type.Int32:
      return new Int32();
    case Type.Int64:
      return new Int64();
    case Type.Uint8:
      return new Uint8();
    case Type.Uint16:
      return new Uint16();
    case Type.Uint32:
      return new Uint32();
    case Type.Uint64:
      return new Uint64();
    case Type.Float16:
      return new Float16();
    case Type.Float32:
      return new Float32();
    case Type.Float64:
      return new Float64();
    case Type.DateMillisecond:
      return new DateMillisecond();
    case Type.DateDay:
      return new DateDay();
    case Type.TimeNanosecond:
      return new TimeNanosecond();
    case Type.TimeMicrosecond:
      return new TimeMicrosecond();
    case Type.TimeMillisecond:
      return new TimeMillisecond();
    case Type.TimeSecond:
      return new TimeSecond();
    case Type.TimestampNanosecond:
      return sanitizeTypedTimestamp(typeLike, TimestampNanosecond);
    case Type.TimestampMicrosecond:
      return sanitizeTypedTimestamp(typeLike, TimestampMicrosecond);
    case Type.TimestampMillisecond:
      return sanitizeTypedTimestamp(typeLike, TimestampMillisecond);
    case Type.TimestampSecond:
      return sanitizeTypedTimestamp(typeLike, TimestampSecond);
    case Type.DenseUnion:
      return sanitizeTypedUnionWithContext(typeLike, DenseUnion, context);
    case Type.SparseUnion:
      return sanitizeTypedUnionWithContext(typeLike, SparseUnion, context);
    case Type.IntervalDayTime:
      return new IntervalDayTime();
    case Type.IntervalYearMonth:
      return new IntervalYearMonth();
    case Type.DurationNanosecond:
      return new DurationNanosecond();
    case Type.DurationMicrosecond:
      return new DurationMicrosecond();
    case Type.DurationMillisecond:
      return new DurationMillisecond();
    case Type.DurationSecond:
      return new DurationSecond();
    default:
      throw new Error("Unrecognized type id in schema: " + typeId);
  }
}

export function sanitizeField(fieldLike: unknown): Field {
  return sanitizeFieldWithContext(fieldLike, createSanitizationContext());
}

function sanitizeFieldWithContext(
  fieldLike: unknown,
  context: SanitizationContext,
): Field {
  if (fieldLike instanceof Field) {
    return fieldLike;
  }
  if (typeof fieldLike !== "object" || fieldLike === null) {
    throw Error("Expected a Field but object was null/undefined");
  }
  if (
    !("type" in fieldLike) ||
    !("name" in fieldLike) ||
    !("nullable" in fieldLike)
  ) {
    throw Error(
      "The field passed in is missing a `type`/`name`/`nullable` property",
    );
  }
  let type: DataType;
  try {
    type = sanitizeTypeWithContext(fieldLike.type, context);
  } catch (error: unknown) {
    throw Error(
      `Unable to sanitize type for field: ${fieldLike.name} due to error: ${error}`,
      { cause: error },
    );
  }
  const name = fieldLike.name;
  if (!(typeof name === "string")) {
    throw Error("The field passed in had a non-string `name` property");
  }
  const nullable = fieldLike.nullable;
  if (!(typeof nullable === "boolean")) {
    throw Error("The field passed in had a non-boolean `nullable` property");
  }
  let metadata;
  if ("metadata" in fieldLike) {
    metadata = sanitizeMetadata(fieldLike.metadata);
  }
  return new Field(name, type, nullable, metadata);
}

/**
 * Convert something schemaLike into a Schema instance
 *
 * This method is often needed even when the caller is using a Schema
 * instance because they might be using a different instance of apache-arrow
 * 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;
  }
  if (typeof schemaLike !== "object" || schemaLike === null) {
    throw Error("Expected a Schema but object was null/undefined");
  }
  if (!("fields" in schemaLike)) {
    throw Error(
      "The schema passed in does not appear to be a schema (no 'fields' property)",
    );
  }
  let metadata;
  if ("metadata" in schemaLike) {
    metadata = sanitizeMetadata(schemaLike.metadata);
  }
  if (!Array.isArray(schemaLike.fields)) {
    throw Error(
      "The schema passed in had a 'fields' property but it was not an array",
    );
  }
  const sanitizedFields = schemaLike.fields.map((field) =>
    sanitizeFieldWithContext(field, context),
  );
  return new Schema(sanitizedFields, metadata);
}

export function sanitizeTable(tableLike: TableLike): Table {
  if (tableLike instanceof Table) {
    return tableLike;
  }
  if (typeof tableLike !== "object" || tableLike === null) {
    throw Error("Expected a Table but object was null/undefined");
  }
  if (!("schema" in tableLike)) {
    throw Error(
      "The table passed in does not appear to be a table (no 'schema' property)",
    );
  }
  if (!("batches" in tableLike)) {
    throw Error(
      "The table passed in does not appear to be a table (no 'columns' property)",
    );
  }
  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,
  context: SanitizationContext,
): RecordBatch {
  if (batchLike instanceof RecordBatch) {
    return batchLike;
  }
  if (typeof batchLike !== "object" || batchLike === null) {
    throw Error("Expected a RecordBatch but object was null/undefined");
  }
  if (!("schema" in batchLike)) {
    throw Error(
      "The record batch passed in does not appear to be a record batch (no 'schema' property)",
    );
  }
  if (!("data" in batchLike)) {
    throw Error(
      "The record batch passed in does not appear to be a record batch (no 'data' property)",
    );
  }
  const schema = sanitizeSchemaWithContext(batchLike.schema, context);
  const data = sanitizeData(batchLike.data, context) as Data<Struct>;
  return new RecordBatch(schema, 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 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,
    {
      [BufferType.OFFSET]: dataLike.valueOffsets,
      [BufferType.DATA]: dataLike.values,
      [BufferType.VALIDITY]: dataLike.nullBitmap,
      [BufferType.TYPE]: dataLike.typeIds,
    },
    dataLike.children.map((child) => sanitizeData(child, context)),
    dictionary,
  );
  context.data.set(dataLike, data);
  return data;
}

const constructorsByTypeName = {
  null: () => new Null(),
  binary: () => new Binary(),
  utf8: () => new Utf8(),
  bool: () => new Bool(),
  int8: () => new Int8(),
  int16: () => new Int16(),
  int32: () => new Int32(),
  int64: () => new Int64(),
  uint8: () => new Uint8(),
  uint16: () => new Uint16(),
  uint32: () => new Uint32(),
  uint64: () => new Uint64(),
  float16: () => new Float16(),
  float32: () => new Float32(),
  float64: () => new Float64(),
  datemillisecond: () => new DateMillisecond(),
  dateday: () => new DateDay(),
  timenanosecond: () => new TimeNanosecond(),
  timemicrosecond: () => new TimeMicrosecond(),
  timemillisecond: () => new TimeMillisecond(),
  timesecond: () => new TimeSecond(),
  intervaldaytime: () => new IntervalDayTime(),
  intervalyearmonth: () => new IntervalYearMonth(),
  durationnanosecond: () => new DurationNanosecond(),
  durationmicrosecond: () => new DurationMicrosecond(),
  durationmillisecond: () => new DurationMillisecond(),
  durationsecond: () => new DurationSecond(),
} as const;

type MappableTypeName = keyof typeof constructorsByTypeName;

export function dataTypeFromName(typeName: string): DataType {
  const normalizedTypeName = typeName.toLowerCase() as MappableTypeName;
  const _constructor = constructorsByTypeName[normalizedTypeName];

  if (!_constructor) {
    throw new Error("Unrecognized type name in schema: " + typeName);
  }

  return _constructor();
}