hyparquet/src/plan.js

Parquet file parser for JavaScript

import { bloomEligibleColumns, readBloomFilter } from './bloom.js'
import { canSkipRowGroup, filterPageRanges, pathsNeededForFilter } from './filter.js'
import { readColumnIndex, readOffsetIndex } from './indexes.js'
import { parquetSchema } from './metadata.js'
import { getPhysicalColumns } from './schema.js'

/**
 * @import {AsyncBuffer, BloomFilter, ByteRange, ChunkPlan, ColumnPageStats, FileMetaData, GroupPlan, PageLocation, PageRanges, ParquetParsers, ParquetQueryFilter, ParquetReadOptions, QueryPlan, RowGroup, SchemaElement, SchemaTree} from '../src/types.js'
 */

// Combine column chunks if less than 2mb
const runLimit = 1 << 21 // 2mb
// When reading selected columns, combine chunks separated by at most this many unrequested bytes.
// Arrow C++ before 18 and Polars write a copy of each column's metadata after its chunk, leaving small gaps.
// Same default as Arrow's CacheOptions hole_size_limit.
const columnGapLimit = 1 << 13 // 8kb

/**
 * Plan which byte ranges to read to satisfy a read request.
 * Metadata must be non-null.
 *
 * @param {ParquetReadOptions & { bloomFiltersByGroup?: Record<string, BloomFilter>[], schemaElements?: Record<string, SchemaElement>, pageRangesByGroup?: (PageRanges | undefined)[], pageLocationsByGroup?: Record<string, PageLocation[]>[] }} options
 * @returns {QueryPlan}
 */
export function parquetPlan(options) {
  const { metadata, rowStart = 0, columns, useOffsetIndex = false } = options
  if (!metadata) throw new Error('parquetPlan requires metadata')
  /** @type {GroupPlan[]} */
  const groups = []
  /** @type {ByteRange[]} */
  const fetches = []
  /** @type {ByteRange[]} */
  const indexes = []
  const scanPlan = parquetPlanGroups(options)
  for (const group of scanPlan.groups) {
    const groupPlan = parquetPlanGroup({ ...group, columns, useOffsetIndex })
    groups.push(...groupPlan.groups)
    fetches.push(...groupPlan.fetches)
    indexes.push(...groupPlan.indexes)
  }
  fetches.push(...indexes)

  return { metadata, rowStart, rowEnd: scanPlan.rowEnd, columns, fetches, groups }
}

/**
 * Select physical row-group ranges without planning column reads.
 *
 * @param {ParquetReadOptions & { bloomFiltersByGroup?: Record<string, BloomFilter>[], schemaElements?: Record<string, SchemaElement>, pageRangesByGroup?: (PageRanges | undefined)[], pageLocationsByGroup?: Record<string, PageLocation[]>[] }} options
 * @returns {{groups: {rowGroup: RowGroup, groupIndex: number, groupStart: number, groupRows: number, ranges: PageRanges, pageRanges?: PageRanges, pageLocations?: Record<string, PageLocation[]>}[], rowEnd: number}}
 */
