Line data Source code
1 : import * as arrow from 'apache-arrow';
2 :
3 : import { binaryVector, boolVector, floatVector, intVector, utf8Vector } from '../../utils/arrow_vectors.js';
4 :
5 : /// An Arrow table that Hyper can read with external(format => 'arrowstream').
6 : export interface HyperArrowInsert {
7 : /// The table with the columns converted to types that Hyper reads
8 : table: arrow.Table;
9 : /// The column descriptor of the external scan, Hyper cannot infer it from the Arrow schema
10 : columns: string;
11 : }
12 :
13 1 : function quoteIdentifier(value: string): string {
14 1 : return `"${value.replace(/"/g, '""')}"`;
15 1 : }
16 :
17 : /// The Hyper type that the metadata of a field names, e.g. of a json or varchar(n) column exported by Hyper.
18 : /// Hyper insists on this type when it reads the field back.
19 1 : function annotatedHyperType(field: arrow.Field): string | null {
20 1 : const length = field.metadata.get('hyper:max_string_length');
21 1 : switch (field.metadata.get('hyper:type')) {
22 1 : case 'JSON':
23 1 : return 'json';
24 1 : case 'Oid':
25 1 : return 'oid';
26 1 : case 'Geography':
27 1 : return 'geography';
28 1 : case 'Char':
29 1 : return length == null ? null : `char(${length})`;
30 1 : }
31 1 : return length == null ? null : `varchar(${length})`;
32 1 : }
33 :
34 : /// The Hyper type of an Arrow field that Hyper reads, null when the column needs converting first.
35 : /// Hyper insists on the type that corresponds to the Arrow type, a wider type is rejected.
36 1 : function readableHyperType(field: arrow.Field): string | null {
37 1 : const type = field.type as arrow.DataType;
38 1 : const annotated = annotatedHyperType(field);
39 1 : if (annotated != null) return annotated;
40 1 : switch (type.typeId) {
41 1 : case arrow.Type.Int: {
42 1 : const int = type as arrow.Int;
43 1 : switch (int.bitWidth) {
44 1 : case 8:
45 1 : case 16:
46 1 : return int.isSigned || int.bitWidth === 8 ? 'smallint' : 'integer';
47 1 : case 32:
48 1 : return int.isSigned ? 'integer' : 'bigint';
49 1 : case 64:
50 1 : return int.isSigned ? 'bigint' : 'numeric(20,0)';
51 1 : }
52 0 : return null;
53 0 : }
54 1 : case arrow.Type.Float:
55 1 : switch ((type as arrow.Float).precision) {
56 1 : case arrow.Precision.SINGLE:
57 1 : return 'real';
58 1 : case arrow.Precision.DOUBLE:
59 1 : return 'double precision';
60 1 : }
61 1 : return null;
62 1 : case arrow.Type.Utf8:
63 1 : return 'text';
64 1 : case arrow.Type.Binary:
65 1 : return 'bytea';
66 1 : case arrow.Type.Bool:
67 1 : return 'boolean';
68 1 : case arrow.Type.Decimal: {
69 1 : const decimal = type as arrow.Decimal;
70 1 : return decimal.bitWidth === 128 ? `numeric(${decimal.precision},${decimal.scale})` : null;
71 1 : }
72 1 : case arrow.Type.Date:
73 1 : return 'date';
74 1 : case arrow.Type.Time:
75 1 : return 'time';
76 1 : case arrow.Type.Timestamp: {
77 1 : const timestamp = type as arrow.Timestamp;
78 1 : if (timestamp.unit === arrow.TimeUnit.NANOSECOND) return null;
79 1 : return timestamp.timezone == null ? 'timestamp' : 'timestamptz';
80 1 : }
81 1 : case arrow.Type.Interval:
82 1 : return 'interval';
83 1 : case arrow.Type.List: {
84 : // The element nullability is part of the array type, e.g. array_agg yields non-null elements.
85 1 : const child = (type as arrow.List).children[0];
86 1 : const element = readableHyperType(child);
87 1 : if (element == null) return null;
88 1 : return child.nullable ? `array(${element})` : `array(${element} not null)`;
89 1 : }
90 1 : default:
91 1 : return null;
92 1 : }
93 1 : }
94 :
95 : /// A vector of values of an Arrow type in a type that Hyper reads, one chunk.
96 : /// The vector is assembled from its buffers, the Arrow builders evaluate code (see arrow_vectors.ts).
97 1 : function readableVector(type: arrow.DataType, values: unknown[]): arrow.Vector {
98 1 : switch (type.typeId) {
99 1 : case arrow.Type.Utf8:
100 1 : case arrow.Type.LargeUtf8:
101 1 : return utf8Vector(values as Array<string | null>);
102 0 : case arrow.Type.Binary:
103 1 : case arrow.Type.LargeBinary:
104 1 : case arrow.Type.FixedSizeBinary:
105 1 : return binaryVector(values as Array<Uint8Array | null>);
106 1 : case arrow.Type.Float:
107 1 : return (type as arrow.Float).precision === arrow.Precision.DOUBLE
108 0 : ? floatVector(new arrow.Float64(), values as Array<number | null>)
109 1 : : floatVector(new arrow.Float32(), values as Array<number | null>);
110 1 : case arrow.Type.Int:
111 1 : return intVector(type as arrow.Int, values as Array<number | bigint | null>);
112 0 : case arrow.Type.Bool:
113 0 : return boolVector(values as Array<boolean | null>);
114 1 : case arrow.Type.Null:
115 1 : return utf8Vector(values.map(() => null));
116 0 : default:
117 0 : throw new Error(`Hyper cannot import Arrow columns of type ${String(type)}`);
118 1 : }
119 1 : }
120 :
121 : /// Convert a column that Hyper cannot read to a type that it reads.
122 1 : function convertColumn(vector: arrow.Vector): arrow.Vector {
123 1 : const type = vector.type as arrow.DataType;
124 1 : switch (type.typeId) {
125 1 : case arrow.Type.Dictionary:
126 1 : return readableVector((type as arrow.Dictionary).dictionary, [...vector]);
127 1 : case arrow.Type.LargeUtf8:
128 1 : case arrow.Type.LargeBinary:
129 1 : case arrow.Type.FixedSizeBinary:
130 1 : case arrow.Type.Float:
131 1 : case arrow.Type.Null:
132 1 : return readableVector(type, [...vector]);
133 1 : case arrow.Type.Timestamp: {
134 : // Hyper keeps microseconds, nanoseconds are truncated.
135 1 : const timezone = (type as arrow.Timestamp).timezone;
136 1 : return new arrow.Vector(vector.data.map(data => arrow.makeData({
137 1 : type: new arrow.TimestampMicrosecond(timezone),
138 1 : length: data.length,
139 1 : nullCount: data.nullCount,
140 1 : nullBitmap: data.nullBitmap,
141 1 : data: (data.values as BigInt64Array).map(value => value / 1000n),
142 1 : })));
143 1 : }
144 1 : default:
145 1 : throw new Error(`Hyper cannot import Arrow columns of type ${String(type)}`);
146 1 : }
147 1 : }
148 :
149 : /// Prepare an Arrow table for an external Arrow stream scan in Hyper.
150 1 : export function prepareHyperArrowInsert(table: arrow.Table): HyperArrowInsert {
151 1 : const columns: Record<string, arrow.Vector> = {};
152 1 : let converted = false;
153 1 : table.schema.fields.forEach((field, index) => {
154 1 : let column = table.getChildAt(index)!;
155 1 : if (readableHyperType(field) == null) {
156 1 : column = convertColumn(column);
157 1 : converted = true;
158 1 : }
159 1 : columns[field.name] = column;
160 1 : });
161 : // Converted columns are single chunks, the table redistributes the other columns into matching batches.
162 : // The rebuilt table drops the field metadata, so the descriptor follows the fields of the scanned table.
163 1 : const scanned = converted ? new arrow.Table(columns) : table;
164 1 : const descriptor = scanned.schema.fields.map(field => {
165 1 : const type = readableHyperType(field);
166 1 : if (type == null) {
167 0 : throw new Error(`Hyper cannot import Arrow columns of type ${String(field.type)}`);
168 : }
169 1 : return `${quoteIdentifier(field.name)} ${type}`;
170 1 : });
171 1 : return { table: scanned, columns: descriptor.join(', ') };
172 1 : }
|
