mirror of https://github.com/apache/cassandra
Improve performance of cqlsh COPY FROM
Patch by Tyler Hobbs; reviewed by Aleksey Yeschenko for CASSANDRA-8225
This commit is contained in:
parent
5cf61e718f
commit
7110904e40
|
|
@ -1,4 +1,5 @@
|
|||
2.1.4
|
||||
* (cqlsh) Greatly improve performance of COPY FROM (CASSANDRA-8225)
|
||||
* IndexSummary effectiveIndexInterval is now a guideline, not a rule (CASSANDRA-8993)
|
||||
* Use correct bounds for page cache eviction of compressed files (CASSANDRA-8746)
|
||||
* SSTableScanner enforces its bounds (CASSANDRA-8946)
|
||||
|
|
|
|||
290
bin/cqlsh
290
bin/cqlsh
|
|
@ -51,6 +51,8 @@ import platform
|
|||
import warnings
|
||||
import csv
|
||||
import getpass
|
||||
from functools import partial
|
||||
import traceback
|
||||
|
||||
|
||||
readline = None
|
||||
|
|
@ -108,6 +110,8 @@ except ImportError, e:
|
|||
from cassandra.cluster import Cluster, PagedResult
|
||||
from cassandra.query import SimpleStatement, ordered_dict_factory
|
||||
from cassandra.policies import WhiteListRoundRobinPolicy
|
||||
from cassandra.protocol import QueryMessage, ResultMessage
|
||||
from cassandra.marshal import int16_pack, int32_pack, uint64_pack
|
||||
from cassandra.metadata import protect_name, protect_names, protect_value
|
||||
from cassandra.auth import PlainTextAuthProvider
|
||||
|
||||
|
|
@ -117,7 +121,7 @@ cqlshlibdir = os.path.join(CASSANDRA_PATH, 'pylib')
|
|||
if os.path.isdir(cqlshlibdir):
|
||||
sys.path.insert(0, cqlshlibdir)
|
||||
|
||||
from cqlshlib import cqlhandling, cql3handling, pylexotron, sslhandling, async_insert, meter
|
||||
from cqlshlib import cqlhandling, cql3handling, pylexotron, sslhandling
|
||||
from cqlshlib.displaying import (RED, BLUE, CYAN, ANSI_RESET, COLUMN_NAME_COLORS,
|
||||
FormattedValue, colorme)
|
||||
from cqlshlib.formatting import format_by_type, formatter_for, format_value_utype
|
||||
|
|
@ -550,6 +554,7 @@ class Shell(cmd.Cmd):
|
|||
self.auth_provider = PlainTextAuthProvider(username=username, password=password)
|
||||
self.username = username
|
||||
self.keyspace = keyspace
|
||||
self.ssl = ssl
|
||||
self.tracing_enabled = tracing_enabled
|
||||
self.expand_enabled = expand_enabled
|
||||
if use_conn:
|
||||
|
|
@ -913,7 +918,6 @@ class Shell(cmd.Cmd):
|
|||
self.handle_statement(st, statementtext)
|
||||
except Exception, e:
|
||||
if self.debug:
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
else:
|
||||
self.printerr(e)
|
||||
|
|
@ -1435,73 +1439,251 @@ class Shell(cmd.Cmd):
|
|||
except IOError, e:
|
||||
self.printerr("Can't open %r for reading: %s" % (fname, e))
|
||||
return 0
|
||||
|
||||
current_record = None
|
||||
try:
|
||||
if header:
|
||||
linesource.next()
|
||||
table_meta = self.get_table_meta(ks, cf)
|
||||
reader = csv.reader(linesource, **dialect_options)
|
||||
from functools import partial
|
||||
rownum, error = \
|
||||
async_insert.insert_concurrent(self.session, enumerate(reader, start=1),
|
||||
partial(
|
||||
self.create_insert_statement,
|
||||
columns, nullval,
|
||||
table_meta))
|
||||
if error:
|
||||
self.printerr(str(error[0]))
|
||||
self.printerr("Aborting import at record #%d. "
|
||||
"Previously-inserted values still present."
|
||||
% error[1])
|
||||
|
||||
from multiprocessing import Process, Pipe, cpu_count
|
||||
|
||||
# Pick a resonable number of child processes. We need to leave at
|
||||
# least one core for the parent process. This doesn't necessarily
|
||||
# need to be capped at 4, but it's currently enough to keep
|
||||
# a single local Cassandra node busy, and I see lower throughput
|
||||
# with more processes.
|
||||
try:
|
||||
num_processes = max(1, min(4, cpu_count() - 1))
|
||||
except NotImplementedError:
|
||||
num_processes = 1
|
||||
|
||||
processes, pipes = [], [],
|
||||
for i in range(num_processes):
|
||||
parent_conn, child_conn = Pipe()
|
||||
pipes.append(parent_conn)
|
||||
processes.append(Process(target=self.multiproc_import, args=(child_conn, ks, cf, columns, nullval)))
|
||||
|
||||
for process in processes:
|
||||
process.start()
|
||||
|
||||
last_checkpoint_time = time.time()
|
||||
current_rate = 0.0
|
||||
for current_record, row in enumerate(reader, start=1):
|
||||
# write to the child process
|
||||
pipes[current_record % num_processes].send((current_record, row))
|
||||
|
||||
# update the progress and current rate periodically
|
||||
if (current_record % 10000) == 0:
|
||||
new_checkpoint_time = time.time()
|
||||
new_rate = 10000.0 / (new_checkpoint_time - last_checkpoint_time)
|
||||
last_checkpoint_time = new_checkpoint_time
|
||||
|
||||
# smooth the rate a bit
|
||||
if current_rate == 0.0:
|
||||
current_rate = new_rate
|
||||
else:
|
||||
current_rate = (current_rate + new_rate) / 2.0
|
||||
|
||||
output = 'Processed %s rows; Write: %.2f rows/s\r' % \
|
||||
(current_record, current_rate)
|
||||
sys.stdout.write(output)
|
||||
sys.stdout.flush()
|
||||
|
||||
# check for any errors reported by the children
|
||||
if (current_record % 100) == 0:
|
||||
if self._check_child_pipes(current_record, pipes):
|
||||
# no errors seen, continue with outer loop
|
||||
continue
|
||||
else:
|
||||
# errors seen, break out of outer loop
|
||||
break
|
||||
except Exception, exc:
|
||||
if current_record is None:
|
||||
# we failed before we started
|
||||
self.printerr("\nError starting import process:\n")
|
||||
self.printerr(str(exc))
|
||||
if self.debug:
|
||||
traceback.print_exc()
|
||||
else:
|
||||
self.printerr("\n" + str(exc))
|
||||
self.printerr("\nAborting import at record #%d. "
|
||||
"Previously inserted records and some records after "
|
||||
"this number may be present."
|
||||
% (current_record,))
|
||||
if self.debug:
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
# send a message that indicates we're done
|
||||
for pipe in pipes:
|
||||
pipe.send((None, None))
|
||||
|
||||
for process in processes:
|
||||
process.join()
|
||||
|
||||
self._check_child_pipes(current_record, pipes)
|
||||
|
||||
for pipe in pipes:
|
||||
pipe.close()
|
||||
|
||||
if do_close:
|
||||
linesource.close()
|
||||
elif self.tty:
|
||||
print
|
||||
return rownum
|
||||
|
||||
def create_insert_statement(self, columns, nullval, table_meta, row):
|
||||
return current_record
|
||||
|
||||
if len(row) != len(columns):
|
||||
raise ValueError(
|
||||
"Record has the wrong number of fields (%d instead of %d)."
|
||||
% (len(row), len(columns)))
|
||||
def _check_child_pipes(self, current_record, pipes):
|
||||
# check the pipes for errors from child processes
|
||||
for pipe in pipes:
|
||||
if pipe.poll():
|
||||
try:
|
||||
(record_num, error) = pipe.recv()
|
||||
self.printerr("\n" + str(error))
|
||||
self.printerr(
|
||||
"Aborting import at record #%d. "
|
||||
"Previously inserted records are still present, "
|
||||
"and some records after that may be present as well."
|
||||
% (record_num,))
|
||||
return False
|
||||
except EOFError:
|
||||
# pipe is closed, nothing to read
|
||||
self.printerr("\nChild process died without notification, "
|
||||
"aborting import at record #%d. Previously "
|
||||
"inserted records are probably still present, "
|
||||
"and some records after that may be present "
|
||||
"as well." % (current_record,))
|
||||
return False
|
||||
return True
|
||||
|
||||
rowmap = {}
|
||||
primary_key_columns = [col.name for col in table_meta.primary_key]
|
||||
for name, value in zip(columns, row):
|
||||
type = table_meta.columns[name].data_type
|
||||
cqltype = table_meta.columns[name].typestring
|
||||
def multiproc_import(self, pipe, ks, cf, columns, nullval):
|
||||
"""
|
||||
This method is where child processes start when doing a COPY FROM
|
||||
operation. The child process will open one connection to the node and
|
||||
interact directly with the connection, bypassing most of the driver
|
||||
code. Because we don't need retries, connection pooling, thread safety,
|
||||
and other fancy features, this is okay.
|
||||
"""
|
||||
|
||||
if value != nullval:
|
||||
if cqltype in ('ascii', 'text', 'timestamp', 'date', 'time', 'inet'):
|
||||
rowmap[name] = protect_value(value)
|
||||
else:
|
||||
rowmap[name] = value
|
||||
elif name in primary_key_columns:
|
||||
# By default, nullval is an empty string. See CASSANDRA-7792 for details.
|
||||
message = "Cannot insert null value for primary key column '%s'." % (name,)
|
||||
if nullval == '':
|
||||
message += " If you want to insert empty strings, consider using " \
|
||||
"the WITH NULL=<marker> option for COPY."
|
||||
self.printerr(message)
|
||||
return False
|
||||
else:
|
||||
rowmap[name] = 'null'
|
||||
# would be nice to be able to use a prepared query here, but in order
|
||||
# to use that interface, we'd need to have all the input as native
|
||||
# values already, reading them from text just like the various
|
||||
# Cassandra cql types do. Better just to submit them all as intact
|
||||
# CQL string literals and let Cassandra do its thing.
|
||||
query = 'INSERT INTO %s.%s (%s) VALUES (%s)' % (
|
||||
protect_name(table_meta.keyspace.name),
|
||||
protect_name(table_meta.name),
|
||||
', '.join(protect_names(rowmap.keys())),
|
||||
', '.join(rowmap.values())
|
||||
)
|
||||
if self.debug:
|
||||
print 'Import using CQL: %s' % query
|
||||
return SimpleStatement(query)
|
||||
# open a new connection for this subprocess
|
||||
new_cluster = Cluster(
|
||||
contact_points=(self.hostname,),
|
||||
port=self.port,
|
||||
cql_version=self.conn.cql_version,
|
||||
protocol_version=DEFAULT_PROTOCOL_VERSION,
|
||||
auth_provider=self.auth_provider,
|
||||
ssl_options=sslhandling.ssl_settings(self.hostname, CONFIG_FILE) if self.ssl else None,
|
||||
load_balancing_policy=WhiteListRoundRobinPolicy([self.hostname]),
|
||||
compression=None)
|
||||
session = new_cluster.connect(self.keyspace)
|
||||
conn = session._pools.values()[0]._connection
|
||||
|
||||
# pre-build as much of the query as we can
|
||||
table_meta = self.get_table_meta(ks, cf)
|
||||
pk_cols = [col.name for col in table_meta.primary_key]
|
||||
cqltypes = [table_meta.columns[name].typestring for name in columns]
|
||||
pk_indexes = [columns.index(col.name) for col in table_meta.primary_key]
|
||||
query = 'INSERT INTO %s.%s (%s) VALUES (%%s)' % (
|
||||
protect_name(ks),
|
||||
protect_name(cf),
|
||||
', '.join(columns))
|
||||
|
||||
# we need to handle some types specially
|
||||
should_escape = [t in ('ascii', 'text', 'timestamp', 'date', 'time', 'inet') for t in cqltypes]
|
||||
|
||||
insert_timestamp = int(time.time() * 1e6)
|
||||
|
||||
def callback(record_num, response):
|
||||
# This is the callback we register for all inserts. Because this
|
||||
# is run on the event-loop thread, we need to hold a lock when
|
||||
# adjusting in_flight.
|
||||
with conn.lock:
|
||||
conn.in_flight -= 1
|
||||
|
||||
if not isinstance(response, ResultMessage):
|
||||
# It's an error. Notify the parent process and let it send
|
||||
# a stop signal to all child processes (including this one).
|
||||
pipe.send((record_num, str(response)))
|
||||
if isinstance(response, Exception) and self.debug:
|
||||
traceback.print_exc(response)
|
||||
|
||||
current_record = 0
|
||||
insert_num = 0
|
||||
try:
|
||||
while True:
|
||||
# To avoid totally maxing out the connection,
|
||||
# defer to the reactor thread when we're close
|
||||
# to capacity
|
||||
if conn.in_flight > (conn.max_request_id * 0.9):
|
||||
conn._readable = True
|
||||
time.sleep(0.05)
|
||||
continue
|
||||
|
||||
try:
|
||||
(current_record, row) = pipe.recv()
|
||||
except EOFError:
|
||||
# the pipe was closed and there's nothing to receive
|
||||
sys.stdout.write('Failed to read from pipe:\n\n')
|
||||
sys.stdout.flush()
|
||||
conn._writable = True
|
||||
conn._readable = True
|
||||
break
|
||||
|
||||
# see if the parent process has signaled that we are done
|
||||
if (current_record, row) == (None, None):
|
||||
conn._writable = True
|
||||
conn._readable = True
|
||||
pipe.close()
|
||||
break
|
||||
|
||||
# format the values in the row
|
||||
for i, value in enumerate(row):
|
||||
if value != nullval:
|
||||
if should_escape[i]:
|
||||
row[i] = protect_value(value)
|
||||
elif i in pk_indexes:
|
||||
# By default, nullval is an empty string. See CASSANDRA-7792 for details.
|
||||
message = "Cannot insert null value for primary key column '%s'." % (pk_cols[i],)
|
||||
if nullval == '':
|
||||
message += " If you want to insert empty strings, consider using " \
|
||||
"the WITH NULL=<marker> option for COPY."
|
||||
pipe.send((current_record, message))
|
||||
return
|
||||
else:
|
||||
row[i] = 'null'
|
||||
|
||||
full_query = query % (','.join(row),)
|
||||
query_message = QueryMessage(
|
||||
full_query, self.consistency_level, serial_consistency_level=None,
|
||||
fetch_size=None, paging_state=None, timestamp=insert_timestamp)
|
||||
|
||||
request_id = conn.get_request_id()
|
||||
binary_message = query_message.to_binary(
|
||||
stream_id=request_id, protocol_version=DEFAULT_PROTOCOL_VERSION, compression=None)
|
||||
|
||||
# add the message directly to the connection's queue
|
||||
with conn.lock:
|
||||
conn.in_flight += 1
|
||||
conn._callbacks[request_id] = partial(callback, current_record)
|
||||
conn.deque.append(binary_message)
|
||||
|
||||
# every 50 records, clear the pending writes queue and read
|
||||
# any responses we have
|
||||
if insert_num % 50 == 0:
|
||||
conn._writable = True
|
||||
conn._readable = True
|
||||
|
||||
insert_num += 1
|
||||
except Exception, exc:
|
||||
pipe.send((current_record, exc))
|
||||
finally:
|
||||
# wait for any pending requests to finish
|
||||
while conn.in_flight > 0:
|
||||
conn._readable = True
|
||||
time.sleep(0.01)
|
||||
|
||||
new_cluster.shutdown()
|
||||
|
||||
def perform_csv_export(self, ks, cf, columns, fname, opts):
|
||||
dialect_options = self.csv_dialect_defaults.copy()
|
||||
|
|
|
|||
|
|
@ -1,115 +0,0 @@
|
|||
# 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 threading import Event, Condition
|
||||
from . import meter
|
||||
import sys
|
||||
|
||||
class _CountDownLatch(object):
|
||||
def __init__(self, counter=1):
|
||||
self._count = counter
|
||||
self._lock = Condition()
|
||||
|
||||
def count_down(self):
|
||||
with self._lock:
|
||||
self._count -= 1
|
||||
if self._count <= 0:
|
||||
self._lock.notifyAll()
|
||||
|
||||
def await(self):
|
||||
with self._lock:
|
||||
while self._count > 0:
|
||||
# use a timeout so that the main thread wakes up occasionally
|
||||
# so it can see keyboard interrupts (CASSANDRA-7815)
|
||||
self._lock.wait(0.5)
|
||||
|
||||
|
||||
class _ChainedWriter(object):
|
||||
|
||||
CONCURRENCY = 100
|
||||
|
||||
def __init__(self, session, enumerated_reader, statement_func):
|
||||
self._sentinel = object()
|
||||
self._session = session
|
||||
self._cancellation_event = Event()
|
||||
self._first_error = None
|
||||
self._task_counter = _CountDownLatch(self.CONCURRENCY)
|
||||
self._enumerated_reader = enumerated_reader
|
||||
self._statement_func = statement_func
|
||||
self._meter = meter.Meter()
|
||||
|
||||
def insert(self):
|
||||
if not self._enumerated_reader:
|
||||
return 0, None
|
||||
|
||||
for i in xrange(self.CONCURRENCY):
|
||||
self._execute_next(self._sentinel, 0)
|
||||
|
||||
try:
|
||||
self._task_counter.await()
|
||||
except KeyboardInterrupt:
|
||||
self._cancellation_event.set()
|
||||
sys.stdout.write('Aborting due to keyboard interrupt\n')
|
||||
self._task_counter.await()
|
||||
self._meter.done()
|
||||
return self._meter.num_finished(), self._first_error
|
||||
|
||||
|
||||
def _abort(self, error, failed_record):
|
||||
if not self._first_error:
|
||||
self._first_error = error, failed_record
|
||||
self._task_counter.count_down()
|
||||
self._cancellation_event.set()
|
||||
|
||||
def _handle_error(self, error, failed_record):
|
||||
self._abort(error, failed_record)
|
||||
|
||||
def _execute_next(self, result, last_completed_record):
|
||||
if self._cancellation_event.is_set():
|
||||
self._task_counter.count_down()
|
||||
return
|
||||
|
||||
if result is not self._sentinel:
|
||||
self._meter.mark_written()
|
||||
|
||||
try:
|
||||
(current_record, row) = next(self._enumerated_reader)
|
||||
except StopIteration:
|
||||
self._task_counter.count_down()
|
||||
return
|
||||
except Exception as exc:
|
||||
self._abort(exc, last_completed_record)
|
||||
return
|
||||
|
||||
if self._cancellation_event.is_set():
|
||||
self._task_counter.count_down()
|
||||
return
|
||||
|
||||
try:
|
||||
statement = self._statement_func(row)
|
||||
future = self._session.execute_async(statement)
|
||||
future.add_callbacks(callback=self._execute_next,
|
||||
callback_args=(current_record,),
|
||||
errback=self._handle_error,
|
||||
errback_args=(current_record,))
|
||||
except Exception as exc:
|
||||
self._abort(exc, current_record)
|
||||
return
|
||||
|
||||
|
||||
def insert_concurrent(session, enumerated_reader, statement_func):
|
||||
return _ChainedWriter(session, enumerated_reader, statement_func).insert()
|
||||
|
||||
|
|
@ -1,59 +0,0 @@
|
|||
# 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 time import time
|
||||
import sys
|
||||
from threading import RLock
|
||||
|
||||
|
||||
class Meter(object):
|
||||
|
||||
def __init__(self):
|
||||
self._num_finished = 0
|
||||
self._last_checkpoint_time = None
|
||||
self._current_rate = 0.0
|
||||
self._lock = RLock()
|
||||
|
||||
def mark_written(self):
|
||||
with self._lock:
|
||||
if not self._last_checkpoint_time:
|
||||
self._last_checkpoint_time = time()
|
||||
self._num_finished += 1
|
||||
|
||||
if self._num_finished % 10000 == 0:
|
||||
previous_checkpoint_time = self._last_checkpoint_time
|
||||
self._last_checkpoint_time = time()
|
||||
new_rate = 10000.0 / (self._last_checkpoint_time - previous_checkpoint_time)
|
||||
if self._current_rate == 0.0:
|
||||
self._current_rate = new_rate
|
||||
else:
|
||||
self._current_rate = (self._current_rate + new_rate) / 2.0
|
||||
|
||||
if self._num_finished % 1000 != 0:
|
||||
return
|
||||
output = 'Processed %s rows; Write: %.2f rows/s\r' % \
|
||||
(self._num_finished, self._current_rate)
|
||||
sys.stdout.write(output)
|
||||
sys.stdout.flush()
|
||||
|
||||
def num_finished(self):
|
||||
with self._lock:
|
||||
return self._num_finished
|
||||
|
||||
def done(self):
|
||||
print ""
|
||||
|
||||
|
||||
Loading…
Reference in New Issue