Skip to content

Commit 0620905

Browse files
platypiiplatypii
authored
Defer parquet scan column planning (#173)
1 parent eb6eb73 commit 0620905

3 files changed

Lines changed: 181 additions & 156 deletions

File tree

‎src/plan.js‎

Lines changed: 116 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -18,39 +18,61 @@ const runLimit = 1 << 21 // 2mb
1818
* @param {ParquetReadOptions & { bloomFiltersByGroup?: Record<string, BloomFilter>[], schemaElements?: Record<string, SchemaElement>, pageRangesByGroup?: (PageRanges | undefined)[], pageLocationsByGroup?: Record<string, PageLocation[]>[] }} options
1919
* @returns {QueryPlan}
2020
*/
21-
export function parquetPlan({ metadata, rowStart = 0, rowEnd = Infinity, columns, filter, filterStrict = true, useOffsetIndex = false, bloomFiltersByGroup, schemaElements, pageRangesByGroup, pageLocationsByGroup }) {
21+
export function parquetPlan(options) {
22+
const { metadata, rowStart = 0, columns, useOffsetIndex = false } = options
2223
if (!metadata) throw new Error('parquetPlan requires metadata')
2324
/** @type {GroupPlan[]} */
2425
const groups = []
2526
/** @type {ByteRange[]} */
2627
const fetches = []
2728
/** @type {ByteRange[]} */
2829
const indexes = []
30+
const scanPlan = parquetPlanGroups(options)
31+
for (const group of scanPlan.groups) {
32+
const groupPlan = parquetPlanGroup({ ...group, columns, useOffsetIndex })
33+
groups.push(...groupPlan.groups)
34+
fetches.push(...groupPlan.fetches)
35+
indexes.push(...groupPlan.indexes)
36+
}
37+
fetches.push(...indexes)
38+
39+
return { metadata, rowStart, rowEnd: scanPlan.rowEnd, columns, fetches, groups }
40+
}
41+
42+
/**
43+
* Select physical row-group ranges without planning column reads.
44+
*
45+
* @param {ParquetReadOptions & { bloomFiltersByGroup?: Record<string, BloomFilter>[], schemaElements?: Record<string, SchemaElement>, pageRangesByGroup?: (PageRanges | undefined)[], pageLocationsByGroup?: Record<string, PageLocation[]>[] }} options
46+
* @returns {{groups: {rowGroup: RowGroup, groupIndex: number, groupStart: number, groupRows: number, ranges: PageRanges, pageRanges?: PageRanges, pageLocations?: Record<string, PageLocation[]>}[], rowEnd: number}}
47+
*/
48+
export function parquetPlanGroups({ metadata, rowStart = 0, rowEnd = Infinity, columns, filter, filterStrict = true, bloomFiltersByGroup, schemaElements, pageRangesByGroup, pageLocationsByGroup }) {
49+
if (!metadata) throw new Error('parquetPlan requires metadata')
2950
const schemaTree = parquetSchema(metadata)
3051
const physicalColumns = getPhysicalColumns(schemaTree)
3152
const elementsByPath = filter ? {
3253
...physicalSchemaElements(schemaTree),
3354
...schemaElements,
3455
} : schemaElements
35-
36-
// find which row groups to read
37-
let groupStart = 0 // first row index of the current group
38-
let rgIdx = 0
39-
for (const rowGroup of metadata.row_groups) {
56+
const groups = []
57+
let groupStart = 0
58+
for (let groupIndex = 0; groupIndex < metadata.row_groups.length; groupIndex++) {
59+
const rowGroup = metadata.row_groups[groupIndex]
4060
const groupRows = Number(rowGroup.num_rows)
4161
const groupEnd = groupStart + groupRows
42-
const bloomFilters = bloomFiltersByGroup?.[rgIdx]
43-
// if row group overlaps with row range, add it to the plan
44-
if (groupRows > 0 && groupEnd > rowStart && groupStart < rowEnd && !canSkipRowGroup({ rowGroup, physicalColumns, filter, strict: filterStrict, bloomFilters, schemaElements: elementsByPath })) {
62+
if (groupRows > 0 && groupEnd > rowStart && groupStart < rowEnd && !canSkipRowGroup({
63+
rowGroup,
64+
physicalColumns,
65+
filter,
66+
strict: filterStrict,
67+
bloomFilters: bloomFiltersByGroup?.[groupIndex],
68+
schemaElements: elementsByPath,
69+
})) {
4570
const selectStart = Math.max(rowStart - groupStart, 0)
4671
const selectEnd = Math.min(rowEnd - groupStart, groupRows)
47-
48-
// page-level pruning: split the group selection into candidate sub-ranges.
49-
// an empty list of sub-ranges skips the group entirely.
50-
const pageRanges = pageRangesByGroup?.[rgIdx]
51-
const pageLocations = pageLocationsByGroup?.[rgIdx]
72+
const pageRanges = pageRangesByGroup?.[groupIndex]
73+
const pageLocations = pageLocationsByGroup?.[groupIndex]
5274
/** @type {PageRanges} */
53-
let subranges = pageRanges
75+
let ranges = pageRanges
5476
? pageRanges
5577
.map(([start, end]) => {
5678
/** @type {[number, number]} */
@@ -60,105 +82,97 @@ export function parquetPlan({ metadata, rowStart = 0, rowEnd = Infinity, columns
6082
.filter(([start, end]) => start < end)
6183
: [[selectStart, selectEnd]]
6284

63-
// splitting requires page reads for every included chunk, or full chunks
64-
// would be fetched once per sub-range; collapse to one covering range otherwise
65-
if (subranges.length > 1) {
85+
if (ranges.length > 1) {
6686
const canSplit = rowGroup.columns.every(chunk => {
6787
const columnName = chunk.meta_data?.path_in_schema[0]
6888
const columnPath = chunk.meta_data?.path_in_schema.join('.')
6989
if (columns && columnName && !columns.includes(columnName)) return true
7090
return !!(chunk.offset_index_offset && chunk.offset_index_length) || !!(columnPath && pageLocations?.[columnPath])
7191
})
72-
if (!canSplit) {
73-
subranges = [[subranges[0][0], subranges[subranges.length - 1][1]]]
74-
} else {
75-
subranges = coalesceOverlappingPageRanges(subranges, rowGroup, columns, pageLocations)
76-
}
92+
ranges = canSplit
93+
? coalesceOverlappingPageRanges(ranges, rowGroup, columns, pageLocations)
94+
: [[ranges[0][0], ranges[ranges.length - 1][1]]]
7795
}
78-
79-
if (subranges.length) {
80-
/** @type {ChunkPlan[]} */
81-
const chunks = []
82-
// Multiple sub-ranges are necessarily narrower than the whole group.
83-
const narrowed = subranges.length > 1 ||
84-
subranges[0][0] > 0 || subranges[0][1] < groupRows
85-
// loop through each column chunk
86-
for (const chunk of rowGroup.columns) {
87-
const meta = chunk.meta_data
88-
if (chunk.file_path) throw new Error('parquet file_path not supported')
89-
if (!meta) throw new Error('parquet column metadata is undefined')
90-
// add included column chunks to the plan
91-
if (!columns || columns.includes(meta.path_in_schema[0])) {
92-
// full column chunk
93-
const columnOffset = meta.dictionary_page_offset || meta.data_page_offset
94-
const startByte = Number(columnOffset)
95-
const endByte = Number(columnOffset + meta.total_compressed_size)
96-
const chunkPageLocations = pageLocations?.[meta.path_in_schema.join('.')]
97-
98-
if (chunkPageLocations && narrowed) {
99-
// page locations already parsed during page index prefetch
100-
chunks.push({
101-
columnMetadata: meta,
102-
pageLocations: chunkPageLocations,
103-
range: { startByte, endByte },
104-
})
105-
} else if ((useOffsetIndex || pageRanges) && chunk.offset_index_offset && chunk.offset_index_length && narrowed) {
106-
const offsetIndexStart = Number(chunk.offset_index_offset)
107-
chunks.push({
108-
columnMetadata: meta,
109-
offsetIndex: {
110-
startByte: offsetIndexStart,
111-
endByte: offsetIndexStart + chunk.offset_index_length,
112-
},
113-
range: { startByte, endByte },
114-
})
115-
} else {
116-
chunks.push({
117-
columnMetadata: meta,
118-
range: { startByte, endByte },
119-
})
120-
}
121-
122-
}
123-
}
124-
125-
for (const [subStart, subEnd] of subranges) {
126-
groups.push({ chunks, rowGroup, groupStart, groupRows, selectStart: subStart, selectEnd: subEnd })
127-
}
128-
129-
// combine runs of column chunks
130-
/** @type {ByteRange | undefined} */
131-
let run
132-
for (const chunk of chunks) {
133-
if ('pageLocations' in chunk) {
134-
// pages are fetched on demand in readRowGroup
135-
} else if ('offsetIndex' in chunk) {
136-
indexes.push(chunk.offsetIndex)
137-
} else {
138-
const { range } = chunk
139-
if (columns) {
140-
fetches.push(range)
141-
} else if (run && range.endByte - run.startByte <= runLimit) {
142-
// extend range
143-
run.endByte = range.endByte
144-
} else {
145-
// new range
146-
if (run) fetches.push(run)
147-
run = { ...range }
148-
}
149-
}
150-
}
151-
if (run) fetches.push(run)
96+
if (ranges.length) {
97+
groups.push({ rowGroup, groupIndex, groupStart, groupRows, ranges, pageRanges, pageLocations })
15298
}
15399
}
154-
155100
groupStart = groupEnd
156-
rgIdx++
157101
}
158-
if (!isFinite(rowEnd)) rowEnd = groupStart
159-
fetches.push(...indexes)
102+
return { groups, rowEnd: isFinite(rowEnd) ? rowEnd : groupStart }
103+
}
104+
105+
/**
106+
* Build byte plans for retained ranges in one row group.
107+
*
108+
* @param {object} options
109+
* @param {RowGroup} options.rowGroup
110+
* @param {number} options.groupStart
111+
* @param {number} options.groupRows
112+
* @param {PageRanges} options.ranges
113+
* @param {string[]} [options.columns]
114+
* @param {boolean} [options.useOffsetIndex]
115+
* @param {PageRanges} [options.pageRanges]
116+
* @param {Record<string, PageLocation[]>} [options.pageLocations]
117+
* @returns {{groups: GroupPlan[], fetches: ByteRange[], indexes: ByteRange[]}}
118+
*/
119+
export function parquetPlanGroup({ rowGroup, groupStart, groupRows, ranges, columns, useOffsetIndex = false, pageRanges, pageLocations }) {
120+
/** @type {ChunkPlan[]} */
121+
const chunks = []
122+
/** @type {ByteRange[]} */
123+
const fetches = []
124+
/** @type {ByteRange[]} */
125+
const indexes = []
126+
const narrowed = ranges.length > 1 || ranges[0][0] > 0 || ranges[0][1] < groupRows
127+
for (const chunk of rowGroup.columns) {
128+
const meta = chunk.meta_data
129+
if (chunk.file_path) throw new Error('parquet file_path not supported')
130+
if (!meta) throw new Error('parquet column metadata is undefined')
131+
if (columns && !columns.includes(meta.path_in_schema[0])) continue
132+
const columnOffset = meta.dictionary_page_offset || meta.data_page_offset
133+
const startByte = Number(columnOffset)
134+
const endByte = Number(columnOffset + meta.total_compressed_size)
135+
const chunkPageLocations = pageLocations?.[meta.path_in_schema.join('.')]
160136

161-
return { metadata, rowStart, rowEnd, columns, fetches, groups }
137+
if (chunkPageLocations && narrowed) {
138+
chunks.push({ columnMetadata: meta, pageLocations: chunkPageLocations, range: { startByte, endByte } })
139+
} else if ((useOffsetIndex || pageRanges) && chunk.offset_index_offset && chunk.offset_index_length && narrowed) {
140+
const startByte = Number(chunk.offset_index_offset)
141+
chunks.push({
142+
columnMetadata: meta,
143+
offsetIndex: { startByte, endByte: startByte + chunk.offset_index_length },
144+
range: { startByte: Number(columnOffset), endByte },
145+
})
146+
} else {
147+
chunks.push({ columnMetadata: meta, range: { startByte, endByte } })
148+
}
149+
}
150+
151+
/** @type {ByteRange | undefined} */
152+
let run
153+
for (const chunk of chunks) {
154+
if ('pageLocations' in chunk) continue
155+
if ('offsetIndex' in chunk) {
156+
indexes.push(chunk.offsetIndex)
157+
} else if (columns) {
158+
fetches.push(chunk.range)
159+
} else if (run && chunk.range.endByte - run.startByte <= runLimit) {
160+
run.endByte = chunk.range.endByte
161+
} else {
162+
if (run) fetches.push(run)
163+
run = { ...chunk.range }
164+
}
165+
}
166+
if (run) fetches.push(run)
167+
const groups = ranges.map(([selectStart, selectEnd]) => ({
168+
chunks,
169+
rowGroup,
170+
groupStart,
171+
groupRows,
172+
selectStart,
173+
selectEnd,
174+
}))
175+
return { groups, fetches, indexes }
162176
}
163177

164178
/**

0 commit comments

Comments
 (0)