export function parquetPlanGroups({ metadata, rowStart = 0, rowEnd = Infinity, columns, filter, filterStrict = true, bloomFiltersByGroup, schemaElements, pageRangesByGroup, pageLocationsByGroup }) {
  if (!metadata) throw new Error('parquetPlan requires metadata')
  const schemaTree = parquetSchema(metadata)
  const physicalColumns = getPhysicalColumns(schemaTree)
  const elementsByPath = filter ? {
    ...physicalSchemaElements(schemaTree),
    ...schemaElements,
  } : schemaElements
  const groups = []
  let groupStart = 0
  for (let groupIndex = 0; groupIndex < metadata.row_groups.length; groupIndex++) {
    const rowGroup = metadata.row_groups[groupIndex]
    const groupRows = Number(rowGroup.num_rows)
    const groupEnd = groupStart + groupRows
    if (groupRows > 0 && groupEnd > rowStart && groupStart < rowEnd && !canSkipRowGroup({
      rowGroup,
      physicalColumns,
      filter,
      strict: filterStrict,
      bloomFilters: bloomFiltersByGroup?.[groupIndex],
      schemaElements: elementsByPath,
    })) {
      const selectStart = Math.max(rowStart - groupStart, 0)
      const selectEnd = Math.min(rowEnd - groupStart, groupRows)
      const pageRanges = pageRangesByGroup?.[groupIndex]
      const pageLocations = pageLocationsByGroup?.[groupIndex]
      /** @type {PageRanges} */
      let ranges = pageRanges
        ? pageRanges
          .map(([start, end]) => {
            /** @type {[number, number]} */
            const range = [Math.max(start, selectStart), Math.min(end, selectEnd)]
            return range
          })
          .filter(([start, end]) => start < end)
        : [[selectStart, selectEnd]]

      if (ranges.length > 1) {
        const canSplit = rowGroup.columns.every(chunk => {
          const columnName = chunk.meta_data?.path_in_schema[0]
          const columnPath = chunk.meta_data?.path_in_schema.join('.')
          if (columns && columnName && !columns.includes(columnName)) return true
          return !!(chunk.offset_index_offset && chunk.offset_index_length) || !!(columnPath && pageLocations?.[columnPath])
        })
        ranges = canSplit
          ? coalesceOverlappingPageRanges(ranges, rowGroup, columns, pageLocations)
          : [[ranges[0][0], ranges[ranges.length - 1][1]]]
      }
      if (ranges.length) {
        groups.push({ rowGroup, groupIndex, groupStart, groupRows, ranges, pageRanges, pageLocations })
      }
    }
    groupStart = groupEnd
  }
  return { groups, rowEnd: isFinite(rowEnd) ? rowEnd : groupStart }
}

/**
 * Build byte plans for retained ranges in one row group.
 *
 * @param {object} options
 * @param {RowGroup} options.rowGroup
 * @param {number} options.groupStart
 * @param {number} options.groupRows
 * @param {PageRanges} options.ranges
 * @param {string[]} [options.columns]
 * @param {boolean} [options.useOffsetIndex]
 * @param {PageRanges} [options.pageRanges]
 * @param {Record<string, PageLocation[]>} [options.pageLocations]
 * @returns {{groups: GroupPlan[], fetches: ByteRange[], indexes: ByteRange[]}}
 */
export function parquetPlanGroup({ rowGroup, groupStart, groupRows, ranges, columns, useOffsetIndex = false, pageRanges, pageLocations }) {
  /** @type {ChunkPlan[]} */
  const chunks = []
  /** @type {ByteRange[]} */
  const fetches = []
  /** @type {ByteRange[]} */
  const indexes = []
  const narrowed = ranges.length > 1 || ranges[0][0] > 0 || ranges[0][1] < groupRows
  for (const chunk of rowGroup.columns) {
    const meta = chunk.meta_data
    if (chunk.file_path) throw new Error('parquet file_path not supported')
    if (!meta) throw new Error('parquet column metadata is undefined')
    if (columns && !columns.includes(meta.path_in_schema[0])) continue
    const columnOffset = meta.dictionary_page_offset || meta.data_page_offset
    const startByte = Number(columnOffset)
    const endByte = Number(columnOffset + meta.total_compressed_size)
    const chunkPageLocations = pageLocations?.[meta.path_in_schema.join('.')]

    if (chunkPageLocations && narrowed) {
      chunks.push({ columnMetadata: meta, pageLocations: chunkPageLocations, range: { startByte, endByte } })
    } else if ((useOffsetIndex || pageRanges) && chunk.offset_index_offset && chunk.offset_index_length && narrowed) {
      const startByte = Number(chunk.offset_index_offset)
      chunks.push({
        columnMetadata: meta,
        offsetIndex: { startByte, endByte: startByte + chunk.offset_index_length },
        range: { startByte: Number(columnOffset), endByte },
      })
    } else {
      chunks.push({ columnMetadata: meta, range: { startByte, endByte } })
    }
  }

  /** @type {ByteRange[]} */
  const columnRanges = []
  /** @type {ByteRange | undefined} */
  let run
  for (const chunk of chunks) {
    if ('pageLocations' in chunk) continue
    if ('offsetIndex' in chunk) {
      indexes.push(chunk.offsetIndex)
    } else if (columns) {
      columnRanges.push(chunk.range)
    } else if (run && chunk.range.endByte - run.startByte <= runLimit) {
      run.endByte = chunk.range.endByte
    } else {
      if (run) fetches.push(run)
      run = { ...chunk.range }
    }
  }
  if (run) fetches.push(run)
  // selected chunks are combined in file order, skipping only small gaps
  fetches.push(...coalesceByteRanges(columnRanges, columnGapLimit, runLimit))
  const groups = ranges.map(([selectStart, selectEnd]) => ({
    chunks,
    rowGroup,
    groupStart,
    groupRows,
    selectStart,
    selectEnd,
  }))
  return { groups, fetches, indexes }
}

