Simplify isListLike and isMapLike

This commit is contained in:
Kenny Daniel 2024-05-06 13:30:14 -07:00
parent 12dc5a47f8
commit 797db7b4df
No known key found for this signature in database
GPG Key ID: 90AB653A8CAD7E45
3 changed files with 19 additions and 24 deletions

@ -157,7 +157,7 @@ async function readRowGroup(options, rowGroup, groupStart) {
throw new Error(`parquet column length ${columnData.length} does not match row group length ${rowGroup.num_rows}`)
}
if (isMapLike(schemaPath)) {
if (isMapLike(schemaPath[schemaPath.length - 3])) {
const name = columnMetadata.path_in_schema.slice(0, -2).join('.')
if (!maps.has(name)) {
maps.set(name, columnData)

@ -95,11 +95,10 @@ export function getMaxDefinitionLevel(schemaPath) {
/**
* Check if a column is list-like.
*
* @param {SchemaTree[]} schemaPath
* @returns {boolean} true if map-like
* @param {SchemaTree} schema
* @returns {boolean} true if list-like
*/
export function isListLike(schemaPath) {
const schema = schemaPath.at(-3)
export function isListLike(schema) {
if (!schema) return false
if (schema.element.converted_type !== 'LIST') return false
if (schema.children.length > 1) return false
@ -108,20 +107,16 @@ export function isListLike(schemaPath) {
if (firstChild.children.length > 1) return false
if (firstChild.element.repetition_type !== 'REPEATED') return false
const secondChild = firstChild.children[0]
if (secondChild.element.repetition_type !== 'REQUIRED') return false
return true
}
/**
* Check if a column is map-like.
*
* @param {SchemaTree[]} schemaPath
* @param {SchemaTree} schema
* @returns {boolean} true if map-like
*/
export function isMapLike(schemaPath) {
const schema = schemaPath.at(-3)
export function isMapLike(schema) {
if (!schema) return false
if (schema.element.converted_type !== 'MAP') return false
if (schema.children.length > 1) return false

@ -64,21 +64,21 @@ describe('Parquet schema utils', () => {
})
it('isListLike', () => {
expect(isListLike(getSchemaPath(schema, []))).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child1']))).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child2']))).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child2', 'list', 'element']))).toBe(true)
expect(isListLike(getSchemaPath(schema, ['child3']))).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child3', 'map', 'key']))).toBe(false)
expect(isListLike(getSchemaPath(schema, [])[1])).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child1'])[1])).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child2'])[1])).toBe(true)
expect(isListLike(getSchemaPath(schema, ['child2', 'list', 'element'])[1])).toBe(true)
expect(isListLike(getSchemaPath(schema, ['child3'])[1])).toBe(false)
expect(isListLike(getSchemaPath(schema, ['child3', 'map', 'key'])[1])).toBe(false)
})
it('isMapLike', () => {
expect(isMapLike(getSchemaPath(schema, []))).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child1']))).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child2']))).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child2', 'list', 'element']))).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child3']))).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child3', 'map', 'key']))).toBe(true)
expect(isMapLike(getSchemaPath(schema, ['child3', 'map', 'value']))).toBe(true)
expect(isMapLike(getSchemaPath(schema, [])[1])).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child1'])[1])).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child2'])[1])).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child2', 'list', 'element'])[1])).toBe(false)
expect(isMapLike(getSchemaPath(schema, ['child3'])[1])).toBe(true)
expect(isMapLike(getSchemaPath(schema, ['child3', 'map', 'key'])[1])).toBe(true)
expect(isMapLike(getSchemaPath(schema, ['child3', 'map', 'value'])[1])).toBe(true)
})
})