Files
impala/tests/shell/test_shell_client.py
Sahil Takiar 3088ca8580 IMPALA-9818: Add fetch size as option to impala shell
Adds the option --fetch_size to the Impala shell. This new option allows
users to specify the fetch size used when issuing fetch RPCs to the
Impala Coordinator (e.g. TFetchResultsReq and BeeswaxService.fetch).
This parameter applies for all client protocols: beeswax, hs2, hs2-http.
The default --fetch_size is set to 10240 (10x the default batch size).

The new --fetch_size parameter is most effective when result spooling is
enabled. When result spooling is disabled, Impala can only return a
single row batch per fetch RPC (so 1024 rows by default). When result
spooling is enabled, Impala can return up to 100 row batches per fetch
request.

Removes some logic in the the impala_client.py file that attempts to
simulate a fetch_size. The code would issue multiple fetch requests to
fullfill the given fetch_size. This logic is no longer needed now that
result spooling is available.

Testing:
* Ran core tests
* Added new tests in test_shell_client.py and test_shell_commandline.py

Change-Id: I8dc7962aada6b38795241d067a99bd94fabca57b
Reviewed-on: http://gerrit.cloudera.org:8080/16041
Reviewed-by: Impala Public Jenkins <impala-public-jenkins@cloudera.com>
Tested-by: Sahil Takiar <stakiar@cloudera.com>
2020-06-10 17:46:21 +00:00

97 lines
3.8 KiB
Python

#!/usr/bin/env impala-python
# -*- coding: utf-8 -*-
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from shell.impala_client import ImpalaBeeswaxClient, ImpalaHS2Client
from tests.common.impala_test_suite import ImpalaTestSuite
from tests.common.test_dimensions import create_client_protocol_dimension
from util import get_impalad_host_port
class TestShellClient(ImpalaTestSuite):
"""Tests for the Impala Shell clients: ImpalaBeeswaxClient and ImpalaHS2Client."""
@classmethod
def get_workload(self):
return 'functional-query'
@classmethod
def add_test_dimensions(cls):
cls.ImpalaTestMatrix.add_dimension(create_client_protocol_dimension())
def test_fetch_size(self, vector):
"""Tests that when result spooling is disabled, setting a small batch_size causes
the shell to fetch a single batch at a time, even when the configured fetch size is
larger than the batch_size."""
handle = None
num_rows = 100
batch_size = 10
query_options = {'batch_size': str(batch_size)}
client = self.__get_shell_client(vector)
try:
client.connect()
handle = client.execute_query(
"select * from functional.alltypes limit {0}".format(num_rows), query_options)
self.__fetch_rows(client.fetch(handle), batch_size, num_rows)
finally:
if handle is not None: client.close_query(handle)
client.close_connection()
def test_fetch_size_result_spooling(self, vector):
"""Tests that when result spooling is enabled, that the exact fetch_size is honored
even if a small batch_size is configured."""
handle = None
fetch_size = 20
num_rows = 100
query_options = {'batch_size': '10', 'spool_query_results': 'true'}
client = self.__get_shell_client(vector, fetch_size)
try:
client.connect()
handle = client.execute_query(
"select * from functional.alltypes limit {0}".format(num_rows), query_options)
self.__fetch_rows(client.fetch(handle), num_rows / fetch_size, num_rows)
finally:
if handle is not None: client.close_query(handle)
client.close_connection()
def __fetch_rows(self, fetch_batches, num_batches, num_rows):
"""Fetches all rows using the given fetch_batches generator. Asserts that num_batches
batches are produced by the generator and that num_rows are returned."""
num_batches_count = 0
rows_per_batch = num_rows / num_batches
for fetch_batch in fetch_batches:
assert len(fetch_batch) == rows_per_batch
num_batches_count += 1
if num_batches_count == num_batches: break
assert num_batches_count == num_batches
def __get_shell_client(self, vector, fetch_size=1024):
"""Returns the client specified by the protocol in the given vector."""
impalad = get_impalad_host_port(vector).split(":")
protocol = vector.get_value("protocol")
if protocol == 'hs2':
return ImpalaHS2Client(impalad, fetch_size, None)
elif protocol == 'hs2-http':
return ImpalaHS2Client(impalad, fetch_size, None,
use_http_base_transport=True, http_path='cliservice')
elif protocol == 'beeswax':
return ImpalaBeeswaxClient(impalad, fetch_size, None)