/**
 * Merge candidate ranges when any selected column would read an overlapping
 * data page for both ranges. Unknown page layouts are merged conservatively.
 *
 * @param {PageRanges} ranges
 * @param {RowGroup} rowGroup
 * @param {string[] | undefined} columns
 * @param {Record<string, PageLocation[]> | undefined} pageLocations
 * @returns {PageRanges}
 */
function coalesceOverlappingPageRanges(ranges, rowGroup, columns, pageLocations) {
  const selectedPageLayouts = rowGroup.columns
    .filter(chunk => !columns || columns.includes(chunk.meta_data?.path_in_schema[0] || ''))
    .map(chunk => pageLocations?.[chunk.meta_data?.path_in_schema.join('.') || ''])

  /** @type {PageRanges} */
  const merged = []
  for (const range of ranges) {
    const last = merged[merged.length - 1]
    const overlapsPage = last && selectedPageLayouts.some(pages => {
      if (!pages) return true
      const lastPages = pagesForRange(last, pages, Number(rowGroup.num_rows))
      const rangePages = pagesForRange(range, pages, Number(rowGroup.num_rows))
      return lastPages[0] <= rangePages[1] && rangePages[0] <= lastPages[1]
    })
    if (last && overlapsPage) last[1] = range[1]
    else merged.push([...range])
  }
  return merged
}

/**
 * Return the inclusive indexes of the first and last pages overlapping a row
 * range. Offset indexes cover every row, so a valid range always finds a page.
 *
 * @param {[number, number]} range
 * @param {PageLocation[]} pages
 * @param {number} groupRows
 * @returns {[number, number]}
 */
function pagesForRange([rangeStart, rangeEnd], pages, groupRows) {
  let first = Infinity
  let last = -Infinity
  for (let i = 0; i < pages.length; i++) {
    const pageStart = Number(pages[i].first_row_index)
    const pageEnd = i + 1 < pages.length ? Number(pages[i + 1].first_row_index) : groupRows
    if (pageEnd > rangeStart && pageStart < rangeEnd) {
      first = Math.min(first, i)
      last = i
    }
  }
  return [first, last]
}

/**
 * Fetch bloom filters for $eq / $in columns of row groups not already provably
 * skippable by statistics alone. Returns an array indexed by row-group ordinal;
 * each entry maps top-level column name → BloomFilter for any chunk whose
 * bloom filter we were able to parse. Adds one round-trip when at least one
 * bloom filter is fetched; otherwise returns synchronously.
 *
 * @param {object} options
 * @param {AsyncBuffer} options.file
 * @param {FileMetaData} options.metadata
 * @param {ParquetQueryFilter} options.filter
 * @param {boolean} [options.filterStrict]
 * @returns {Promise<Record<string, BloomFilter>[]>}
 */
