feat: optionally expand MERGED_CELL blocks from table analysis (#66)

* feat: optionally expand MERGED_CELL blocks from table analysis

* Minify test json

src/Document.ts

@@ -28,13 +28,15 @@ export class Document {
   pages: Page[];
   blocks: BlockStruct[];
   blockMap: BlockMap;
+  options: { expandMergedTableCells?: boolean } | undefined;
 
   /**
    * @param blocks - Block objects as returned by Textract
    * @throws {ParseError} - If the blocks are not valid
    * @throws {UnknownError} - If an unknown error occurs
    */
-  constructor(blocks: BlockStruct[]) {
+  constructor(blocks: BlockStruct[], options?: { expandMergedTableCells?: boolean }) {
+    this.options = options;
     try {
       const ret = BlockStructSchema.array().safeParse(blocks);
       if (!ret.success) {
@@ -74,14 +76,14 @@ export class Document {
     for (const b of this.blocks) {
       if (b.BlockType === 'PAGE') {
         if (blocks.length > 0) {
-          this.pages.push(new Page(blocks, this.blockMap));
+          this.pages.push(new Page(blocks, this.blockMap, this.options));
         }
         blocks = [];
       }
       blocks.push(b);
     }
     if (blocks.length > 0) {
-      this.pages.push(new Page(blocks, this.blockMap));
+      this.pages.push(new Page(blocks, this.blockMap, this.options));
     }
   }
 

src/Page.ts

@@ -14,8 +14,14 @@ export class Page {
   geometry?: Geometry;
   id: string;
   content: (Line | Table | Form | Field)[];
+  options: { expandMergedTableCells?: boolean } | undefined;
 
-  constructor(blocks: BlockStruct[], blockMap: BlockMap) {
+  constructor(
+    blocks: BlockStruct[],
+    blockMap: BlockMap,
+    options?: { expandMergedTableCells?: boolean }
+  ) {
+    this.options = options;
     this.blocks = blocks;
     this.text = '';
     this.lines = [];
@@ -48,7 +54,7 @@ export class Page {
           break;
         }
         case 'TABLE': {
-          const tbl = new Table(b, blockMap);
+          const tbl = new Table(b, blockMap, this.options);
           this.tables.push(tbl);
           this.content.push(tbl);
           break;

src/Table/Table.ts

@@ -2,16 +2,19 @@ import { Geometry } from '../Geometry/index.js';
 import { Cell } from './Cell.js';
 import { Row } from './Row.js';
 
-import type { BlockMap, CellBlock, TableBlock } from '../BlockStruct.js';
+import type { BlockMap, CellBlock, MergedCellBlock, TableBlock } from '../BlockStruct.js';
 
 export class Table {
   block: TableBlock;
   confidence: number;
   geometry: Geometry;
   id: string;
   rows: Row[];
-
-  constructor(block: TableBlock, blockMap: BlockMap) {
+  constructor(
+    block: TableBlock,
+    blockMap: BlockMap,
+    options?: { expandMergedTableCells?: boolean }
+  ) {
     this.block = block;
     this.confidence = block.Confidence;
     this.geometry = new Geometry(block.Geometry);
@@ -39,6 +42,46 @@ export class Table {
         }
       }
     }
+
+    if (options?.expandMergedTableCells) {
+      this.expandMergedCells(block, blockMap);
+    }
+  }
+
+  expandMergedCells(block: TableBlock, blockMap: BlockMap) {
+    const cellIDToContent = new Map();
+
+    const mergedCells = block.Relationships.filter(({ Type }) => Type === 'MERGED_CELL').flatMap(
+      ({ Ids }) =>
+        Ids.map((id) => blockMap.get(id)).filter(
+          (block): block is MergedCellBlock => block?.BlockType === 'MERGED_CELL'
+        )
+    );
+
+    for (const mergedCell of mergedCells) {
+      const cells = mergedCell.Relationships?.filter((rs) => rs.Type === 'CHILD')
+        .flatMap((rs) => rs.Ids.map((id) => blockMap.get(id)))
+        .filter((block): block is CellBlock => block?.BlockType === 'CELL');
+      if (!cells) {
+        continue;
+      }
+      const mergedContent: string[] = [];
+      for (const cell of cells) {
+        mergedContent.push(new Cell(cell, blockMap).text);
+      }
+      for (const cell of cells) {
+        cellIDToContent.set(cell.Id, mergedContent.join(' '));
+      }
+    }
+
+    for (const row of this.rows) {
+      for (const cell of row.cells) {
+        const overrideContent = cellIDToContent.get(cell.id);
+        if (overrideContent) {
+          cell.text = overrideContent;
+        }
+      }
+    }
   }
 
   toString() {

tests/all.test.ts

@@ -40,3 +40,23 @@ describe('Rows from table created correctly', () => {
     }
   });
 });
+
+describe('MERGED_CELL expansion', () => {
+  it('ignores MERGED_CELL blocks by default', async () => {
+    const expected = await readFile(join(root, 'merged', 'withoutExpansion.txt'), 'utf8');
+    const blocks = JSON.parse(await readFile(join(root, 'merged', 'blocks.json'), 'utf8'));
+
+    const doc = new Document(blocks);
+
+    expect(doc.pages?.[0]?.tables?.[0]?.toString()).toEqual(expected);
+  });
+
+  it('merges content of cells with `expandMergedTableCells` option', async () => {
+    const expected = await readFile(join(root, 'merged', 'withExpansion.txt'), 'utf8');
+    const blocks = JSON.parse(await readFile(join(root, 'merged', 'blocks.json'), 'utf8'));
+
+    const doc = new Document(blocks, { expandMergedTableCells: true });
+
+    expect(doc.pages?.[0]?.tables?.[0]?.toString()).toEqual(expected);
+  });
+});

tests/merged/withExpansion.txt

@@ -0,0 +1,17 @@
+Table
+==========
+Row
+==========
+[][Old][Old][New][New]
+Row
+==========
+[][Revenue][Profit][Revenue][Profit]
+Row
+==========
+[January][1][2][3][4]
+Row
+==========
+[February][5][6][7][8]
+Row
+==========
+[February][9][10][11][12]

tests/merged/withoutExpansion.txt

@@ -0,0 +1,17 @@
+Table
+==========
+Row
+==========
+[][Old][][New][]
+Row
+==========
+[][Revenue][Profit][Revenue][Profit]
+Row
+==========
+[January][1][2][3][4]
+Row
+==========
+[February][5][6][7][8]
+Row
+==========
+[][9][10][11][12]