ceph/src/arrow/js/src/schema.ts

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied.  See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 import { DataType } from './type';
  19
  20 export class Schema<T extends { [key: string]: DataType } = any> {
  21
  22     public readonly fields: Field<T[keyof T]>[];
  23     public readonly metadata: Map<string, string>;
  24     public readonly dictionaries: Map<number, DataType>;
  25
  26     constructor(fields: Field[] = [],
  27                 metadata?: Map<string, string> | null,
  28                 dictionaries?: Map<number, DataType> | null) {
  29         this.fields = (fields || []) as Field<T[keyof T]>[];
  30         this.metadata = metadata || new Map();
  31         if (!dictionaries) {
  32             dictionaries = generateDictionaryMap(fields);
  33         }
  34         this.dictionaries = dictionaries;
  35     }
  36     public get [Symbol.toStringTag]() { return 'Schema'; }
  37     public toString() {
  38         return `Schema<{ ${this.fields.map((f, i) => `${i}: ${f}`).join(', ')} }>`;
  39     }
  40
  41     public select<K extends keyof T = any>(...columnNames: K[]) {
  42         const names = columnNames.reduce((xs, x) => (xs[x] = true) && xs, Object.create(null));
  43         return new Schema<{ [P in K]: T[P] }>(this.fields.filter((f) => names[f.name]), this.metadata);
  44     }
  45     public selectAt<K extends T[keyof T] = any>(...columnIndices: number[]) {
  46         return new Schema<{ [key: string]: K }>(columnIndices.map((i) => this.fields[i]).filter(Boolean), this.metadata);
  47     }
  48
  49     public assign<R extends { [key: string]: DataType } = any>(schema: Schema<R>): Schema<T & R>;
  50     public assign<R extends { [key: string]: DataType } = any>(...fields: (Field<R[keyof R]> | Field<R[keyof R]>[])[]): Schema<T & R>;
  51     public assign<R extends { [key: string]: DataType } = any>(...args: (Schema<R> | Field<R[keyof R]> | Field<R[keyof R]>[])[]) {
  52
  53         const other = (args[0] instanceof Schema
  54             ? args[0] as Schema<R>
  55             : Array.isArray(args[0])
  56                 ? new Schema<R>(<Field<R[keyof R]>[]> args[0])
  57                 : new Schema<R>(<Field<R[keyof R]>[]> args));
  58
  59         const curFields = [...this.fields] as Field[];
  60         const metadata = mergeMaps(mergeMaps(new Map(), this.metadata), other.metadata);
  61         const newFields = other.fields.filter((f2) => {
  62             const i = curFields.findIndex((f) => f.name === f2.name);
  63             return ~i ? (curFields[i] = f2.clone({
  64                 metadata: mergeMaps(mergeMaps(new Map(), curFields[i].metadata), f2.metadata)
  65             })) && false : true;
  66         }) as Field[];
  67
  68         const newDictionaries = generateDictionaryMap(newFields, new Map());
  69
  70         return new Schema<T & R>(
  71             [...curFields, ...newFields], metadata,
  72             new Map([...this.dictionaries, ...newDictionaries])
  73         );
  74     }
  75 }
  76
  77 export class Field<T extends DataType = any> {
  78
  79     public static new<T extends DataType = any>(props: { name: string | number; type: T; nullable?: boolean; metadata?: Map<string, string> | null }): Field<T>;
  80     public static new<T extends DataType = any>(name: string | number | Field<T>, type: T, nullable?: boolean, metadata?: Map<string, string> | null): Field<T>;
  81     /** @nocollapse */
  82     public static new<T extends DataType = any>(...args: any[]) {
  83         let [name, type, nullable, metadata] = args;
  84         if (args[0] && typeof args[0] === 'object') {
  85             ({ name } = args[0]);
  86             (type === undefined) && (type = args[0].type);
  87             (nullable === undefined) && (nullable = args[0].nullable);
  88             (metadata === undefined) && (metadata = args[0].metadata);
  89         }
  90         return new Field<T>(`${name}`, type, nullable, metadata);
  91     }
  92
  93     public readonly type: T;
  94     public readonly name: string;
  95     public readonly nullable: boolean;
  96     public readonly metadata: Map<string, string>;
  97
  98     constructor(name: string, type: T, nullable = false, metadata?: Map<string, string> | null) {
  99         this.name = name;
 100         this.type = type;
 101         this.nullable = nullable;
 102         this.metadata = metadata || new Map();
 103     }
 104
 105     public get typeId() { return this.type.typeId; }
 106     public get [Symbol.toStringTag]() { return 'Field'; }
 107     public toString() { return `${this.name}: ${this.type}`; }
 108     public clone<R extends DataType = T>(props: { name?: string | number; type?: R; nullable?: boolean; metadata?: Map<string, string> | null }): Field<R>;
 109     public clone<R extends DataType = T>(name?: string | number | Field<T>, type?: R, nullable?: boolean, metadata?: Map<string, string> | null): Field<R>;
 110     public clone<R extends DataType = T>(...args: any[]) {
 111         let [name, type, nullable, metadata] = args;
 112         (!args[0] || typeof args[0] !== 'object')
 113             ? ([name = this.name, type = this.type, nullable = this.nullable, metadata = this.metadata] = args)
 114             : ({name = this.name, type = this.type, nullable = this.nullable, metadata = this.metadata} = args[0]);
 115         return Field.new<R>(name, type, nullable, metadata);
 116     }
 117 }
 118
 119 /** @ignore */
 120 function mergeMaps<TKey, TVal>(m1?: Map<TKey, TVal> | null, m2?: Map<TKey, TVal> | null): Map<TKey, TVal> {
 121     return new Map([...(m1 || new Map()), ...(m2 || new Map())]);
 122 }
 123
 124 /** @ignore */
 125 function generateDictionaryMap(fields: Field[], dictionaries = new Map<number, DataType>()): Map<number, DataType> {
 126
 127     for (let i = -1, n = fields.length; ++i < n;) {
 128         const field = fields[i];
 129         const type = field.type;
 130         if (DataType.isDictionary(type)) {
 131             if (!dictionaries.has(type.id)) {
 132                 dictionaries.set(type.id, type.dictionary);
 133             } else if (dictionaries.get(type.id) !== type.dictionary) {
 134                 throw new Error(`Cannot create Schema containing two different dictionaries with the same Id`);
 135             }
 136         }
 137         if (type.children && type.children.length > 0) {
 138             generateDictionaryMap(type.children, dictionaries);
 139         }
 140     }
 141
 142     return dictionaries;
 143 }
 144
 145 // Add these here so they're picked up by the externs creator
 146 // in the build, and closure-compiler doesn't minify them away
 147 (Schema.prototype as any).fields = null;
 148 (Schema.prototype as any).metadata = null;
 149 (Schema.prototype as any).dictionaries = null;
 150
 151 (Field.prototype as any).type = null;
 152 (Field.prototype as any).name = null;
 153 (Field.prototype as any).nullable = null;
 154 (Field.prototype as any).metadata = null;