export async function prefetchBloomFilters({ file, metadata, filter, filterStrict = true }) {
  const result = metadata.row_groups.map(() => /** @type {Record<string, BloomFilter>} */ ({}))
  const eligibleCols = bloomEligibleColumns(filter)
  if (eligibleCols.size === 0) return result
  const physicalColumns = getPhysicalColumns(parquetSchema(metadata))

  /** @type {Promise<void>[]} */
  const tasks = []
  metadata.row_groups.forEach((rowGroup, rgIdx) => {
    if (canSkipRowGroup({ rowGroup, physicalColumns, filter, strict: filterStrict })) return
    for (const colName of eligibleCols) {
      const columnIdx = physicalColumns.indexOf(colName)
      if (columnIdx === -1) continue
      const meta = rowGroup.columns[columnIdx]?.meta_data
      if (!meta?.bloom_filter_offset || !meta.bloom_filter_length) continue
      const start = Number(meta.bloom_filter_offset)
      const end = start + meta.bloom_filter_length
      tasks.push((async () => {
        const buffer = await file.slice(start, end)
        const bloom = readBloomFilter({ view: new DataView(buffer), offset: 0 })
        if (bloom) result[rgIdx][colName] = bloom
      })())
    }
  })

  if (tasks.length) await Promise.all(tasks)
  return result
}

/**
 * Fetch page indexes (column index + offset index) for filter columns of row
 * groups that survive row-group-level pruning, and compute candidate row
 * ranges per group from the per-page min/max statistics.
 *
 * Returns pageRangesByGroup indexed by row-group ordinal: sorted disjoint
 * [start, end) row ranges (relative to the group) that could match the filter.
 * An empty array means the group provably contains no matching rows; undefined
 * means no page-level information was available for that group.
 * Also returns pageLocationsByGroup, mapping row-group ordinals and physical
 * leaf paths to page locations, so the read path can reuse the parsed offset
 * indexes without refetching them.
 *
 * @param {object} options
 * @param {AsyncBuffer} options.file
 * @param {FileMetaData} options.metadata
 * @param {ParquetQueryFilter} [options.filter]
 * @param {boolean} [options.filterStrict]
 * @param {number} [options.rowStart]
 * @param {number} [options.rowEnd]
 * @param {string[]} [options.columns]
 * @param {Record<string, BloomFilter>[]} [options.bloomFiltersByGroup]
 * @param {Record<string, SchemaElement>} [options.schemaElements]
 * @param {Partial<ParquetParsers>} [options.parsers]
 * @returns {Promise<{pageRangesByGroup: (PageRanges | undefined)[], pageLocationsByGroup: Record<string, PageLocation[]>[]}>}
 */
