File size: 4,201 Bytes
9ae1216 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | /**
* Build a tree from the schema elements.
*
* @param {SchemaElement[]} schema
* @param {number} rootIndex index of the root element
* @param {string[]} path path to the element
* @returns {SchemaTree} tree of schema elements
*/
function schemaTree(schema, rootIndex, path) {
const element = schema[rootIndex]
const children = []
let count = 1
// Read the specified number of children
if (element.num_children) {
while (children.length < element.num_children) {
const childElement = schema[rootIndex + count]
const child = schemaTree(schema, rootIndex + count, [...path, childElement.name])
count += child.count
children.push(child)
}
}
return { count, element, children, path }
}
/**
* Get schema elements from the root to the given element name.
*
* @param {SchemaElement[]} schema
* @param {string[]} name path to the element
* @returns {SchemaTree[]} list of schema elements
*/
export function getSchemaPath(schema, name) {
let tree = schemaTree(schema, 0, [])
const path = [tree]
for (const part of name) {
const child = tree.children.find(child => child.element.name === part)
if (!child) throw new Error(`parquet schema element not found: ${name}`)
path.push(child)
tree = child
}
return path
}
/**
* Get all physical (leaf) column names.
*
* @param {SchemaTree} schemaTree
* @returns {string[]} list of physical column names
*/
export function getPhysicalColumns(schemaTree) {
/** @type {string[]} */
const columns = []
/** @param {SchemaTree} node */
function traverse(node) {
if (node.children.length) {
for (const child of node.children) {
traverse(child)
}
} else {
columns.push(node.path.join('.'))
}
}
traverse(schemaTree)
return columns
}
/**
* Get the max repetition level for a given schema path.
*
* @param {SchemaTree[]} schemaPath
* @returns {number} max repetition level
*/
export function getMaxRepetitionLevel(schemaPath) {
let maxLevel = 0
for (const { element } of schemaPath) {
if (element.repetition_type === 'REPEATED') {
maxLevel++
}
}
return maxLevel
}
/**
* Get the max definition level for a given schema path.
*
* @param {SchemaTree[]} schemaPath
* @returns {number} max definition level
*/
export function getMaxDefinitionLevel(schemaPath) {
let maxLevel = 0
for (const { element } of schemaPath.slice(1)) {
if (element.repetition_type !== 'REQUIRED') {
maxLevel++
}
}
return maxLevel
}
/**
* Check if a column is list-like.
*
* @param {SchemaTree} schema
* @returns {boolean} true if list-like
*/
export function isListLike(schema) {
if (!schema) return false
if (schema.element.converted_type !== 'LIST') return false
if (schema.children.length > 1) return false
const firstChild = schema.children[0]
if (firstChild.children.length > 1) return false
if (firstChild.element.repetition_type !== 'REPEATED') return false
return true
}
/**
* Check if a column is map-like.
*
* @param {SchemaTree} schema
* @returns {boolean} true if map-like
*/
export function isMapLike(schema) {
if (!schema) return false
if (schema.element.converted_type !== 'MAP') return false
if (schema.children.length > 1) return false
const firstChild = schema.children[0]
if (firstChild.children.length !== 2) return false
if (firstChild.element.repetition_type !== 'REPEATED') return false
const keyChild = firstChild.children.find(child => child.element.name === 'key')
if (keyChild?.element.repetition_type === 'REPEATED') return false
const valueChild = firstChild.children.find(child => child.element.name === 'value')
if (valueChild?.element.repetition_type === 'REPEATED') return false
return true
}
/**
* Returns true if a column is non-nested.
*
* @param {SchemaTree[]} schemaPath
* @returns {boolean}
*/
export function isFlatColumn(schemaPath) {
if (schemaPath.length !== 2) return false
const [, column] = schemaPath
if (column.element.repetition_type === 'REPEATED') return false
if (column.children.length) return false
return true
}
/**
* @import {SchemaElement, SchemaTree} from '../src/types.js'
*/
|