From 00a8dbf57f377e4486b6c6b5dfaa4afb49d051d7 Mon Sep 17 00:00:00 2001
From: Kenny Daniel <platypii@gmail.com>
Date: Sun, 14 Jan 2024 15:36:58 -0800
Subject: [PATCH] Schema util tests

---
 src/constants.js    |  6 ++++++
 src/schema.js       |  2 +-
 test/schema.test.js | 51 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)
 create mode 100644 test/schema.test.js

diff --git a/src/constants.js b/src/constants.js
index 302f974..61e785c 100644
--- a/src/constants.js
+++ b/src/constants.js
@@ -20,3 +20,9 @@ export const ParquetEncoding = {
   RLE_DICTIONARY: 8,
   BYTE_STREAM_SPLIT: 9,
 }
+
+export const FieldRepetitionType = {
+  REQUIRED: 0,
+  OPTIONAL: 1,
+  REPEATED: 2,
+}
diff --git a/src/schema.js b/src/schema.js
index 7b2e0fd..9a45fcd 100644
--- a/src/schema.js
+++ b/src/schema.js
@@ -1,4 +1,4 @@
-import { FieldRepetitionType } from './types.js'
+import { FieldRepetitionType } from './constants.js'
 
 /**
  * @typedef {import('./types.js').SchemaElement} SchemaElement
diff --git a/test/schema.test.js b/test/schema.test.js
new file mode 100644
index 0000000..c53dc93
--- /dev/null
+++ b/test/schema.test.js
@@ -0,0 +1,51 @@
+import { describe, expect, it } from 'vitest'
+import { FieldRepetitionType } from '../src/constants.js'
+import {
+  getMaxDefinitionLevel,
+  getMaxRepetitionLevel,
+  isRequired,
+  schemaElement,
+  skipDefinitionBytes,
+} from '../src/schema.js'
+
+describe('Parquet schema utils', () => {
+  const schema = [
+    { name: 'root', num_children: 2, repetition_type: FieldRepetitionType.REQUIRED },
+    { name: 'child1', repetition_type: FieldRepetitionType.OPTIONAL },
+    { name: 'child2', repetition_type: FieldRepetitionType.REPEATED },
+  ]
+
+  describe('schemaElement', () => {
+    it('should return the correct schema element', () => {
+      expect(schemaElement(schema, ['child1'])).toEqual(schema[1])
+    })
+
+    it('should throw an error if element not found', () => {
+      expect(() => schemaElement(schema, ['nonexistent']))
+        .toThrow('parquet schema element not found: nonexistent')
+    })
+  })
+
+  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('getMaxRepetitionLevel should return the correct max repetition level', () => {
+    expect(getMaxRepetitionLevel(schema, ['child2'])).toBe(1)
+  })
+
+  it('getMaxDefinitionLevel should return the correct max definition level', () => {
+    expect(getMaxDefinitionLevel(schema, ['child1'])).toBe(1)
+  })
+
+  it('skipDefinitionBytes should return the correct number of bytes to skip', () => {
+    expect(skipDefinitionBytes(100)).toBe(6)
+    expect(skipDefinitionBytes(1000)).toBe(7)
+  })
+})