export async function prefetchPageIndexes({ file, metadata, filter, filterStrict = true, rowStart = 0, rowEnd = Infinity, columns, bloomFiltersByGroup, schemaElements, parsers }) {
  /** @type {(PageRanges | undefined)[]} */
  const pageRangesByGroup = metadata.row_groups.map(() => undefined)
  const pageLocationsByGroup = metadata.row_groups.map(() => /** @type {Record<string, PageLocation[]>} */ ({}))
  if (filter && '$nor' in filter && Array.isArray(filter.$nor)) {
    return { pageRangesByGroup, pageLocationsByGroup }
  }
  const filterColumns = pathsNeededForFilter(filter)
  if (!filterColumns.length) return { pageRangesByGroup, pageLocationsByGroup }
  const schemaTree = parquetSchema(metadata)
  const physicalColumns = getPhysicalColumns(schemaTree)
  const elementsByPath = {
    ...physicalSchemaElements(schemaTree),
    ...schemaElements,
  }

  /** @type {ByteRange[]} */
  const indexRanges = []
  /** @type {((file: AsyncBuffer) => Promise<void>)[]} */
  const indexTasks = []
  /** @type {{rgIdx: number, groupRows: number, columnPages: Record<string, ColumnPageStats>}[]} */
  const candidateGroups = []
  let groupStart = 0
  metadata.row_groups.forEach((rowGroup, rgIdx) => {
    const groupRows = Number(rowGroup.num_rows)
    const groupEnd = groupStart + groupRows
    const overlaps = groupRows > 0 && groupEnd > rowStart && groupStart < rowEnd
    groupStart = groupEnd
    if (!overlaps) return
    if (canSkipRowGroup({ rowGroup, physicalColumns, filter, strict: filterStrict, bloomFilters: bloomFiltersByGroup?.[rgIdx], schemaElements: elementsByPath })) return

    /** @type {Record<string, ColumnPageStats>} */
    const columnPages = {}
    let columnTaskCount = 0
    const scheduledOffsetPaths = new Set()
    let hasFilterIndex = false
    for (const columnName of filterColumns) {
      const columnIdx = physicalColumns.indexOf(columnName)
      if (columnIdx === -1) continue
      const chunk = rowGroup.columns[columnIdx]
      const meta = chunk?.meta_data
      if (!meta) continue
      if (!chunk.column_index_offset || !chunk.column_index_length) continue
      if (!chunk.offset_index_offset || !chunk.offset_index_length) continue
      const element = elementsByPath[columnName]
      if (!element) continue
      hasFilterIndex = true
      scheduledOffsetPaths.add(columnName)
      const columnIndexStart = Number(chunk.column_index_offset)
      const offsetIndexStart = Number(chunk.offset_index_offset)
      const columnIndexEnd = columnIndexStart + chunk.column_index_length
      const offsetIndexEnd = offsetIndexStart + chunk.offset_index_length
      indexRanges.push(
        { startByte: columnIndexStart, endByte: columnIndexEnd },
        { startByte: offsetIndexStart, endByte: offsetIndexEnd }
      )
      columnTaskCount++
      indexTasks.push(async prefetchedFile => {
        const [columnIndexBuffer, offsetIndexBuffer] = await Promise.all([
          prefetchedFile.slice(columnIndexStart, columnIndexEnd),
          prefetchedFile.slice(offsetIndexStart, offsetIndexEnd),
        ])
        const columnIndex = readColumnIndex({ view: new DataView(columnIndexBuffer), offset: 0 }, element, parsers)
        const offsetIndex = readOffsetIndex({ view: new DataView(offsetIndexBuffer), offset: 0 })
        pageLocationsByGroup[rgIdx][columnName] = offsetIndex.page_locations
        columnPages[columnName] = {
          minValues: columnIndex.min_values,
          maxValues: columnIndex.max_values,
          nullPages: columnIndex.null_pages,
          nullCounts: columnIndex.null_counts,
          pageStarts: offsetIndex.page_locations.map(page => Number(page.first_row_index)),
          element,
        }
      })
    }
    // Candidate ranges are shared by every selected column. Load their offset
    // indexes too so the planner can merge ranges that map to the same coarse
    // page instead of fetching and decoding that page more than once.
    if (hasFilterIndex) {
      for (const chunk of rowGroup.columns) {
        const meta = chunk.meta_data
        if (!meta) continue
        const columnName = meta.path_in_schema[0]
        const columnPath = meta.path_in_schema.join('.')
        if (columns && !columns.includes(columnName)) continue
        if (scheduledOffsetPaths.has(columnPath)) continue
        if (!chunk.offset_index_offset || !chunk.offset_index_length) continue
        scheduledOffsetPaths.add(columnPath)
        const offsetIndexStart = Number(chunk.offset_index_offset)
        const offsetIndexEnd = offsetIndexStart + chunk.offset_index_length
        indexRanges.push({ startByte: offsetIndexStart, endByte: offsetIndexEnd })
        columnTaskCount++
        indexTasks.push(async prefetchedFile => {
          const offsetIndexBuffer = await prefetchedFile.slice(offsetIndexStart, offsetIndexEnd)
          const offsetIndex = readOffsetIndex({ view: new DataView(offsetIndexBuffer), offset: 0 })
          pageLocationsByGroup[rgIdx][columnPath] = offsetIndex.page_locations
        })
      }
    }
    if (columnTaskCount) {
      candidateGroups.push({ rgIdx, groupRows, columnPages })
    }
  })

  if (indexTasks.length) {
    const prefetchedFile = prefetchAsyncBuffer(file, { fetches: coalesceByteRanges(indexRanges) })
    await Promise.all(indexTasks.map(task => task(prefetchedFile)))
    for (const { rgIdx, groupRows, columnPages } of candidateGroups) {
      pageRangesByGroup[rgIdx] = filterPageRanges(filter, columnPages, groupRows, filterStrict)
    }
  }
  return { pageRangesByGroup, pageLocationsByGroup }
}

