def _execute(self, code: str) -> Output:
if self.session_id is None:
raise ValueError("session not yet started")
statement = self.client.create_statement(self.session_id, code)
intervals = polling_intervals([0.1, 0.2, 0.3, 0.5], 1.0)
def waiting_for_output(statement):
not_finished = statement.state in {
StatementState.WAITING,
StatementState.RUNNING,
}
available = statement.state == StatementState.AVAILABLE
return not_finished or (available and statement.output is None)
while waiting_for_output(statement):
time.sleep(next(intervals))
statement = self.client.get_statement(
statement.session_id, statement.statement_id
)
if statement.output is None:
raise RuntimeError("statement had no output")
return statement.output
def wait(self) -> SessionState:
"""Wait for the batch session to finish."""
intervals = polling_intervals([0.1, 0.5, 1.0, 3.0], 5.0)
while True:
state = self.state
if state in SESSION_STATE_FINISHED:
break
time.sleep(next(intervals))
return state
def start(self) -> None:
"""Create the remote Spark session and wait for it to be ready."""
session = self.client.create_session(
self.kind, self.proxy_user, self.jars, self.py_files, self.files,
self.driver_memory, self.driver_cores, self.executor_memory,
self.executor_cores, self.num_executors, self.archives, self.queue,
self.name, self.spark_conf, self.logger)
self.session_id = session.session_id
intervals = polling_intervals([0.1, 0.2, 0.3, 0.5], 1.0)
while self.state in SESSION_STATE_NOT_READY:
time.sleep(next(intervals))
def wait(self) -> None:
"""Wait for the session to be ready."""
intervals = polling_intervals([0.1, 0.2, 0.3, 0.5], 1.0)
while self.state in SESSION_STATE_NOT_READY:
time.sleep(next(intervals))