@dotdo/postgres 0.1.0

package/src/iceberg/parquet.test.ts

@@ -0,0 +1,899 @@
+/**
+ * Parquet Writer Module Tests
+ * Task: postgres-8gdz.1 - [RED] Audit source files vs test files for iceberg
+ *
+ * Dedicated tests for parquet.ts module exports:
+ * - WAL_PARQUET_SCHEMA constant
+ * - ParquetWriter class
+ * - walEntryToParquetRecord()
+ * - generateDataFileMetadata()
+ * - createWALParquetWriter()
+ * - WALParquetBatchWriter class
+ * - createWALParquetBatchWriter()
+ * - bloomFilterMightContain()
+ *
+ * Note: Some basic tests exist in iceberg.test.ts.
+ * This file adds more comprehensive coverage.
+ */
+
+import { describe, it, expect, beforeEach, vi } from 'vitest'
+import {
+  WAL_PARQUET_SCHEMA,
+  ParquetWriter,
+  walEntryToParquetRecord,
+  generateDataFileMetadata,
+  createWALParquetWriter,
+  WALParquetBatchWriter,
+  createWALParquetBatchWriter,
+  bloomFilterMightContain,
+  type ParquetWriterConfig,
+  type WALParquetRecord,
+  type BloomFilter,
+  type WALParquetBatchConfig,
+} from './parquet'
+import { createWALSchema } from './schema'
+import type { SerializedWALEntry } from '../worker/wal'
+
+// ============================================================================
+// Mock R2Bucket
+// ============================================================================
+
+function createMockBucket(): R2Bucket {
+  const storage = new Map<string, { data: string | ArrayBuffer; metadata?: Record<string, string> }>()
+
+  return {
+    put: vi.fn(async (key: string, data: string | ArrayBuffer | ReadableStream, options?: R2PutOptions) => {
+      storage.set(key, {
+        data: data as string | ArrayBuffer,
+        metadata: options?.customMetadata,
+      })
+      return {
+        key,
+        version: 'test-version',
+        size: typeof data === 'string' ? data.length : (data as ArrayBuffer).byteLength,
+        etag: 'test-etag',
+        httpEtag: '"test-etag"',
+        uploaded: new Date(),
+        httpMetadata: options?.httpMetadata || {},
+        customMetadata: options?.customMetadata || {},
+      } as R2Object
+    }),
+    get: vi.fn(async (key: string) => {
+      const item = storage.get(key)
+      if (!item) return null
+
+      return {
+        key,
+        version: 'test-version',
+        size: typeof item.data === 'string' ? item.data.length : (item.data as ArrayBuffer).byteLength,
+        etag: 'test-etag',
+        httpEtag: '"test-etag"',
+        uploaded: new Date(),
+        httpMetadata: {},
+        customMetadata: item.metadata || {},
+        text: async () => typeof item.data === 'string' ? item.data : new TextDecoder().decode(item.data as ArrayBuffer),
+        json: async () => JSON.parse(typeof item.data === 'string' ? item.data : new TextDecoder().decode(item.data as ArrayBuffer)),
+        arrayBuffer: async () => typeof item.data === 'string' ? new TextEncoder().encode(item.data).buffer : item.data,
+        bytes: async () => typeof item.data === 'string' ? new TextEncoder().encode(item.data) : new Uint8Array(item.data as ArrayBuffer),
+        blob: async () => new Blob([item.data]),
+        body: null,
+        bodyUsed: false,
+        writeHttpMetadata: () => {},
+      } as R2ObjectBody
+    }),
+    head: vi.fn(async () => null),
+    delete: vi.fn(async () => {}),
+    list: vi.fn(async () => ({ objects: [], truncated: false, delimitedPrefixes: [] } as R2Objects)),
+    createMultipartUpload: vi.fn(),
+    resumeMultipartUpload: vi.fn(),
+  } as unknown as R2Bucket
+}
+
+// ============================================================================
+// WAL_PARQUET_SCHEMA Tests
+// ============================================================================
+
+describe('WAL_PARQUET_SCHEMA', () => {
+  it('should define all required WAL columns', () => {
+    const columnNames = WAL_PARQUET_SCHEMA.map((c) => c.name)
+
+    expect(columnNames).toContain('do_id')
+    expect(columnNames).toContain('lsn')
+    expect(columnNames).toContain('timestamp')
+    expect(columnNames).toContain('operation')
+    expect(columnNames).toContain('schema_name')
+    expect(columnNames).toContain('table_name')
+    expect(columnNames).toContain('primary_key')
+    expect(columnNames).toContain('new_row')
+    expect(columnNames).toContain('old_row')
+    expect(columnNames).toContain('changed_columns')
+    expect(columnNames).toContain('transaction_id')
+    expect(columnNames).toContain('sql')
+  })
+
+  it('should have correct Parquet types', () => {
+    const timestampCol = WAL_PARQUET_SCHEMA.find((c) => c.name === 'timestamp')
+    expect(timestampCol?.type).toBe('INT64')
+    expect(timestampCol?.logicalType).toBe('TIMESTAMP_MILLIS')
+
+    const stringCols = WAL_PARQUET_SCHEMA.filter((c) =>
+      ['do_id', 'lsn', 'operation'].includes(c.name)
+    )
+    stringCols.forEach((col) => {
+      expect(col.type).toBe('BYTE_ARRAY')
+      expect(col.logicalType).toBe('STRING')
+    })
+  })
+
+  it('should have correct repetition (required vs optional)', () => {
+    const requiredCols = WAL_PARQUET_SCHEMA.filter((c) => c.repetition === 'REQUIRED')
+    const optionalCols = WAL_PARQUET_SCHEMA.filter((c) => c.repetition === 'OPTIONAL')
+
+    expect(requiredCols.length).toBeGreaterThan(0)
+    expect(optionalCols.length).toBeGreaterThan(0)
+
+    // Check specific fields
+    expect(WAL_PARQUET_SCHEMA.find((c) => c.name === 'do_id')?.repetition).toBe('REQUIRED')
+    expect(WAL_PARQUET_SCHEMA.find((c) => c.name === 'new_row')?.repetition).toBe('OPTIONAL')
+  })
+
+  it('should have field IDs for Iceberg mapping', () => {
+    WAL_PARQUET_SCHEMA.forEach((col) => {
+      expect(col.fieldId).toBeDefined()
+      expect(typeof col.fieldId).toBe('number')
+    })
+  })
+
+  it('should have dictionary encoding for low-cardinality columns', () => {
+    const dictionaryColumns = WAL_PARQUET_SCHEMA.filter((c) => c.useDictionary)
+
+    expect(dictionaryColumns.map((c) => c.name)).toContain('do_id')
+    expect(dictionaryColumns.map((c) => c.name)).toContain('operation')
+    expect(dictionaryColumns.map((c) => c.name)).toContain('schema_name')
+    expect(dictionaryColumns.map((c) => c.name)).toContain('table_name')
+  })
+
+  it('should have bloom filters for equality-predicate columns', () => {
+    const bloomFilterColumns = WAL_PARQUET_SCHEMA.filter((c) => c.bloomFilter)
+
+    expect(bloomFilterColumns.map((c) => c.name)).toContain('do_id')
+    expect(bloomFilterColumns.map((c) => c.name)).toContain('lsn')
+    expect(bloomFilterColumns.map((c) => c.name)).toContain('operation')
+  })
+})
+
+// ============================================================================
+// walEntryToParquetRecord Tests
+// ============================================================================
+
+describe('walEntryToParquetRecord', () => {
+  it('should convert INSERT entry correctly', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12345',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      newRow: { id: 1, name: 'Alice' },
+      timestamp: 1704067200000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.do_id).toBe('test-do')
+    expect(record.lsn).toBe('12345')
+    expect(record.operation).toBe('INSERT')
+    expect(record.schema_name).toBe('public')
+    expect(record.table_name).toBe('users')
+    expect(record.new_row).toBe('{"id":1,"name":"Alice"}')
+    expect(record.old_row).toBeNull()
+    expect(record.timestamp).toBe(1704067200000)
+  })
+
+  it('should convert UPDATE entry correctly', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12346',
+      operation: 'UPDATE',
+      schema: 'public',
+      table: 'users',
+      oldRow: { id: 1, name: 'Alice' },
+      newRow: { id: 1, name: 'Bob' },
+      changedColumns: ['name'],
+      timestamp: 1704067260000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.operation).toBe('UPDATE')
+    expect(record.old_row).toBe('{"id":1,"name":"Alice"}')
+    expect(record.new_row).toBe('{"id":1,"name":"Bob"}')
+    expect(record.changed_columns).toBe('["name"]')
+  })
+
+  it('should convert DELETE entry correctly', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12347',
+      operation: 'DELETE',
+      schema: 'public',
+      table: 'users',
+      oldRow: { id: 1, name: 'Alice' },
+      timestamp: 1704067320000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.operation).toBe('DELETE')
+    expect(record.old_row).toBe('{"id":1,"name":"Alice"}')
+    expect(record.new_row).toBeNull()
+  })
+
+  it('should handle primary key', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12348',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      primaryKey: { id: 1 },
+      newRow: { id: 1, name: 'Alice' },
+      timestamp: 1704067380000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.primary_key).toBe('{"id":1}')
+  })
+
+  it('should handle transaction ID', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12349',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      transactionId: 'tx-123',
+      timestamp: 1704067440000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.transaction_id).toBe('tx-123')
+  })
+
+  it('should handle SQL statement', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12350',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      sql: "INSERT INTO users VALUES (1, 'Alice')",
+      timestamp: 1704067500000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.sql).toBe("INSERT INTO users VALUES (1, 'Alice')")
+  })
+
+  it('should handle missing optional fields', () => {
+    const entry: SerializedWALEntry = {
+      lsn: '12351',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'users',
+      timestamp: 1704067560000,
+    }
+
+    const record = walEntryToParquetRecord(entry, 'test-do')
+
+    expect(record.primary_key).toBeNull()
+    expect(record.new_row).toBeNull()
+    expect(record.old_row).toBeNull()
+    expect(record.changed_columns).toBeNull()
+    expect(record.transaction_id).toBeNull()
+    expect(record.sql).toBeNull()
+  })
+})
+
+// ============================================================================
+// ParquetWriter Tests
+// ============================================================================
+
+describe('ParquetWriter', () => {
+  let writer: ParquetWriter
+
+  const createRecord = (overrides: Partial<WALParquetRecord> = {}): WALParquetRecord => ({
+    do_id: 'test-do',
+    lsn: '1',
+    timestamp: Date.now(),
+    operation: 'INSERT',
+    schema_name: 'public',
+    table_name: 'test',
+    primary_key: null,
+    new_row: '{}',
+    old_row: null,
+    changed_columns: null,
+    transaction_id: null,
+    sql: null,
+    ...overrides,
+  })
+
+  beforeEach(() => {
+    writer = createWALParquetWriter()
+  })
+
+  describe('addRecord', () => {
+    it('should add a single record', () => {
+      writer.addRecord(createRecord())
+
+      expect(writer.getRecordCount()).toBe(1)
+    })
+
+    it('should track row group size', () => {
+      const initialSize = writer.getCurrentRowGroupSize()
+
+      writer.addRecord(createRecord({ new_row: '{"large": "data"}' }))
+
+      expect(writer.getCurrentRowGroupSize()).toBeGreaterThan(initialSize)
+    })
+  })
+
+  describe('addRecords', () => {
+    it('should add multiple records', () => {
+      const records = [createRecord({ lsn: '1' }), createRecord({ lsn: '2' }), createRecord({ lsn: '3' })]
+
+      writer.addRecords(records)
+
+      expect(writer.getRecordCount()).toBe(3)
+    })
+  })
+
+  describe('getRecordCount', () => {
+    it('should return 0 for empty writer', () => {
+      expect(writer.getRecordCount()).toBe(0)
+    })
+  })
+
+  describe('shouldFlushRowGroup', () => {
+    it('should return false for small data', () => {
+      writer.addRecord(createRecord())
+
+      expect(writer.shouldFlushRowGroup()).toBe(false)
+    })
+
+    it('should return true when row group size exceeds target', () => {
+      // Create writer with very small target row group size
+      const smallRowGroupWriter = createWALParquetWriter(undefined, {
+        targetRowGroupSize: 100, // 100 bytes - very small
+      })
+
+      // Add records with large data until we exceed target
+      const largeRecord: WALParquetRecord = {
+        do_id: 'test-do',
+        lsn: '1',
+        timestamp: Date.now(),
+        operation: 'INSERT',
+        schema_name: 'public',
+        table_name: 'test',
+        primary_key: null,
+        new_row: JSON.stringify({ data: 'x'.repeat(50) }), // Large payload
+        old_row: null,
+        changed_columns: null,
+        transaction_id: null,
+        sql: null,
+      }
+
+      smallRowGroupWriter.addRecord(largeRecord)
+      smallRowGroupWriter.addRecord(largeRecord)
+      smallRowGroupWriter.addRecord(largeRecord)
+
+      // With 3 large records, should exceed 100 byte target
+      expect(smallRowGroupWriter.shouldFlushRowGroup()).toBe(true)
+    })
+  })
+
+  describe('getColumnStats', () => {
+    it('should compute statistics for each column', () => {
+      for (let i = 0; i < 10; i++) {
+        writer.addRecord(createRecord({ lsn: String(i) }))
+      }
+
+      const stats = writer.getColumnStats()
+
+      expect(stats.get('do_id')?.valueCount).toBe(10)
+      expect(stats.get('do_id')?.nullCount).toBe(0)
+      expect(stats.get('lsn')?.valueCount).toBe(10)
+    })
+
+    it('should track null counts', () => {
+      for (let i = 0; i < 5; i++) {
+        writer.addRecord(createRecord({ new_row: null }))
+      }
+
+      const stats = writer.getColumnStats()
+
+      expect(stats.get('new_row')?.nullCount).toBe(5)
+    })
+
+    it('should compute min/max for numeric columns', () => {
+      const baseTime = 1704067200000
+      for (let i = 0; i < 10; i++) {
+        writer.addRecord(createRecord({ timestamp: baseTime + i * 1000 }))
+      }
+
+      const stats = writer.getColumnStats()
+
+      expect(stats.get('timestamp')?.minValue).toBe(baseTime)
+      expect(stats.get('timestamp')?.maxValue).toBe(baseTime + 9000)
+    })
+
+    it('should compute min/max for string columns', () => {
+      writer.addRecord(createRecord({ do_id: 'aaa' }))
+      writer.addRecord(createRecord({ do_id: 'zzz' }))
+      writer.addRecord(createRecord({ do_id: 'mmm' }))
+
+      const stats = writer.getColumnStats()
+
+      expect(stats.get('do_id')?.minValue).toBe('aaa')
+      expect(stats.get('do_id')?.maxValue).toBe('zzz')
+    })
+  })
+
+  describe('getBloomFilters', () => {
+    it('should return empty map before build', () => {
+      const filters = writer.getBloomFilters()
+
+      expect(filters.size).toBe(0)
+    })
+
+    it('should have filters after build', () => {
+      writer.addRecord(createRecord())
+      writer.build()
+
+      const filters = writer.getBloomFilters()
+
+      expect(filters.size).toBeGreaterThan(0)
+    })
+  })
+
+  describe('bloomFilterMightContain', () => {
+    it('should return null for column without bloom filter', () => {
+      writer.addRecord(createRecord())
+      writer.build()
+
+      const result = writer.bloomFilterMightContain('timestamp', 1704067200000)
+
+      // timestamp doesn't have bloom filter enabled
+      expect(result === null || result === true).toBe(true)
+    })
+
+    it('should return true for existing value', () => {
+      writer.addRecord(createRecord({ do_id: 'test-do-123' }))
+      writer.build()
+
+      const result = writer.bloomFilterMightContain('do_id', 'test-do-123')
+
+      expect(result).toBe(true)
+    })
+
+    it('should return false for definitely non-existing value', () => {
+      writer.addRecord(createRecord({ do_id: 'test-do-123' }))
+      writer.build()
+
+      // Bloom filter might have false positives but not false negatives
+      const result = writer.bloomFilterMightContain('do_id', 'definitely-not-exists-xyz')
+
+      // Could be true (false positive) or false (correctly rejected)
+      expect(typeof result).toBe('boolean')
+    })
+  })
+
+  describe('build', () => {
+    it('should throw error on empty writer', () => {
+      expect(() => writer.build()).toThrow('Cannot build Parquet file with no records')
+    })
+
+    it('should produce valid Parquet bytes with magic header', () => {
+      writer.addRecord(createRecord())
+
+      const bytes = writer.build()
+
+      expect(bytes).toBeInstanceOf(Uint8Array)
+      // PAR1 magic bytes
+      expect(bytes[0]).toBe(0x50) // P
+      expect(bytes[1]).toBe(0x41) // A
+      expect(bytes[2]).toBe(0x52) // R
+      expect(bytes[3]).toBe(0x31) // 1
+    })
+
+    it('should have PAR1 magic at end', () => {
+      writer.addRecord(createRecord())
+
+      const bytes = writer.build()
+
+      const len = bytes.length
+      expect(bytes[len - 4]).toBe(0x50) // P
+      expect(bytes[len - 3]).toBe(0x41) // A
+      expect(bytes[len - 2]).toBe(0x52) // R
+      expect(bytes[len - 1]).toBe(0x31) // 1
+    })
+  })
+
+  describe('clear', () => {
+    it('should reset record count', () => {
+      writer.addRecord(createRecord())
+      expect(writer.getRecordCount()).toBe(1)
+
+      writer.clear()
+
+      expect(writer.getRecordCount()).toBe(0)
+    })
+
+    it('should reset row group size', () => {
+      writer.addRecord(createRecord())
+
+      writer.clear()
+
+      expect(writer.getCurrentRowGroupSize()).toBe(0)
+    })
+
+    it('should clear bloom filters', () => {
+      writer.addRecord(createRecord())
+      writer.build()
+      expect(writer.getBloomFilters().size).toBeGreaterThan(0)
+
+      writer.clear()
+
+      expect(writer.getBloomFilters().size).toBe(0)
+    })
+  })
+
+  describe('getConfig', () => {
+    it('should return writer configuration', () => {
+      const config = writer.getConfig()
+
+      expect(config.compression).toBeDefined()
+      expect(config.targetRowGroupSize).toBeDefined()
+      expect(config.enableBloomFilters).toBeDefined()
+    })
+  })
+
+  describe('getCompressionRatioEstimate', () => {
+    it('should return estimate based on compression type', () => {
+      const estimate = writer.getCompressionRatioEstimate()
+
+      expect(estimate).toBeGreaterThan(1)
+    })
+  })
+})
+
+// ============================================================================
+// createWALParquetWriter Tests
+// ============================================================================
+
+describe('createWALParquetWriter', () => {
+  it('should create writer with default config', () => {
+    const writer = createWALParquetWriter()
+
+    expect(writer).toBeInstanceOf(ParquetWriter)
+  })
+
+  it('should create writer with custom metadata', () => {
+    const writer = createWALParquetWriter({ 'custom.key': 'value' })
+
+    expect(writer).toBeInstanceOf(ParquetWriter)
+  })
+
+  it('should create writer with custom config', () => {
+    const config: ParquetWriterConfig = {
+      compression: 'SNAPPY',
+      targetRowGroupSize: 64 * 1024 * 1024,
+    }
+    const writer = createWALParquetWriter(undefined, config)
+
+    expect(writer.getConfig().compression).toBe('SNAPPY')
+  })
+})
+
+// ============================================================================
+// generateDataFileMetadata Tests
+// ============================================================================
+
+describe('generateDataFileMetadata', () => {
+  it('should generate Iceberg data file metadata', () => {
+    const writer = createWALParquetWriter()
+    writer.addRecord({
+      do_id: 'test-do',
+      lsn: '1',
+      timestamp: 1704067200000,
+      operation: 'INSERT',
+      schema_name: 'public',
+      table_name: 'test',
+      primary_key: null,
+      new_row: '{}',
+      old_row: null,
+      changed_columns: null,
+      transaction_id: null,
+      sql: null,
+    })
+
+    const schema = createWALSchema()
+    const metadata = generateDataFileMetadata(
+      writer,
+      'data/test/file.parquet',
+      5000,
+      { do_id_partition: 'test-do' },
+      schema
+    )
+
+    expect(metadata['file-path']).toBe('data/test/file.parquet')
+    expect(metadata['file-format']).toBe('parquet')
+    expect(metadata['file-size-in-bytes']).toBe(5000)
+    expect(metadata['record-count']).toBe(1)
+    expect(metadata.content).toBe('data')
+    expect(metadata.partition).toEqual({ do_id_partition: 'test-do' })
+  })
+
+  it('should include column statistics', () => {
+    const writer = createWALParquetWriter()
+    for (let i = 0; i < 10; i++) {
+      writer.addRecord({
+        do_id: 'test-do',
+        lsn: String(i),
+        timestamp: 1704067200000 + i * 1000,
+        operation: 'INSERT',
+        schema_name: 'public',
+        table_name: 'test',
+        primary_key: null,
+        new_row: null,
+        old_row: null,
+        changed_columns: null,
+        transaction_id: null,
+        sql: null,
+      })
+    }
+
+    const schema = createWALSchema()
+    const metadata = generateDataFileMetadata(writer, 'file.parquet', 5000, {}, schema)
+
+    expect(metadata['value-counts']).toBeDefined()
+    expect(metadata['null-value-counts']).toBeDefined()
+    expect(metadata['lower-bounds']).toBeDefined()
+    expect(metadata['upper-bounds']).toBeDefined()
+  })
+})
+
+// ============================================================================
+// bloomFilterMightContain Tests
+// ============================================================================
+
+describe('bloomFilterMightContain', () => {
+  it('should return true for values in the filter', () => {
+    // Create a simple bloom filter manually
+    const filter: BloomFilter = {
+      bits: new Uint8Array(64),
+      numHashFunctions: 3,
+      numBits: 512,
+    }
+
+    // We can't easily test this without the internal filter creation
+    // This is more of a smoke test
+    const result = bloomFilterMightContain(filter, 'test')
+
+    expect(typeof result).toBe('boolean')
+  })
+
+  it('should handle string values', () => {
+    const filter: BloomFilter = {
+      bits: new Uint8Array(64),
+      numHashFunctions: 3,
+      numBits: 512,
+    }
+
+    const result = bloomFilterMightContain(filter, 'hello-world')
+
+    expect(typeof result).toBe('boolean')
+  })
+
+  it('should handle non-string values by JSON serialization', () => {
+    const filter: BloomFilter = {
+      bits: new Uint8Array(64),
+      numHashFunctions: 3,
+      numBits: 512,
+    }
+
+    const result = bloomFilterMightContain(filter, { id: 123 })
+
+    expect(typeof result).toBe('boolean')
+  })
+})
+
+// ============================================================================
+// WALParquetBatchWriter Tests
+// ============================================================================
+
+describe('WALParquetBatchWriter', () => {
+  let batchWriter: WALParquetBatchWriter
+
+  const createEntry = (overrides: Partial<SerializedWALEntry> = {}): SerializedWALEntry => ({
+    lsn: '1',
+    operation: 'INSERT',
+    schema: 'public',
+    table: 'test',
+    timestamp: Date.now(),
+    ...overrides,
+  })
+
+  beforeEach(() => {
+    batchWriter = createWALParquetBatchWriter()
+  })
+
+  describe('addEntry', () => {
+    it('should add entry to batch', async () => {
+      await batchWriter.addEntry(createEntry(), 'test-do')
+
+      const stats = batchWriter.getStats()
+      expect(stats.totalRecords).toBe(1)
+    })
+
+    it('should track partition count', async () => {
+      await batchWriter.addEntry(createEntry({ timestamp: 1704067200000 }), 'do-1')
+      await batchWriter.addEntry(createEntry({ timestamp: 1704153600000 }), 'do-2')
+
+      const stats = batchWriter.getStats()
+      expect(stats.partitionCount).toBeGreaterThanOrEqual(1)
+    })
+  })
+
+  describe('addEntries', () => {
+    it('should add multiple entries', async () => {
+      const entries = [createEntry({ lsn: '1' }), createEntry({ lsn: '2' }), createEntry({ lsn: '3' })]
+
+      await batchWriter.addEntries(entries, 'test-do')
+
+      const stats = batchWriter.getStats()
+      expect(stats.totalRecords).toBe(3)
+    })
+  })
+
+  describe('flush', () => {
+    it('should flush all partitions', async () => {
+      await batchWriter.addEntry(createEntry(), 'test-do')
+
+      const result = await batchWriter.flush()
+
+      expect(result.recordsWritten).toBe(1)
+      expect(result.filesCreated).toBeGreaterThanOrEqual(1)
+    })
+
+    it('should return batch write result', async () => {
+      await batchWriter.addEntry(createEntry(), 'test-do')
+
+      const result = await batchWriter.flush()
+
+      expect(result).toHaveProperty('recordsWritten')
+      expect(result).toHaveProperty('filesCreated')
+      expect(result).toHaveProperty('totalBytesUncompressed')
+      expect(result).toHaveProperty('totalBytesCompressed')
+      expect(result).toHaveProperty('compressionRatio')
+      expect(result).toHaveProperty('durationMs')
+      expect(result).toHaveProperty('filePaths')
+    })
+
+    it('should handle empty batch', async () => {
+      const result = await batchWriter.flush()
+
+      expect(result.recordsWritten).toBe(0)
+      expect(result.filesCreated).toBe(0)
+    })
+  })
+
+  describe('getStats', () => {
+    it('should return current statistics', () => {
+      const stats = batchWriter.getStats()
+
+      expect(stats).toHaveProperty('totalRecords')
+      expect(stats).toHaveProperty('partitionCount')
+      expect(stats).toHaveProperty('filesCreated')
+      expect(stats).toHaveProperty('estimatedCompressionRatio')
+    })
+  })
+
+  describe('reset', () => {
+    it('should clear all state', async () => {
+      await batchWriter.addEntry(createEntry(), 'test-do')
+
+      batchWriter.reset()
+
+      const stats = batchWriter.getStats()
+      expect(stats.totalRecords).toBe(0)
+      expect(stats.partitionCount).toBe(0)
+    })
+  })
+})
+
+// ============================================================================
+// createWALParquetBatchWriter Tests
+// ============================================================================
+
+describe('createWALParquetBatchWriter', () => {
+  it('should create batch writer with default config', () => {
+    const batchWriter = createWALParquetBatchWriter()
+
+    expect(batchWriter).toBeInstanceOf(WALParquetBatchWriter)
+  })
+
+  it('should create batch writer with custom config', () => {
+    const config: WALParquetBatchConfig = {
+      compression: 'GZIP',
+      maxRecordsPerFile: 1000,
+    }
+    const batchWriter = createWALParquetBatchWriter(config)
+
+    expect(batchWriter).toBeInstanceOf(WALParquetBatchWriter)
+  })
+
+  it('should use provided R2 bucket for storage', async () => {
+    const mockBucket = createMockBucket()
+    const batchWriter = createWALParquetBatchWriter({
+      bucket: mockBucket,
+      basePath: 'iceberg/wal/data',
+      compression: 'SNAPPY',
+    })
+
+    const entry: SerializedWALEntry = {
+      lsn: '1',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'test',
+      timestamp: Date.now(),
+    }
+
+    await batchWriter.addEntry(entry, 'test-do')
+    const result = await batchWriter.flush()
+
+    // Should have written to R2
+    expect(result.filesCreated).toBe(1)
+    expect(result.recordsWritten).toBe(1)
+    expect(mockBucket.put).toHaveBeenCalled()
+
+    // Verify the file path structure
+    const filePath = result.filePaths[0]
+    expect(filePath).toContain('iceberg/wal/data')
+    expect(filePath).toContain('.parquet')
+  })
+
+  it('should call onFlush callback when flushing', async () => {
+    const flushedFiles: Array<{ data: Uint8Array; partition: Record<string, unknown> }> = []
+
+    const batchWriter = createWALParquetBatchWriter({
+      compression: 'SNAPPY',
+      onFlush: async (data: Uint8Array, partition: Record<string, unknown>) => {
+        flushedFiles.push({ data, partition })
+        return `custom/path/${partition.do_id_partition}/${Date.now()}.parquet`
+      },
+    })
+
+    const entry: SerializedWALEntry = {
+      lsn: '1',
+      operation: 'INSERT',
+      schema: 'public',
+      table: 'test',
+      timestamp: Date.now(),
+    }
+
+    await batchWriter.addEntry(entry, 'test-do-123')
+    const result = await batchWriter.flush()
+
+    // Should have called onFlush callback
+    expect(flushedFiles.length).toBe(1)
+    expect(flushedFiles[0]?.partition).toHaveProperty('do_id_partition', 'test-do-123')
+
+    // The data should be valid Parquet (starts with PAR1 magic)
+    const data = flushedFiles[0]?.data
+    expect(data?.[0]).toBe(0x50) // P
+    expect(data?.[1]).toBe(0x41) // A
+    expect(data?.[2]).toBe(0x52) // R
+    expect(data?.[3]).toBe(0x31) // 1
+
+    // Result should have the custom path
+    expect(result.filePaths[0]).toContain('custom/path/test-do-123')
+  })
+})