From 6ebb6a9ee875f1c88c3f6a721f9cc4b93099fa44 Mon Sep 17 00:00:00 2001
From: Kenny Daniel <platypii@gmail.com>
Date: Mon, 29 Apr 2024 16:47:52 -0700
Subject: [PATCH] isListLike and isMapLike tests

---
 src/schema.js       |  2 +-
 test/schema.test.js | 60 +++++++++++++++++++++++++++++++++------------
 2 files changed, 46 insertions(+), 16 deletions(-)

diff --git a/src/schema.js b/src/schema.js
index d0769e9..f5d58ee 100644
--- a/src/schema.js
+++ b/src/schema.js
@@ -140,7 +140,7 @@ export function getColumnName(schema, path) {
  * @param {string[]} path column path
  * @returns {boolean} true if map-like
  */
-function isListLike(schemaElements, path) {
+export function isListLike(schemaElements, path) {
   const schema = schemaElement(schemaElements, path.slice(0, -2))
   if (path.length < 3) return false
   if (schema.element.converted_type !== 'LIST') return false
diff --git a/test/schema.test.js b/test/schema.test.js
index 39e7b9f..5fbbf87 100644
--- a/test/schema.test.js
+++ b/test/schema.test.js
@@ -2,6 +2,8 @@ import { describe, expect, it } from 'vitest'
 import {
   getMaxDefinitionLevel,
   getMaxRepetitionLevel,
+  isListLike,
+  isMapLike,
   isRequired,
   schemaElement,
   skipDefinitionBytes,
@@ -13,13 +15,19 @@ describe('Parquet schema utils', () => {
    * @type {SchemaElement[]}
    */
   const schema = [
-    { name: 'root', num_children: 2, repetition_type: 'REQUIRED' },
+    { name: 'root', num_children: 3 },
     { name: 'child1', repetition_type: 'OPTIONAL' },
-    { name: 'child2', repetition_type: 'REPEATED' },
+    { name: 'child2', repetition_type: 'OPTIONAL', num_children: 1, converted_type: 'LIST' },
+    { name: 'list', repetition_type: 'REPEATED', num_children: 1 },
+    { name: 'element', repetition_type: 'REQUIRED' },
+    { name: 'child3', repetition_type: 'OPTIONAL', num_children: 1, converted_type: 'MAP' },
+    { name: 'map', repetition_type: 'REPEATED', num_children: 2 },
+    { name: 'key', repetition_type: 'REQUIRED' },
+    { name: 'value', repetition_type: 'OPTIONAL' },
   ]
 
   describe('schemaElement', () => {
-    it('should return the correct schema element', () => {
+    it('should return the schema element', () => {
       expect(schemaElement(schema, ['child1'])).toEqual({
         children: [],
         count: 1,
@@ -33,26 +41,48 @@ describe('Parquet schema utils', () => {
     })
   })
 
-  describe('isRequired', () => {
-    it('should return true for required elements', () => {
-      expect(isRequired(schema, [])).toBe(true)
-    })
-
-    it('should return false for optional or repeated elements', () => {
-      expect(isRequired(schema, ['child1'])).toBe(false)
-    })
+  it('isRequired', () => {
+    expect(isRequired(schema, [])).toBe(true)
+    expect(isRequired(schema, ['child1'])).toBe(false)
+    expect(isRequired(schema, ['child2'])).toBe(false)
+    expect(isRequired(schema, ['child3'])).toBe(false)
   })
 
-  it('getMaxRepetitionLevel should return the correct max repetition level', () => {
-    expect(getMaxRepetitionLevel(schema, ['child2'])).toBe(1)
+  it('getMaxRepetitionLevel', () => {
+    expect(getMaxRepetitionLevel(schema, ['child1'])).toBe(0)
+    expect(getMaxRepetitionLevel(schema, ['child2'])).toBe(0)
+    expect(getMaxRepetitionLevel(schema, ['child2', 'list', 'element'])).toBe(1)
+    expect(getMaxRepetitionLevel(schema, ['child3'])).toBe(0)
+    expect(getMaxRepetitionLevel(schema, ['child3', 'map', 'key'])).toBe(1)
   })
 
-  it('getMaxDefinitionLevel should return the correct max definition level', () => {
+  it('getMaxDefinitionLevel', () => {
     expect(getMaxDefinitionLevel(schema, ['child1'])).toBe(1)
+    expect(getMaxDefinitionLevel(schema, ['child2'])).toBe(1)
+    expect(getMaxDefinitionLevel(schema, ['child3'])).toBe(1)
   })
 
-  it('skipDefinitionBytes should return the correct number of bytes to skip', () => {
+  it('skipDefinitionBytes', () => {
     expect(skipDefinitionBytes(100)).toBe(6)
     expect(skipDefinitionBytes(1000)).toBe(7)
   })
+
+  it('isListLike', () => {
+    expect(isListLike(schema, [])).toBe(false)
+    expect(isListLike(schema, ['child1'])).toBe(false)
+    expect(isListLike(schema, ['child2'])).toBe(false)
+    expect(isListLike(schema, ['child2', 'list', 'element'])).toBe(true)
+    expect(isListLike(schema, ['child3'])).toBe(false)
+    expect(isListLike(schema, ['child3', 'map', 'key'])).toBe(false)
+  })
+
+  it('isMapLike', () => {
+    expect(isMapLike(schema, [])).toBe(false)
+    expect(isMapLike(schema, ['child1'])).toBe(false)
+    expect(isMapLike(schema, ['child2'])).toBe(false)
+    expect(isMapLike(schema, ['child2', 'list', 'element'])).toBe(false)
+    expect(isMapLike(schema, ['child3'])).toBe(false)
+    expect(isMapLike(schema, ['child3', 'map', 'key'])).toBe(true)
+    expect(isMapLike(schema, ['child3', 'map', 'value'])).toBe(true)
+  })
 })