/**
 * Merge overlapping or touching byte ranges. By default no bytes are added;
 * pass maxGap to also merge ranges separated by up to maxGap unrequested bytes,
 * as long as the merged range stays within maxSize.
 *
 * @param {ByteRange[]} ranges
 * @param {number} [maxGap]
 * @param {number} [maxSize]
 * @returns {ByteRange[]}
 */
function coalesceByteRanges(ranges, maxGap = 0, maxSize = Infinity) {
  const sorted = ranges
    .map(range => ({ ...range }))
    .sort((a, b) => a.startByte - b.startByte || a.endByte - b.endByte)
  /** @type {ByteRange[]} */
  const merged = []
  for (const range of sorted) {
    const last = merged[merged.length - 1]
    if (last && range.startByte <= last.endByte + maxGap && Math.max(last.endByte, range.endByte) - last.startByte <= maxSize) {
      last.endByte = Math.max(last.endByte, range.endByte)
    } else {
      merged.push(range)
    }
  }
  return merged
}

/**
 * Build a lookup from physical paths to leaf schema elements.
 *
 * @param {SchemaTree} schemaTree
 * @returns {Record<string, SchemaElement>}
 */
function physicalSchemaElements(schemaTree) {
  /** @type {Record<string, SchemaElement>} */
  const elements = {}
  /** @param {SchemaTree} node */
  function traverse(node) {
    if (node.children.length) {
      for (const child of node.children) traverse(child)
    } else {
      elements[node.path.join('.')] = node.element
    }
  }
  traverse(schemaTree)
  return elements
}

/**
 * Prefetch byte ranges from an AsyncBuffer.
 *
 * @param {AsyncBuffer} file
 * @param {{fetches: ByteRange[]}} options
 * @returns {AsyncBuffer}
 */
export function prefetchAsyncBuffer(file, { fetches }) {
  // fetch byte ranges from the file
  const promises = fetches.map(({ startByte, endByte }) => file.slice(startByte, endByte))
  return {
    byteLength: file.byteLength,
    slice(start, end = file.byteLength) {
      // find matching slice
      const index = fetches.findIndex(({ startByte, endByte }) => startByte <= start && end <= endByte)
      if (index < 0) {
        // fallback to direct read
        return file.slice(start, end)
      }
      if (fetches[index].startByte !== start || fetches[index].endByte !== end) {
        // slice a subrange of the prefetch
        const startOffset = start - fetches[index].startByte
        const endOffset = end - fetches[index].startByte
        if (promises[index] instanceof Promise) {
          return promises[index].then(buffer => buffer.slice(startOffset, endOffset))
        } else {
          return promises[index].slice(startOffset, endOffset)
        }
      } else {
        return promises[index]
      }
    },
  }
}