Complete Fetch Phase (EXTERNAL_LINKS disposition and ARROW format)
#598
Conversation
Signed-off-by: varun-edachali-dbx <[email protected]>
covered by #588 Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
only relevant in Fetch phase Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
Signed-off-by: varun-edachali-dbx <[email protected]>
varun-edachali-dbx
had a problem deploying
to
azure-prod
GitHub Actions
Failure
| super().__init__( | ||
| max_download_threads=max_download_threads, | ||
| ssl_options=ssl_options, | ||
| schema_bytes=None, | ||
| lz4_compressed=lz4_compressed, | ||
| description=description, | ||
| ) | ||
|
|
||
| self._sea_client = sea_client | ||
| self._statement_id = statement_id | ||
| self._total_chunk_count = total_chunk_count | ||
|
|
||
| logger.debug( | ||
| "SeaCloudFetchQueue: Initialize CloudFetch loader for statement {}, total chunks: {}".format( | ||
| statement_id, total_chunk_count | ||
| ) | ||
| ) | ||
|
|
||
| initial_links = result_data.external_links or [] | ||
| first_link = next((l for l in initial_links if l.chunk_index == 0), None) | ||
| if not first_link: | ||
| # possibly an empty response | ||
| return None | ||
|
|
||
| # Track the current chunk we're processing | ||
| self._current_chunk_index = 0 | ||
| # Initialize table and position | ||
| self.table = self._create_table_from_link(first_link) | ||
|
|
||
| def _convert_to_thrift_link(self, link: ExternalLink) -> TSparkArrowResultLink: | ||
| """Convert SEA external links to Thrift format for compatibility with existing download manager.""" | ||
| # Parse the ISO format expiration time | ||
| expiry_time = int(dateutil.parser.parse(link.expiration).timestamp()) | ||
| return TSparkArrowResultLink( | ||
| fileLink=link.external_link, | ||
| expiryTime=expiry_time, | ||
| rowCount=link.row_count, | ||
| bytesNum=link.byte_count, | ||
| startRowOffset=link.row_offset, | ||
| httpHeaders=link.http_headers or {}, | ||
| ) | ||
|
|
||
| def _get_chunk_link(self, chunk_index: int) -> Optional[ExternalLink]: | ||
| """Progress to the next chunk link.""" | ||
| if chunk_index >= self._total_chunk_count: | ||
| return None | ||
|
|
||
| try: | ||
| return self._sea_client.get_chunk_link(self._statement_id, chunk_index) | ||
| except Exception as e: | ||
| raise ServerOperationError( | ||
| f"Error fetching link for chunk {chunk_index}: {e}", | ||
| { | ||
| "operation-id": self._statement_id, | ||
| "diagnostic-info": None, | ||
| }, | ||
| ) | ||
|
|
||
| def _create_table_from_link( | ||
| self, link: ExternalLink | ||
| ) -> Union["pyarrow.Table", None]: | ||
| """Create a table from a link.""" | ||
|
|
||
| thrift_link = self._convert_to_thrift_link(link) | ||
| self.download_manager.add_link(thrift_link) | ||
|
|
||
| row_offset = link.row_offset | ||
| arrow_table = self._create_table_at_offset(row_offset) | ||
|
|
||
| return arrow_table | ||
|
|
||
| def _create_next_table(self) -> Union["pyarrow.Table", None]: | ||
| """Create next table by retrieving the logical next downloaded file.""" | ||
| self._current_chunk_index += 1 | ||
| next_chunk_link = self._get_chunk_link(self._current_chunk_index) | ||
| if not next_chunk_link: | ||
| return None | ||
| return self._create_table_from_link(next_chunk_link) |
There was a problem hiding this comment.
Is this multi-threaded? _create_table_from_link adds link one-by-one. _create_next_table calls _create_table_from_link. This seems to be serial execution
There was a problem hiding this comment.
first you will consume arrow table from one link, then call create_next_table to get table from next link and so on. link gets added to queue once the previous link is consumed. isn't this wasteful?
There was a problem hiding this comment.
will this be solved in future PRs where you introduce a separate link downloader?
There was a problem hiding this comment.
Yes, this will be solved in the future PR where I introduce a separate link downloader. There, I use the multithreading capabilities of the ResultFileDownloadManager.
I intentionally kept this one quite simple and essentially single threaded so link expiry is not a concern, we download a link and create the arrow table the moment we fetch it. That way, this PR in isolation is correct.
varun-edachali-dbx
temporarily deployed
to
azure-prod
GitHub Actions
Inactive
| if not self.schema_bytes: | ||
| return pyarrow.Table.from_pydict({}) |
There was a problem hiding this comment.
why was this added to existing code? Is this just precautionary?
There was a problem hiding this comment.
Thrift provides us with arrow schema bytes, SEA does not. For the Thrift (empty) table, we can augment the column schema information into the pyarrow table, but for SEA, we return a completely empty table without this information.
This code essentially augments the schema information if it is available (which is the case for Thrift, but not SEA).
There was a problem hiding this comment.
Got it. For SEA we don't have easy way to put schema information from ResultManifest into an empty table?
There was a problem hiding this comment.
It is doable, but introduces some complexity - we would have to map each type given to us in the description (the 7 tuple containing name, type, etc.) to a corresponding pyarrow type (pa.int64(), pa.string(), etc.) and use that to define our pyarrow table schema, and then return an empty table. I do not see a way to do this without creating a manual mapping ("STRING" -> pa.string() etc.) so it could introduce some complexity and take some time.
If this is worth doing, I'll create a ticket and begin working on it soon.
What type of PR is this?
Description
Complete the Fetch phase of the SEA implementation for
EXTERNAL_LINKSdisposition andARROW_STREAMformat. Introduce a newCloudFetchQueue.Note that this is not yet feature complete, with some known failures in the following domains (for the SEA backend only):
These will be resolved in future PRs.
How is this tested?
examples/experimental/sea_connector_test.py.The coverage of the key classes by the unit tests are as below:
backend/sea/queue.py(SeaCloudFetchQueue and SeaResultSetQueueFactory classes)backend/sea/result_set.py(SeaResultSet class)Related Tickets & Documents
Design Doc
PECOBLR-553