[query] Add vds.read_dense_mt to execute a single pass densify on read

[chrisvittal]

Change Description

Add read_dense_mt, a single pass vds densify method that reads a VDS as a dense matrix table. Using query_matrix_table_rows, we read the partitions of the variant data, alongside reading the reference data at the variant partition start less ref_block_max_len, making sure we have all the data that could produce a reference element for that partition.

Part of #14499

Security Assessment

Delete all except the correct answer:

• This change has no security impact

Impact Description

Query only

[chrisvittal]

The current error

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
Cell In[4], line 1
----> 1 dense.write('/tmp/dense.mt')

File <decorator-gen-1482>:2, in write(self, output, overwrite, stage_locally, _codec_spec, _partitions)

File ~/src/hail/hail/python/hail/typecheck/check.py:585, in _make_dec.<locals>.wrapper(__original_func, *args, **kwargs)
    582 @decorator
    583 def wrapper(__original_func: Callable[..., T], *args, **kwargs) -> T:
    584     args_, kwargs_ = check_all(__original_func, args, kwargs, checkers, is_method=is_method)
--> 585     return __original_func(*args_, **kwargs_)

File ~/src/hail/hail/python/hail/matrixtable.py:2810, in MatrixTable.write(self, output, overwrite, stage_locally, _codec_spec, _partitions)
   2807     _partitions_type = None
   2809 writer = ir.MatrixNativeWriter(output, overwrite, stage_locally, _codec_spec, _partitions, _partitions_type)
-> 2810 Env.backend().execute(ir.MatrixWrite(self._mir, writer))

File ~/src/hail/hail/python/hail/backend/spark_backend.py:217, in SparkBackend.execute(self, ir, timed)
    214     except Exception as fatal:
    215         raise err from fatal
--> 217 raise err

File ~/src/hail/hail/python/hail/backend/spark_backend.py:209, in SparkBackend.execute(self, ir, timed)
    207 def execute(self, ir: BaseIR, timed: bool = False) -> Any:
    208     try:
--> 209         return super().execute(ir, timed)
    210     except Exception as err:
    211         if self._copy_log_on_error:

File ~/src/hail/hail/python/hail/backend/backend.py:177, in Backend.execute(self, ir, timed)
    176 def execute(self, ir: BaseIR, timed: bool = False) -> Any:
--> 177     payload = ExecutePayload(self._render_ir(ir), '{"name":"StreamBufferSpec"}', timed)
    178     try:
    179         result, timings = self._rpc(ActionTag.EXECUTE, payload)

File ~/src/hail/hail/python/hail/backend/backend.py:194, in Backend._render_ir(self, ir)
    192 def _render_ir(self, ir):
    193     r = CSERenderer()
--> 194     return r(finalize_randomness(ir))

File ~/src/hail/hail/python/hail/ir/renderer.py:135, in CSERenderer.__call__(self, root)
    134 def __call__(self, root: 'ir.BaseIR') -> str:
--> 135     binding_sites = CSEAnalysisPass(self)(root)
    136     return CSEPrintPass(self)(root, binding_sites)

File ~/src/hail/hail/python/hail/ir/renderer.py:199, in CSEAnalysisPass.__call__(self, root)
    196 child_frame = frame.make_child_frame(len(stack))
    198 if isinstance(child, ir.IR):
--> 199     bind_depth = child_frame.bind_depth()
    200     lets = None
    201     if bind_depth < len(stack):

File ~/src/hail/hail/python/hail/ir/renderer.py:309, in CSEAnalysisPass.StackFrame.bind_depth(self)
    307 bind_depth = self.min_binding_depth
    308 if len(self.node.free_vars) > 0:
--> 309     bind_depth = max(bind_depth, *(self.context[0][var] for var in self.node.free_vars))
    310 if len(self.node.free_agg_vars) > 0:
    311     bind_depth = max(bind_depth, *(self.context[1][var] for var in self.node.free_agg_vars))

File ~/src/hail/hail/python/hail/ir/renderer.py:309, in <genexpr>(.0)
    307 bind_depth = self.min_binding_depth
    308 if len(self.node.free_vars) > 0:
--> 309     bind_depth = max(bind_depth, *(self.context[0][var] for var in self.node.free_vars))
    310 if len(self.node.free_agg_vars) > 0:
    311     bind_depth = max(bind_depth, *(self.context[1][var] for var in self.node.free_agg_vars))

KeyError: 'global'

[patrick-schultz]

Oh yay, it's cse. I'll take a look.

[patrick-schultz]

I've got it reproducing locally. Will dig in tomorrow morning. I always need to use the debugger to figure out cse issues.

[patrick-schultz]

Very nicely done. Just a couple small questions.

[patrick-schultz]

Is this testing anything not covered by the above equivalentce test?

[chrisvittal]

Not really.

[patrick-schultz]

What is this referring to?

[chrisvittal]

Like the _intervals argument for reads. Reading the data with a specified partitioning. It's something I want to implements, but haven't yet.

[patrick-schultz]

Ah, got it. Should we leave off the _intervals argument if it doesn't do anything yet?

[chrisvittal]

@grohli poke

[grohli]

Sorry about the delay on this--LGTM, thanks Chris!