You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@cassandra.apache.org by dj...@apache.org on 2023/05/19 22:35:42 UTC
[cassandra-analytics] branch trunk created (now 1633cd9)
This is an automated email from the ASF dual-hosted git repository.
djoshi pushed a change to branch trunk
in repository https://gitbox.apache.org/repos/asf/cassandra-analytics.git
at 1633cd9 CEP-28: Apache Cassandra Analytics
This branch includes the following new commits:
new 1633cd9 CEP-28: Apache Cassandra Analytics
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@cassandra.apache.org
For additional commands, e-mail: commits-help@cassandra.apache.org
[cassandra-analytics] 01/01: CEP-28: Apache Cassandra Analytics
Posted by dj...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
djoshi pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/cassandra-analytics.git
commit 1633cd9c6c3d88d5c66825fab76a369266509f7e
Author: Dinesh Joshi <dj...@apache.org>
AuthorDate: Fri May 19 14:57:47 2023 -0700
CEP-28: Apache Cassandra Analytics
This is the initial commit for the Apache Cassandra Analytics project
where we support reading and writing bulk data from Apache Cassandra from
Spark.
Patch by James Berragan, Doug Rohrer; Reviewed by Dinesh Joshi, Yifan Cai for CASSANDRA-16222
Co-authored-by: James Berragan <jb...@apple.com>
Co-authored-by: Doug Rohrer <dr...@apple.com>
Co-authored-by: Saranya Krishnakumar <sa...@apple.com>
Co-authored-by: Francisco Guerrero <fr...@apple.com>
Co-authored-by: Yifan Cai <yc...@apache.org>
Co-authored-by: Jyothsna Konisa <jk...@apple.com>
Co-authored-by: Yuriy Semchyshyn <ys...@apple.com>
Co-authored-by: Dinesh Joshi <dj...@apache.org>
---
.gitignore | 101 +
DEV-README.md | 71 +
LICENSE.txt | 202 ++
README.md | 37 +
build.gradle | 144 ++
cassandra-analytics-core-example/README.md | 231 ++
cassandra-analytics-core-example/build.gradle | 56 +
.../spark/example/SampleCassandraJob.java | 221 ++
cassandra-analytics-core/build.gradle | 179 ++
.../cassandra/bridge/CassandraBridgeFactory.java | 135 ++
.../cassandra/bridge/CassandraVersionFeatures.java | 198 ++
.../bridge/PostDelegationClassLoader.java | 71 +
.../apache/cassandra/clients/ExecutorHolder.java | 42 +
.../java/org/apache/cassandra/clients/Sidecar.java | 405 ++++
.../cassandra/clients/SidecarInstanceImpl.java | 122 ++
.../clients/SidecarStreamConsumerAdapter.java | 107 +
.../org/apache/cassandra/clients/SslConfig.java | 486 ++++
.../org/apache/cassandra/spark/KryoRegister.java | 108 +
.../cassandra/spark/bulkwriter/BulkSparkConf.java | 464 ++++
.../spark/bulkwriter/BulkWriteValidator.java | 169 ++
.../spark/bulkwriter/BulkWriterContext.java | 37 +
.../cassandra/spark/bulkwriter/CancelJobEvent.java | 38 +
.../bulkwriter/CassandraBulkSourceRelation.java | 185 ++
.../bulkwriter/CassandraBulkWriterContext.java | 194 ++
.../spark/bulkwriter/CassandraClusterInfo.java | 499 +++++
.../spark/bulkwriter/CassandraContext.java | 99 +
.../spark/bulkwriter/CassandraJobInfo.java | 119 +
.../spark/bulkwriter/CassandraRingMonitor.java | 135 ++
.../spark/bulkwriter/CassandraSchemaInfo.java | 37 +
.../cassandra/spark/bulkwriter/ClusterInfo.java | 52 +
.../spark/bulkwriter/CommitCoordinator.java | 223 ++
.../cassandra/spark/bulkwriter/CommitError.java | 48 +
.../cassandra/spark/bulkwriter/CommitResult.java | 66 +
.../spark/bulkwriter/CqlTableInfoProvider.java | 175 ++
.../spark/bulkwriter/DataTransferApi.java | 63 +
.../cassandra/spark/bulkwriter/DecoratedKey.java | 92 +
.../spark/bulkwriter/InstanceAvailability.java | 53 +
.../apache/cassandra/spark/bulkwriter/JobInfo.java | 54 +
.../bulkwriter/LoadNotSupportedException.java | 28 +
.../cassandra/spark/bulkwriter/RecordWriter.java | 205 ++
.../cassandra/spark/bulkwriter/RingInstance.java | 144 ++
.../cassandra/spark/bulkwriter/RowBufferMode.java | 26 +
.../cassandra/spark/bulkwriter/SSTableWriter.java | 188 ++
.../spark/bulkwriter/SSTableWriterFactory.java | 50 +
.../cassandra/spark/bulkwriter/SchemaInfo.java | 27 +
.../spark/bulkwriter/SidecarDataTransferApi.java | 153 ++
.../spark/bulkwriter/SqlToCqlTypeConverter.java | 726 ++++++
.../cassandra/spark/bulkwriter/StreamError.java | 34 +
.../cassandra/spark/bulkwriter/StreamResult.java | 59 +
.../cassandra/spark/bulkwriter/StreamSession.java | 300 +++
.../spark/bulkwriter/TableInfoProvider.java | 48 +
.../cassandra/spark/bulkwriter/TableSchema.java | 234 ++
.../spark/bulkwriter/TokenPartitioner.java | 243 ++
.../cassandra/spark/bulkwriter/Tokenizer.java | 47 +
.../cassandra/spark/bulkwriter/WriteMode.java | 26 +
.../cassandra/spark/bulkwriter/WriterOption.java | 25 +
.../cassandra/spark/bulkwriter/WriterOptions.java | 47 +
.../spark/bulkwriter/token/CassandraRing.java | 291 +++
.../spark/bulkwriter/token/ConsistencyLevel.java | 175 ++
.../spark/bulkwriter/token/RangeUtils.java | 148 ++
.../token/ReplicaAwareFailureHandler.java | 134 ++
.../spark/bulkwriter/token/TokenUtils.java | 164 ++
.../spark/bulkwriter/util/FastByteOperations.java | 267 +++
.../spark/bulkwriter/util/SbwJavaSerializer.java | 111 +
.../spark/bulkwriter/util/SbwKryoRegistrator.java | 76 +
.../spark/bulkwriter/util/ThreadUtil.java | 52 +
.../cassandra/spark/cdc/CdcCellIterator.java | 51 +
.../apache/cassandra/spark/cdc/CdcRowIterator.java | 53 +
.../spark/cdc/watermarker/InMemoryWatermarker.java | 301 +++
.../org/apache/cassandra/spark/common/MD5Hash.java | 43 +
.../apache/cassandra/spark/common/SSTables.java | 36 +
.../spark/common/client/ClientException.java | 37 +
.../spark/common/client/InstanceState.java | 33 +
.../spark/common/client/InstanceStatus.java | 30 +
.../cassandra/spark/common/model/BulkFeatures.java | 32 +
.../spark/common/model/CassandraInstance.java | 29 +
.../cassandra/spark/common/schema/BooleanType.java | 51 +
.../cassandra/spark/common/schema/BytesType.java | 55 +
.../spark/common/schema/CollectionType.java | 49 +
.../cassandra/spark/common/schema/ColumnType.java | 51 +
.../cassandra/spark/common/schema/ColumnTypes.java | 108 +
.../cassandra/spark/common/schema/ColumnUtil.java | 54 +
.../cassandra/spark/common/schema/DoubleType.java | 50 +
.../cassandra/spark/common/schema/IntegerType.java | 44 +
.../cassandra/spark/common/schema/ListType.java | 70 +
.../cassandra/spark/common/schema/LongType.java | 44 +
.../cassandra/spark/common/schema/MapType.java | 70 +
.../cassandra/spark/common/schema/SetType.java | 45 +
.../cassandra/spark/common/schema/StringType.java | 48 +
.../spark/common/schema/StringUuidType.java | 45 +
.../cassandra/spark/common/schema/UuidType.java | 56 +
.../apache/cassandra/spark/data/BasicSupplier.java | 46 +
.../cassandra/spark/data/CassandraDataLayer.java | 1169 ++++++++++
.../spark/data/CassandraDataSourceHelper.java | 131 ++
.../apache/cassandra/spark/data/DefaultSizing.java | 47 +
.../cassandra/spark/data/FileSystemSSTable.java | 102 +
.../cassandra/spark/data/FileSystemSource.java | 157 ++
.../cassandra/spark/data/LocalCommitLog.java | 109 +
.../cassandra/spark/data/LocalDataLayer.java | 479 ++++
.../cassandra/spark/data/PartitionedDataLayer.java | 580 +++++
.../apache/cassandra/spark/data/ReplicaSet.java | 92 +
.../spark/data/SidecarProvisionedSSTable.java | 262 +++
.../org/apache/cassandra/spark/data/Sizing.java | 34 +
.../spark/data/partitioner/CassandraRing.java | 339 +++
.../spark/data/partitioner/ConsistencyLevel.java | 124 ++
.../spark/data/partitioner/MultipleReplicas.java | 153 ++
.../partitioner/NotEnoughReplicasException.java | 59 +
.../spark/data/partitioner/SingleReplica.java | 186 ++
.../spark/data/partitioner/TokenPartitioner.java | 289 +++
.../spark/sparksql/AbstractSparkRowIterator.java | 160 ++
.../spark/sparksql/CassandraDataSink.java | 105 +
.../cassandra/spark/sparksql/FullRowBuilder.java | 148 ++
.../spark/sparksql/SparkCellIterator.java | 412 ++++
.../apache/cassandra/spark/utils/BuildInfo.java | 121 +
.../org/apache/cassandra/spark/utils/CqlUtils.java | 291 +++
.../apache/cassandra/spark/utils/FilterUtils.java | 114 +
.../apache/cassandra/spark/utils/Properties.java | 66 +
.../apache/cassandra/spark/utils/RangeUtils.java | 168 ++
.../cassandra/spark/utils/ScalaFunctions.java | 53 +
.../cassandra/spark/utils/TableIdentifier.java | 81 +
.../org/apache/cassandra/spark/utils/UUIDs.java | 90 +
...org.apache.spark.sql.sources.DataSourceRegister | 2 +
.../spark/sparksql/CassandraDataSource.java | 220 ++
.../cassandra/spark/sparksql/LocalDataSource.java | 44 +
.../cassandra/spark/sparksql/SparkRowIterator.java | 77 +
.../spark/sparksql/CassandraDataSource.java | 47 +
.../spark/sparksql/CassandraInputPartition.java | 37 +
.../spark/sparksql/CassandraMicroBatchStream.java | 147 ++
.../sparksql/CassandraPartitionReaderFactory.java | 69 +
.../spark/sparksql/CassandraPartitioning.java | 60 +
.../spark/sparksql/CassandraScanBuilder.java | 156 ++
.../cassandra/spark/sparksql/CassandraTable.java | 69 +
.../spark/sparksql/CassandraTableProvider.java | 67 +
.../cassandra/spark/sparksql/LocalDataSource.java | 43 +
.../cassandra/spark/sparksql/SparkRowIterator.java | 198 ++
.../bridge/CassandraVersionFeaturesTest.java | 59 +
.../apache/cassandra/bridge/CompressionTests.java | 65 +
.../org/apache/cassandra/bridge/VersionRunner.java | 51 +
.../cassandra/clients/SidecarClientConfigTest.java | 101 +
.../cassandra/clients/SidecarInstanceImplTest.java | 74 +
.../cassandra/clients/SidecarInstanceTest.java | 85 +
.../apache/cassandra/clients/SslConfigTest.java | 216 ++
.../org/apache/cassandra/spark/EndToEndTests.java | 2312 ++++++++++++++++++++
.../cassandra/spark/KryoSerializationTests.java | 359 +++
.../org/apache/cassandra/spark/TestDataLayer.java | 159 ++
.../org/apache/cassandra/spark/TestRunnable.java | 31 +
.../java/org/apache/cassandra/spark/TestUtils.java | 355 +++
.../java/org/apache/cassandra/spark/Tester.java | 547 +++++
.../spark/bulkwriter/BulkSparkConfTest.java | 218 ++
.../spark/bulkwriter/CassandraRingMonitorTest.java | 137 ++
.../spark/bulkwriter/CqlTableInfoProviderTest.java | 127 ++
.../spark/bulkwriter/DecoratedKeyTest.java | 63 +
.../spark/bulkwriter/MockBulkWriterContext.java | 405 ++++
.../bulkwriter/MockScheduledExecutorService.java | 182 ++
.../spark/bulkwriter/MockTableWriter.java | 102 +
.../bulkwriter/NonValidatingTestSSTableWriter.java | 36 +
.../spark/bulkwriter/RecordWriterTest.java | 159 ++
.../bulkwriter/RingInstanceSerializationTest.java | 57 +
.../spark/bulkwriter/RingInstanceTest.java | 229 ++
.../cassandra/spark/bulkwriter/RingUtils.java | 112 +
.../spark/bulkwriter/SSTableWriterTest.java | 103 +
.../cassandra/spark/bulkwriter/SSTables.java | 134 ++
.../bulkwriter/SqlToCqlTypeConverterTest.java | 148 ++
.../bulkwriter/StreamSessionConsistencyTest.java | 217 ++
.../spark/bulkwriter/StreamSessionTest.java | 304 +++
.../spark/bulkwriter/TableSchemaNormalizeTest.java | 314 +++
.../spark/bulkwriter/TableSchemaTest.java | 199 ++
.../spark/bulkwriter/TableSchemaTestCommon.java | 348 +++
.../spark/bulkwriter/TokenPartitionerTest.java | 155 ++
.../cassandra/spark/bulkwriter/UploadRequest.java | 50 +
.../benchmarks/ByteOperationsBenchmark.java | 78 +
.../spark/cdc/watermarker/WatermarkerTests.java | 205 ++
.../data/CassandraDataSourceHelperCacheTest.java | 183 ++
.../cassandra/spark/data/CqlFieldComparators.java | 260 +++
.../apache/cassandra/spark/data/CqlFieldTests.java | 256 +++
.../DataLayerUnsupportedPushDownFiltersTest.java | 250 +++
.../cassandra/spark/data/DefaultSizingTest.java | 40 +
.../cassandra/spark/data/LocalDataLayerTests.java | 94 +
.../spark/data/PartitionedDataLayerTests.java | 389 ++++
.../spark/data/ReplicationFactorTests.java | 113 +
.../apache/cassandra/spark/data/VersionRunner.java | 54 +
.../spark/data/partitioner/CassandraRingTests.java | 446 ++++
.../data/partitioner/ConsistencyLevelTests.java | 112 +
.../data/partitioner/JDKSerializationTests.java | 312 +++
.../data/partitioner/MultipleReplicasTests.java | 208 ++
.../spark/data/partitioner/SingleReplicaTests.java | 238 ++
.../data/partitioner/TokenPartitionerTests.java | 83 +
.../spark/reader/CassandraBridgeTests.java | 57 +
.../spark/reader/DataTypeSerializationTests.java | 569 +++++
.../cassandra/spark/reader/OverlapTests.java | 196 ++
.../cassandra/spark/reader/PartitionKeyTests.java | 77 +
.../apache/cassandra/spark/reader/RidTests.java | 136 ++
.../apache/cassandra/spark/reader/SchemaTests.java | 584 +++++
.../spark/reader/TombstoneWriterTests.java | 204 ++
.../reader/common/CompressionInputStreamTests.java | 125 ++
.../cassandra/spark/utils/BuildInfoTest.java | 78 +
.../spark/utils/ByteBufferUtilsTests.java | 100 +
.../apache/cassandra/spark/utils/CqlUtilsTest.java | 678 ++++++
.../cassandra/spark/utils/FilterUtilsTests.java | 80 +
.../cassandra/spark/utils/RangeUtilsTest.java | 267 +++
.../spark/utils/SSTableInputStreamHttpTests.java | 312 +++
.../spark/utils/SSTableInputStreamTests.java | 396 ++++
.../cassandra/spark/utils/SerializationUtils.java | 97 +
.../src/test/resources/cql/fullSchema.cql | 105 +
.../spark/bulkwriter/TestTaskContext.java | 169 ++
.../spark/sparksql/PartitionKeyFilterTests.java | 116 +
.../spark/sparksql/SparkRangeFilterTests.java | 53 +
.../spark/sparksql/SparkRowIteratorTests.java | 274 +++
.../spark/bulkwriter/TestTaskContext.java | 180 ++
.../org/apache/cassandra/spark/cdc/CdcTester.java | 402 ++++
.../org/apache/cassandra/spark/cdc/CdcTests.java | 547 +++++
.../cassandra/spark/cdc/CdcTombstoneTests.java | 710 ++++++
.../cdc/RequireTwoReplicasLocalDataSource.java | 45 +
.../cassandra/spark/sparksql/CdcOffsetTests.java | 42 +
cassandra-bridge/build.gradle | 32 +
.../apache/cassandra/bridge/BigNumberConfig.java | 54 +
.../cassandra/bridge/BigNumberConfigImpl.java | 187 ++
.../apache/cassandra/bridge/CassandraBridge.java | 561 +++++
.../apache/cassandra/bridge/CassandraVersion.java | 69 +
.../apache/cassandra/bridge/RangeTombstone.java | 44 +
.../org/apache/cassandra/bridge/SSTableWriter.java | 28 +
.../org/apache/cassandra/bridge/TokenRange.java | 154 ++
.../org/apache/cassandra/spark/cdc/CommitLog.java | 176 ++
.../cassandra/spark/cdc/CommitLogProvider.java | 27 +
.../spark/cdc/IPartitionUpdateWrapper.java | 27 +
.../apache/cassandra/spark/cdc/TableIdLookup.java | 42 +
.../cdc/watermarker/DoNothingWatermarker.java | 83 +
.../spark/cdc/watermarker/Watermarker.java | 89 +
.../cassandra/spark/config/SchemaFeature.java | 114 +
.../cassandra/spark/config/SchemaFeatureSet.java | 182 ++
.../org/apache/cassandra/spark/data/CqlField.java | 465 ++++
.../org/apache/cassandra/spark/data/CqlTable.java | 329 +++
.../org/apache/cassandra/spark/data/DataLayer.java | 356 +++
.../org/apache/cassandra/spark/data/FileType.java | 83 +
.../spark/data/IncompleteSSTableException.java | 43 +
.../cassandra/spark/data/ReplicationFactor.java | 224 ++
.../org/apache/cassandra/spark/data/SSTable.java | 102 +
.../cassandra/spark/data/SSTablesSupplier.java | 46 +
.../spark/data/partitioner/CassandraInstance.java | 121 +
.../spark/data/partitioner/MurmurHash.java | 159 ++
.../spark/data/partitioner/Partitioner.java | 114 +
.../exceptions/TransportFailureException.java | 97 +
.../cassandra/spark/reader/EmptyStreamScanner.java | 47 +
.../org/apache/cassandra/spark/reader/Rid.java | 238 ++
.../cassandra/spark/reader/SparkSSTableReader.java | 54 +
.../cassandra/spark/reader/StreamScanner.java | 79 +
.../reader/common/AbstractCompressionMetadata.java | 101 +
.../spark/reader/common/BigLongArray.java | 84 +
.../spark/reader/common/ChunkCorruptException.java | 30 +
.../spark/reader/common/RawInputStream.java | 251 +++
.../reader/common/SSTableStreamException.java | 60 +
.../org/apache/cassandra/spark/sparksql/Cell.java | 43 +
.../sparksql/CellTombstonesInComplexDecorator.java | 96 +
.../sparksql/LastModifiedTimestampDecorator.java | 70 +
.../spark/sparksql/NoMatchFoundException.java | 32 +
.../cassandra/spark/sparksql/RangeTombstone.java | 33 +
.../spark/sparksql/RangeTombstoneDecorator.java | 94 +
.../spark/sparksql/RangeTombstoneMarker.java | 48 +
.../cassandra/spark/sparksql/RowBuilder.java | 58 +
.../spark/sparksql/RowBuilderDecorator.java | 123 ++
.../apache/cassandra/spark/sparksql/Tombstone.java | 35 +
.../spark/sparksql/TombstonesInComplex.java | 41 +
.../spark/sparksql/UpdateFlagDecorator.java | 64 +
.../sparksql/UpdatedFieldsIndicatorDecorator.java | 75 +
.../spark/sparksql/filters/CdcOffset.java | 116 +
.../spark/sparksql/filters/CdcOffsetFilter.java | 91 +
.../spark/sparksql/filters/PartitionKeyFilter.java | 84 +
.../spark/sparksql/filters/PruneColumnFilter.java | 52 +
.../spark/sparksql/filters/SparkRangeFilter.java | 64 +
.../org/apache/cassandra/spark/stats/Stats.java | 667 ++++++
.../apache/cassandra/spark/utils/ArrayUtils.java | 42 +
.../cassandra/spark/utils/ByteBufferUtils.java | 225 ++
.../apache/cassandra/spark/utils/ColumnTypes.java | 130 ++
.../cassandra/spark/utils/ComparisonUtils.java | 203 ++
.../apache/cassandra/spark/utils/FutureUtils.java | 165 ++
.../org/apache/cassandra/spark/utils/IOUtils.java | 105 +
.../apache/cassandra/spark/utils/LoggerHelper.java | 152 ++
.../org/apache/cassandra/spark/utils/MapUtils.java | 213 ++
.../apache/cassandra/spark/utils/RandomUtils.java | 107 +
.../spark/utils/SparkClassLoaderOverride.java | 85 +
.../cassandra/spark/utils/TemporaryDirectory.java | 48 +
.../cassandra/spark/utils/ThrowableUtils.java | 70 +
.../org/apache/cassandra/spark/utils/Throwing.java | 68 +
.../apache/cassandra/spark/utils/TimeProvider.java | 34 +
.../spark/utils/streaming/SSTableInputStream.java | 607 +++++
.../spark/utils/streaming/SSTableSource.java | 102 +
.../spark/utils/streaming/StreamBuffer.java | 83 +
.../spark/utils/streaming/StreamConsumer.java | 48 +
.../cassandra/spark/utils/test/TestSSTable.java | 133 ++
.../cassandra/spark/utils/test/TestSchema.java | 879 ++++++++
.../spark/utils/ScalaConversionUtils.java | 56 +
.../spark/utils/ScalaConversionUtils.java | 56 +
.../spark/utils/ScalaConversionUtils.java | 56 +
.../spark/utils/ScalaConversionUtils.java | 56 +
.../cassandra/bridge/BigNumberConfigImplTest.java | 48 +
.../cassandra/spark/utils/ArrayUtilsTest.java | 66 +
.../cassandra/spark/utils/LoggerHelperTests.java | 43 +
.../cassandra/spark/utils/ThrowableUtilTests.java | 112 +
cassandra-four-zero/build.gradle | 87 +
.../bridge/CassandraBridgeImplementation.java | 1051 +++++++++
.../apache/cassandra/bridge/CassandraSchema.java | 61 +
.../apache/cassandra/bridge/CollectionElement.java | 46 +
.../bridge/SSTableWriterImplementation.java | 82 +
.../org/apache/cassandra/bridge/TestCommitLog.java | 92 +
.../cql3/functions/types/TupleHelper.java | 58 +
.../cql3/functions/types/UserTypeHelper.java | 59 +
.../db/commitlog/BufferingCommitLogReader.java | 747 +++++++
.../db/commitlog/PartitionUpdateWrapper.java | 219 ++
.../commitlog/SeekableCommitLogSegmentReader.java | 185 ++
.../io/sstable/SSTableTombstoneWriter.java | 512 +++++
.../cassandra/io/util/CdcRandomAccessReader.java | 258 +++
.../org/apache/cassandra/spark/data/CqlType.java | 197 ++
.../apache/cassandra/spark/data/NativeType.java | 180 ++
.../spark/data/complex/CqlCollection.java | 248 +++
.../cassandra/spark/data/complex/CqlFrozen.java | 241 ++
.../cassandra/spark/data/complex/CqlList.java | 180 ++
.../cassandra/spark/data/complex/CqlMap.java | 215 ++
.../cassandra/spark/data/complex/CqlSet.java | 133 ++
.../cassandra/spark/data/complex/CqlTuple.java | 244 +++
.../cassandra/spark/data/complex/CqlUdt.java | 535 +++++
.../apache/cassandra/spark/data/types/Ascii.java | 47 +
.../apache/cassandra/spark/data/types/BigInt.java | 54 +
.../cassandra/spark/data/types/BinaryBased.java | 60 +
.../apache/cassandra/spark/data/types/Blob.java | 76 +
.../apache/cassandra/spark/data/types/Boolean.java | 99 +
.../apache/cassandra/spark/data/types/Counter.java | 39 +
.../apache/cassandra/spark/data/types/Date.java | 129 ++
.../apache/cassandra/spark/data/types/Decimal.java | 113 +
.../apache/cassandra/spark/data/types/Double.java | 93 +
.../cassandra/spark/data/types/Duration.java | 39 +
.../apache/cassandra/spark/data/types/Empty.java | 98 +
.../apache/cassandra/spark/data/types/Float.java | 93 +
.../apache/cassandra/spark/data/types/Inet.java | 86 +
.../org/apache/cassandra/spark/data/types/Int.java | 93 +
.../cassandra/spark/data/types/LongBased.java | 65 +
.../cassandra/spark/data/types/SmallInt.java | 93 +
.../cassandra/spark/data/types/StringBased.java | 102 +
.../apache/cassandra/spark/data/types/Text.java | 47 +
.../apache/cassandra/spark/data/types/Time.java | 54 +
.../cassandra/spark/data/types/TimeUUID.java | 54 +
.../cassandra/spark/data/types/Timestamp.java | 105 +
.../apache/cassandra/spark/data/types/TinyInt.java | 98 +
.../apache/cassandra/spark/data/types/UUID.java | 75 +
.../apache/cassandra/spark/data/types/VarChar.java | 47 +
.../apache/cassandra/spark/data/types/VarInt.java | 106 +
.../spark/reader/AbstractStreamScanner.java | 606 +++++
.../cassandra/spark/reader/CdcScannerBuilder.java | 536 +++++
.../spark/reader/CompactionStreamScanner.java | 152 ++
.../spark/reader/CompressedRawInputStream.java | 286 +++
.../spark/reader/CompressionMetadata.java | 110 +
.../cassandra/spark/reader/CompressionUtils.java | 84 +
.../cassandra/spark/reader/IndexDbUtils.java | 173 ++
.../apache/cassandra/spark/reader/ReaderUtils.java | 547 +++++
.../cassandra/spark/reader/SSTableCache.java | 155 ++
.../cassandra/spark/reader/SSTableReader.java | 875 ++++++++
.../apache/cassandra/spark/reader/Scannable.java | 27 +
.../cassandra/spark/reader/SchemaBuilder.java | 547 +++++
.../cassandra/spark/reader/SummaryDbUtils.java | 201 ++
.../RangeTombstoneMarkerImplementation.java | 133 ++
.../cassandra/spark/data/types/DateTypeTests.java | 57 +
.../cassandra/spark/reader/CdcScannerTests.java | 152 ++
.../cassandra/spark/reader/IndexDbTests.java | 281 +++
.../cassandra/spark/reader/IndexOffsetTests.java | 190 ++
.../spark/reader/KryoSerializationTests.java | 102 +
.../cassandra/spark/reader/PartitionKeyTests.java | 94 +
.../cassandra/spark/reader/ReaderUtilsTests.java | 367 ++++
.../cassandra/spark/reader/SSTableCacheTests.java | 148 ++
.../cassandra/spark/reader/SSTableReaderTests.java | 1162 ++++++++++
.../cassandra/spark/reader/SchemaBuilderTests.java | 117 +
.../cassandra/spark/reader/SummaryDbTests.java | 176 ++
cassandra-three-zero/TODO.md | 7 +
code_version.sh | 43 +
config/checkstyle/checkstyle.xml | 224 ++
config/checkstyle/suppressions.xml | 40 +
githooks/pre-push | 28 +
gradle.properties | 32 +
gradle/wrapper/gradle-wrapper.jar | Bin 0 -> 58702 bytes
gradle/wrapper/gradle-wrapper.properties | 22 +
gradlew | 167 ++
ide/idea/Project_Default.xml | 151 ++
ide/idea/codeStyleSettings.xml | 349 +++
profiles/scala-2.11-spark-2-jdk-1.8.gradle | 28 +
profiles/scala-2.12-spark-2-jdk-1.8.gradle | 28 +
profiles/scala-2.12-spark-3-jdk-11.gradle | 28 +
profiles/scala-2.13-spark-3-jdk-11.gradle | 28 +
settings.gradle | 25 +
386 files changed, 63670 insertions(+)
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..c8e73d0
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,101 @@
+# 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.
+
+# Compiled class files
+*.class
+
+# Log files
+*.log
+
+# BlueJ files
+*.ctxt
+
+# Mobile Tools for Java (J2ME)
+.mtj.tmp/
+
+# Package files
+*.war
+*.nar
+*.ear
+*.zip
+*.tar.gz
+*.rar
+
+# Virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml
+hs_err_pid*
+
+# Build output directories
+.dist
+/target
+*/target
+/build
+*/build
+
+# IntelliJ files and directories
+out
+.idea
+*.ipr
+*.iws
+*.iml
+atlassian-ide-plugin.xml
+
+# Gradle
+.gradle/
+.out/
+
+# Docker tests
+nohup.out
+dockertests/libs
+test_results/
+
+# generted certs
+sbr_keystore
+sbr_keystore_pass
+sbr_truststore
+sbr_truststore_pass
+truststore.jks
+
+*.hprof
+
+__pycache__/
+
+.vscode/
+
+cdc/bin
+
+# Miscellaneous
+/env/**
+build.log
+.gradle
+.packager
+.shelf
+logs
+generated
+**/build/**
+metastore_db/
+*.pyc
+.cache/
+ttr.log
+buildcache/
+.testdata*/
+dev-lo.properties
+pytests/.env/**/*
+pytests/.pytest_cache
+cassandra-analytics-build.properties
+
+# MacOS
+.DS_Store
diff --git a/DEV-README.md b/DEV-README.md
new file mode 100644
index 0000000..d85755c
--- /dev/null
+++ b/DEV-README.md
@@ -0,0 +1,71 @@
+# Cassandra Analytics
+
+Cassandra Analytics supports Spark 2 (Scala 2.11 and 2.12) and Spark 3 (Scala 2.12).
+
+This project uses Gradle as the dependency management and build framework.
+
+## Building
+
+Cassandra Analytics will build for Spark 2 and Scala 2.11 by default.
+
+Navigate to the top-level directory for this project:
+
+```shell
+./gradlew clean package
+```
+
+### Spark 2 and Scala 2.12
+
+To build for Scala 2.12, set the profile by exporting `SCALA_VERSION=2.12`:
+
+```shell
+export SCALA_VERSION=2.12
+./gradlew clean package
+```
+
+### Spark 3 and Scala 2.12
+
+To build for Spark 3 and Scala 2.12, export both `SCALA_VERSION=2.12` and `SPARK_VERSION=3`:
+
+```shell
+export SCALA_VERSION=2.12
+export SPARK_VERSION=3
+./gradlew clean package
+```
+
+### Git hooks (optional)
+
+To enable git hooks, run the following command at project root.
+
+```shell
+git config core.hooksPath githooks
+```
+
+## IntelliJ
+
+The project is well-supported in IntelliJ.
+
+Run the following profile to copy code style used for this project:
+
+```shell
+./gradlew copyCodeStyle
+```
+
+The project has different sources for Spark 2 and Spark 3.
+
+Spark 2 uses the `org.apache.spark.sql.sources.v2` APIs that have been deprecated in Spark 3.
+
+Spark 3 uses new APIs that live in the `org.apache.spark.sql.connector.read` namespace.
+
+By default, the project will load Spark 2 sources, but you can switch between sources by modifying the `gradle.properties` file.
+
+For Spark 3, use the following in `gradle.properties`:
+
+```properties
+scala=2.12
+spark=3
+```
+
+And then load Gradle changes (on Mac, the shortcut to load Gradle changes is <kbd>Command</kbd> + <kbd>Shift</kbd> + <kbd>I</kbd>).
+
+This will make the IDE pick up the Spark 3 sources, and you should now be able to develop against Spark 3 as well.
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000..d645695
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,202 @@
+
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
+
+ APPENDIX: How to apply the Apache License to your work.
+
+ To apply the Apache License to your work, attach the following
+ boilerplate notice, with the fields enclosed by brackets "[]"
+ replaced with your own identifying information. (Don't include
+ the brackets!) The text should be enclosed in the appropriate
+ comment syntax for the file format. We also recommend that a
+ file or class name and description of purpose be included on the
+ same "printed page" as the copyright notice for easier
+ identification within third-party archives.
+
+ Copyright [yyyy] [name of copyright owner]
+
+ Licensed 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.
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..3820361
--- /dev/null
+++ b/README.md
@@ -0,0 +1,37 @@
+# Cassandra Analytics
+
+## Cassandra Spark Bulk Reader
+
+The open-source repository for the Cassandra Spark Bulk Reader. This library allows integration between Cassandra and Spark job, allowing users to run arbitrary Spark jobs against a Cassandra cluster securely and consistently.
+
+This project contains the necessary open-source implementations to connect to a Cassandra cluster and read the data into Spark.
+
+For example usage, see the example repository; sample steps:
+
+```scala
+import org.apache.cassandra.spark.sparksql.CassandraDataSource
+import org.apache.spark.sql.SparkSession
+
+val sparkSession = SparkSession.builder.getOrCreate()
+val df = sparkSession.read.format("org.apache.cassandra.spark.sparksql.CassandraDataSource")
+ .option("sidecar_instances", "localhost,localhost2,localhost3")
+ .option("keyspace", "sbr_tests")
+ .option("table", "basic_test")
+ .option("DC", "datacenter1")
+ .option("createSnapshot", true)
+ .option("numCores", 4)
+ .load()
+```
+
+## Cassandra Spark Bulk Writer
+
+The Cassandra Spark Bulk Writer allows for high-speed data ingest to Cassandra clusters running Cassandra 3.0 and 4.0.
+
+If you are a consumer of the Cassandra Spark Bulk Writer, please see our end-user documentation: usage instructions, FAQs, troubleshooting guides, and release notes.
+
+Developers interested in contributing to the SBW, please see the [DEV-README](DEV-README.md).
+
+## Getting Started
+
+For example usage, see the [example repository](cassandra-analytics-core-example/README.md). This example covers both
+setting up Cassandra 4.0, Apache Sidecar, and running a Spark Bulk Reader and Spark Bulk Writer job.
diff --git a/build.gradle b/build.gradle
new file mode 100644
index 0000000..ff72ca1
--- /dev/null
+++ b/build.gradle
@@ -0,0 +1,144 @@
+/*
+ * 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.
+ */
+
+buildscript {
+ dependencies {
+ classpath(group: 'com.github.jengelman.gradle.plugins', name: 'shadow', version: '6.1.0')
+ }
+}
+
+plugins {
+ id 'java'
+ id 'java-library'
+ id 'com.github.johnrengelman.shadow' version '5.1.0'
+}
+
+ext.scalaLabel = System.getenv("SCALA_VERSION") ?: "${scala}"
+println("Scala version: ${ext.scalaLabel}")
+
+ext.sparkLabel = (System.getenv("SPARK_VERSION") ?: "${spark}").split(/\./)[0]
+println("Spark version: ${ext.sparkLabel}")
+
+ext.jdkLabel = System.getenv("JDK_VERSION") ?: "${analyticsJDKLevel}"
+println("Java source/target compatibility level: ${ext.jdkLabel}")
+
+def profile = "profiles/scala-${ext.scalaLabel}-spark-${ext.sparkLabel}-jdk-${ext.jdkLabel}.gradle"
+if (!file(profile).exists()) {
+ throw new InvalidUserDataException("Profile ${profile} does not exist, which indicates this combination of Scala, Spark, and Java is unsupported.\n" +
+ "Please either add an appropriate profile if this combination should work, or use a different one.")
+}
+
+println("Using profile ${profile}")
+apply(from: profile)
+
+apply(plugin: 'idea')
+
+task copyCodeStyle(type: Copy) {
+ from('ide/idea/codeStyleSettings.xml')
+ into('.idea')
+}
+
+tasks.idea.dependsOn(tasks.copyCodeStyle)
+
+task copyInspections(type: Copy) {
+ from('ide/idea/Project_Default.xml')
+ into('.idea/inspectionProfiles')
+}
+tasks.idea.dependsOn(tasks.copyInspections)
+
+subprojects {
+ apply(plugin: 'java-library')
+ apply(plugin: 'checkstyle')
+
+ sourceCompatibility = "${project.rootProject.ext.jdkLabel}"
+ targetCompatibility = "${project.rootProject.ext.jdkLabel}"
+
+ tasks.withType(JavaCompile) {
+ if ("${project.rootProject.ext.jdkLabel}" == '1.8') {
+ // Unfortunately, we can't use release here. We can only use public APIs when using release.
+ // org.apache.cassandra.spark.bulkwriter.util.FastByteOperations uses sun.misc.Unsafe which is causing the build
+ // to fail.
+ options.compilerArgs += ['-source', '8', '-target', '8']
+ } else {
+ options.release = 11
+ }
+ }
+
+ archivesBaseName = "${project.name}"
+ if ("${project.rootProject.ext.sparkLabel}" == '3') {
+ archivesBaseName = "${archivesBaseName}_spark3"
+ }
+ archivesBaseName = "${archivesBaseName}_${scalaMajorVersion}"
+
+ if ("${project.rootProject.ext.jdkLabel}" == '1.8') {
+ if ("${version}".contains('-SNAPSHOT')) {
+ version = "${version}".replace('-SNAPSHOT', '-jdk8-SNAPSHOT')
+ } else {
+ version = "${version}-jdk8"
+ }
+ }
+
+ repositories {
+ // uncomment the line below for local development
+ mavenLocal()
+ }
+
+ dependencies {
+ compileOnly(group: 'com.intellij', name: 'annotations', version: "${project.rootProject.intellijVersion}")
+ testCompileOnly(group: 'com.intellij', name: 'annotations', version: "${project.rootProject.intellijVersion}")
+ }
+
+ sourceSets {
+ main {
+ java {
+ srcDirs += ["src/main/spark${project.rootProject.ext.sparkLabel}",
+ "src/main/scala-${scalaMajorVersion}-spark-${sparkMajorVersion}"]
+ }
+ }
+ test {
+ java {
+ srcDirs += ["src/test/spark${project.rootProject.ext.sparkLabel}",
+ "src/test/scala-${scalaMajorVersion}-spark-${sparkMajorVersion}"]
+ }
+ }
+ }
+
+ test {
+ if (JavaVersion.current().isJava11Compatible()) {
+ def JDK11_OPTIONS = ['-Djdk.attach.allowAttachSelf=true',
+ '--add-exports', 'java.base/jdk.internal.misc=ALL-UNNAMED',
+ '--add-exports', 'java.base/jdk.internal.ref=ALL-UNNAMED',
+ '--add-exports', 'java.base/sun.nio.ch=ALL-UNNAMED',
+ '--add-exports', 'java.management.rmi/com.sun.jmx.remote.internal.rmi=ALL-UNNAMED',
+ '--add-exports', 'java.rmi/sun.rmi.registry=ALL-UNNAMED',
+ '--add-exports', 'java.rmi/sun.rmi.server=ALL-UNNAMED',
+ '--add-exports', 'java.sql/java.sql=ALL-UNNAMED',
+ '--add-opens', 'java.base/java.lang.module=ALL-UNNAMED',
+ '--add-opens', 'java.base/jdk.internal.loader=ALL-UNNAMED',
+ '--add-opens', 'java.base/jdk.internal.ref=ALL-UNNAMED',
+ '--add-opens', 'java.base/jdk.internal.reflect=ALL-UNNAMED',
+ '--add-opens', 'java.base/jdk.internal.math=ALL-UNNAMED',
+ '--add-opens', 'java.base/jdk.internal.module=ALL-UNNAMED',
+ '--add-opens', 'java.base/jdk.internal.util.jar=ALL-UNNAMED',
+ '--add-opens', 'jdk.management/com.sun.management.internal=ALL-UNNAMED']
+ jvmArgs(JDK11_OPTIONS)
+ println("JVM arguments for $project.name are $allJvmArgs")
+ }
+ }
+}
diff --git a/cassandra-analytics-core-example/README.md b/cassandra-analytics-core-example/README.md
new file mode 100644
index 0000000..4e219a8
--- /dev/null
+++ b/cassandra-analytics-core-example/README.md
@@ -0,0 +1,231 @@
+# Sample Cassandra Data with Spark Bulk Analytics Job
+
+This sub-project showcases the Cassandra Spark Bulk Analytics read and write functionality. The job writes data to a
+Cassandra table using the bulk writer functionality, and then it reads the data using the bulk reader functionality.
+
+## Requirements
+
+- Java 11
+- A running Cassandra 4.0 cluster
+- A running Cassandra Sidecar service
+
+## Setup
+
+This example requires a Cassandra 4.0 cluster running, we recommend a 3-node cluster. Additionally, it requires the
+latest version of Cassandra Sidecar that we will run from source.
+
+### Step 1: Cassandra using CCM
+
+The easiest way to provision a Cassandra cluster is by using Cassandra Cluster Manager or
+[CCM](https://github.com/riptano/ccm). Follow the configuration instructions for CCM and provision a new test cluster.
+
+For example,
+
+```shell
+ccm create test --version=4.1.1 --nodes=3
+```
+
+> **Note**
+> If you are using macOS, explicit configurations of interface aliases are needed as follows:
+
+```shell
+sudo ifconfig lo0 alias 127.0.0.2
+sudo ifconfig lo0 alias 127.0.0.3
+sudo bash -c 'echo "127.0.0.2 localhost2" >> /etc/hosts'
+sudo bash -c 'echo "127.0.0.3 localhost3" >> /etc/hosts'
+```
+
+After configuring the network interfaces run:
+
+```shell
+ccm start
+```
+
+Verify that your Cassandra cluster is configured and running successfully.
+
+### Step 2: Configure and Run Cassandra Sidecar
+
+In this step, we will clone and configure the Cassandra Sidecar project. Finally, we will run Sidecar which will be
+connecting to our local Cassandra 3-node cluster.
+
+```shell
+git clone https://github.com/frankgh/cassandra-sidecar/tree/CEP-28-bulk-apis
+cd cassandra-sidecar
+```
+
+Configure the `main/dist/sidecar.yaml` file for your local environment. You will most likely only need to configure
+the `cassandra_instances` section in your file pointing to your local Cassandra data directories. Here is what my
+configuration looks like for this tutorial:
+
+```yaml
+cassandra_instances:
+ - id: 1
+ host: localhost
+ port: 9042
+ data_dirs: <your_ccm_parent_path>/.ccm/test/node1/data0
+ uploads_staging_dir: <your_ccm_parent_path>/.ccm/test/node1/sstable-staging
+ jmx_host: 127.0.0.1
+ jmx_port: 7100
+ jmx_ssl_enabled: false
+ - id: 2
+ host: localhost2
+ port: 9042
+ data_dirs: <your_ccm_parent_path>/.ccm/test/node2/data0
+ uploads_staging_dir: <your_ccm_parent_path>/.ccm/test/node2/sstable-staging
+ jmx_host: 127.0.0.1
+ jmx_port: 7200
+ jmx_ssl_enabled: false
+ - id: 3
+ host: localhost3
+ port: 9042
+ data_dirs: <your_ccm_parent_path>/.ccm/test/node3/data0
+ uploads_staging_dir: <your_ccm_parent_path>/.ccm/test/node3/sstable-staging
+ jmx_host: 127.0.0.1
+ jmx_port: 7300
+ jmx_ssl_enabled: false
+```
+
+I have a 3 node setup, so I configure Sidecar for those 3 nodes. CCM creates the Cassandra cluster under
+`${HOME}/.ccm/test`, so I update my `data_dirs` and `uploads_staging_dir` configuration to use my local path.
+
+Next, create the `uploads_staging_dir` where Sidecar will stage SSTables coming from Cassandra Spark bulk writer.
+In my case, I have decided to keep the `sstable-staging` directory inside each of the node's directories.
+
+```shell
+mkdir -p ${HOME}/.ccm/test/node1/sstable-staging
+mkdir -p ${HOME}/.ccm/test/node2/sstable-staging
+mkdir -p ${HOME}/.ccm/test/node3/sstable-staging
+```
+
+Finally, run Cassandra Sidecar, we skip running integration tests because we need docker for integration tests. You
+can opt to run integration tests if you have docker running in your local environment.
+
+```shell
+user:~$ ./gradlew run -x integrationTest
+> Task :common:compileJava UP-TO-DATE
+> Task :cassandra40:compileJava UP-TO-DATE
+> Task :compileJava UP-TO-DATE
+> Task :processResources UP-TO-DATE
+> Task :classes UP-TO-DATE
+> Task :jar UP-TO-DATE
+> Task :startScripts UP-TO-DATE
+> Task :cassandra40:processResources NO-SOURCE
+> Task :cassandra40:classes UP-TO-DATE
+> Task :cassandra40:jar UP-TO-DATE
+> Task :common:processResources NO-SOURCE
+> Task :common:classes UP-TO-DATE
+> Task :common:jar UP-TO-DATE
+> Task :distTar
+> Task :distZip
+> Task :assemble
+> Task :common:compileTestFixturesJava UP-TO-DATE
+> Task :compileTestFixturesJava UP-TO-DATE
+> Task :compileTestJava UP-TO-DATE
+> Task :processTestResources UP-TO-DATE
+> Task :testClasses UP-TO-DATE
+> Task :compileIntegrationTestJava UP-TO-DATE
+> Task :processIntegrationTestResources NO-SOURCE
+> Task :integrationTestClasses UP-TO-DATE
+> Task :checkstyleIntegrationTest UP-TO-DATE
+> Task :checkstyleMain UP-TO-DATE
+> Task :checkstyleTest UP-TO-DATE
+> Task :processTestFixturesResources NO-SOURCE
+> Task :testFixturesClasses UP-TO-DATE
+> Task :checkstyleTestFixtures UP-TO-DATE
+> Task :testFixturesJar UP-TO-DATE
+> Task :common:processTestFixturesResources NO-SOURCE
+> Task :common:testFixturesClasses UP-TO-DATE
+> Task :common:testFixturesJar UP-TO-DATE
+> Task :test UP-TO-DATE
+> Task :jacocoTestReport UP-TO-DATE
+> Task :spotbugsIntegrationTest UP-TO-DATE
+> Task :spotbugsMain UP-TO-DATE
+> Task :spotbugsTest UP-TO-DATE
+> Task :spotbugsTestFixtures UP-TO-DATE
+> Task :check
+> Task :copyJolokia UP-TO-DATE
+> Task :installDist
+> Task :copyDist
+> Task :docs:asciidoctor UP-TO-DATE
+> Task :copyDocs UP-TO-DATE
+> Task :generateReDoc NO-SOURCE
+> Task :generateSwaggerUI NO-SOURCE
+> Task :build
+
+> Task :run
+Could not start Jolokia agent: java.net.BindException: Address already in use
+WARNING: An illegal reflective access operation has occurred
+WARNING: Please consider reporting this to the maintainers of com.google.inject.internal.cglib.core.$ReflectUtils$1
+WARNING: Use --illegal-access=warn to enable warnings of further illegal reflective access operations
+WARNING: All illegal access operations will be denied in a future release
+INFO [main] 2023-04-13 13:13:35,939 YAMLSidecarConfiguration.java:211 - Reading configuration from file://${HOME}/workspace/cassandra-sidecar-server/conf/sidecar.yaml
+INFO [main] 2023-04-13 13:13:35,997 CQLSessionProvider.java:63 - Connecting to localhost on port 9042
+INFO [main] 2023-04-13 13:13:36,001 CQLSessionProvider.java:63 - Connecting to localhost2 on port 9042
+INFO [main] 2023-04-13 13:13:36,014 CQLSessionProvider.java:63 - Connecting to localhost3 on port 9042
+INFO [main] 2023-04-13 13:13:36,030 CacheFactory.java:83 - Building SSTable Import Cache with expireAfterAccess=PT2H, maxSize=10000
+ _____ _ _____ _ _
+/ __ \ | | / ___(_) | |
+| / \/ __ _ ___ ___ __ _ _ __ __| |_ __ __ _ \ `--. _ __| | ___ ___ __ _ _ __
+| | / _` / __/ __|/ _` | '_ \ / _` | '__/ _` | `--. \ |/ _` |/ _ \/ __/ _` | '__|
+| \__/\ (_| \__ \__ \ (_| | | | | (_| | | | (_| | /\__/ / | (_| | __/ (_| (_| | |
+ \____/\__,_|___/___/\__,_|_| |_|\__,_|_| \__,_| \____/|_|\__,_|\___|\___\__,_|_|
+
+
+INFO [main] 2023-04-13 13:13:36,229 CassandraSidecarDaemon.java:75 - Starting Cassandra Sidecar on 0.0.0.0:9043
+```
+
+There we have it, Cassandra Sidecar is now running and connected to all 3 Cassandra nodes on my local machine.
+
+### Step 3: Run the Sample Job
+
+To be able to run the [Sample Job](./src/main/java/org/apache/cassandra/spark/example/SampleCassandraJob.java), you
+need to create the keyspace and table used for the test.
+
+Connect to your local Cassandra cluster using CCM:
+
+```shell
+ccm node1 cqlsh
+```
+
+Then run the following command to create the keyspace:
+
+```cassandraql
+CREATE KEYSPACE spark_test WITH replication = {'class': 'NetworkTopologyStrategy', 'datacenter1': '3'}
+ AND durable_writes = true;
+```
+
+Then run the following command to create the table:
+
+```cassandraql
+CREATE TABLE spark_test.test
+(
+ id BIGINT PRIMARY KEY,
+ course BLOB,
+ marks BIGINT
+);
+```
+
+The Cassandra Spark Bulk Analytics project depends on the Sidecar client. Since the Sidecar dependencies for the
+Analytics have not yet been published to maven central, we need to build them locally:
+
+```shell
+cd ${SIDECAR_REPOSITORY_HOME}
+./gradlew -Pversion=1.0.0-local :common:publishToMavenLocal :client:publishToMavenLocal :vertx-client:publishToMavenLocal :vertx-client-shaded:publishToMavenLocal
+```
+
+We use the suffix `-local` to denote that these artifacts have been built locally.
+
+Finally, we are ready to run the example spark job:
+
+```shell
+cd ${ANALYTICS_REPOSITORY_HOME}
+./gradlew :cassandra-analytics-core-example:run
+```
+
+## Tear down
+
+Stop the Cassandra Sidecar project and tear down the ccm cluster
+
+```shell
+ccm remove test
+```
diff --git a/cassandra-analytics-core-example/build.gradle b/cassandra-analytics-core-example/build.gradle
new file mode 100644
index 0000000..93bce75
--- /dev/null
+++ b/cassandra-analytics-core-example/build.gradle
@@ -0,0 +1,56 @@
+/*
+ * 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.
+ */
+
+plugins {
+ id 'java'
+ id 'application'
+}
+
+repositories {
+ mavenCentral()
+}
+
+dependencies {
+ implementation(project(':cassandra-analytics-core'))
+ implementation(group: 'net.java.dev.jna', name: 'jna', version: '5.9.0')
+ implementation(group: 'org.apache.spark', name: 'spark-core_2.12', version: '3.2.1')
+ implementation(group: 'org.apache.spark', name: 'spark-sql_2.12', version: '3.2.1')
+}
+
+application {
+ mainClass = 'org.apache.cassandra.spark.example.SampleCassandraJob'
+ applicationDefaultJvmArgs = ["-Dfile.encoding=UTF-8",
+ "-Djdk.attach.allowAttachSelf=true",
+ "--add-exports", "java.base/jdk.internal.misc=ALL-UNNAMED",
+ "--add-exports", "java.base/jdk.internal.ref=ALL-UNNAMED",
+ "--add-exports", "java.base/sun.nio.ch=ALL-UNNAMED",
+ "--add-exports", "java.management.rmi/com.sun.jmx.remote.internal.rmi=ALL-UNNAMED",
+ "--add-exports", "java.rmi/sun.rmi.registry=ALL-UNNAMED",
+ "--add-exports", "java.rmi/sun.rmi.server=ALL-UNNAMED",
+ "--add-exports", "java.sql/java.sql=ALL-UNNAMED",
+ "--add-opens", "java.base/java.lang.module=ALL-UNNAMED",
+ "--add-opens", "java.base/jdk.internal.loader=ALL-UNNAMED",
+ "--add-opens", "java.base/jdk.internal.ref=ALL-UNNAMED",
+ "--add-opens", "java.base/jdk.internal.reflect=ALL-UNNAMED",
+ "--add-opens", "java.base/jdk.internal.math=ALL-UNNAMED",
+ "--add-opens", "java.base/jdk.internal.module=ALL-UNNAMED",
+ "--add-opens", "java.base/jdk.internal.util.jar=ALL-UNNAMED",
+ "--add-opens", "jdk.management/com.sun.management.internal=ALL-UNNAMED"]
+
+}
diff --git a/cassandra-analytics-core-example/src/main/java/org/apache/cassandra/spark/example/SampleCassandraJob.java b/cassandra-analytics-core-example/src/main/java/org/apache/cassandra/spark/example/SampleCassandraJob.java
new file mode 100644
index 0000000..4e21749
--- /dev/null
+++ b/cassandra-analytics-core-example/src/main/java/org/apache/cassandra/spark/example/SampleCassandraJob.java
@@ -0,0 +1,221 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.example;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.LongStream;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.KryoRegister;
+import org.apache.cassandra.spark.bulkwriter.BulkSparkConf;
+import org.apache.spark.SparkConf;
+import org.apache.spark.SparkContext;
+import org.apache.spark.api.java.JavaRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.api.java.function.Function2;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.RowFactory;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.types.StructType;
+
+import static org.apache.spark.sql.types.DataTypes.BinaryType;
+import static org.apache.spark.sql.types.DataTypes.LongType;
+
+/**
+ * Example showcasing the Cassandra Spark Analytics write and read capabilities
+ * <p>
+ * Prepare your environment by creating the following keyspace and table
+ * <p>
+ * Schema for the {@code keyspace}:
+ * <pre>
+ * CREATE KEYSPACE spark_test WITH replication = {'class': 'NetworkTopologyStrategy', 'datacenter1': '3'}
+ * AND durable_writes = true;
+ * </pre>
+ * <p>
+ * Schema for the {@code table}:
+ * <pre>
+ * CREATE TABLE spark_test.test (
+ * id BIGINT PRIMARY KEY,
+ * course BLOB,
+ * marks BIGINT
+ * );
+ * </pre>
+ */
+public final class SampleCassandraJob
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(SampleCassandraJob.class);
+
+ private SampleCassandraJob()
+ {
+ throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
+ }
+
+ public static void main(String[] args)
+ {
+ LOGGER.info("Starting Spark job with args={}", Arrays.toString(args));
+
+ SparkConf sparkConf = new SparkConf().setAppName("Sample Spark Cassandra Bulk Reader Job")
+ .set("spark.master", "local[8]");
+
+ // Add SBW-specific settings
+ // TODO: Simplify setting up spark conf
+ BulkSparkConf.setupSparkConf(sparkConf, true);
+ KryoRegister.setup(sparkConf);
+
+ SparkSession spark = SparkSession
+ .builder()
+ .config(sparkConf)
+ .getOrCreate();
+ SparkContext sc = spark.sparkContext();
+ SQLContext sql = spark.sqlContext();
+ LOGGER.info("Spark Conf: " + sparkConf.toDebugString());
+ int rowCount = 10_000;
+
+ try
+ {
+ Dataset<Row> written = write(rowCount, sparkConf, sql, sc);
+ Dataset<Row> read = read(rowCount, sparkConf, sql, sc);
+
+ checkSmallDataFrameEquality(written, read);
+ LOGGER.info("Finished Spark job, shutting down...");
+ sc.stop();
+ }
+ catch (Throwable throwable)
+ {
+ LOGGER.error("Unexpected exception executing Spark job", throwable);
+ try
+ {
+ sc.stop();
+ }
+ catch (Throwable ignored)
+ {
+ }
+ }
+ }
+
+ private static Dataset<Row> write(long rowCount, SparkConf sparkConf, SQLContext sql, SparkContext sc)
+ {
+ StructType schema = new StructType()
+ .add("id", LongType, false)
+ .add("course", BinaryType, false)
+ .add("marks", LongType, false);
+
+ JavaSparkContext javaSparkContext = JavaSparkContext.fromSparkContext(sc);
+ int parallelism = sc.defaultParallelism();
+ JavaRDD<Row> rows = genDataset(javaSparkContext, rowCount, parallelism);
+ Dataset<Row> df = sql.createDataFrame(rows, schema);
+
+ df.write()
+ .format("org.apache.cassandra.spark.sparksql.CassandraDataSink")
+ .option("sidecar_instances", "localhost,localhost2,localhost3")
+ .option("keyspace", "spark_test")
+ .option("table", "test")
+ .option("local_dc", "datacenter1")
+ .option("bulk_writer_cl", "LOCAL_QUORUM")
+ .option("number_splits", "-1")
+ .mode("append")
+ .save();
+ return df;
+ }
+
+ private static Dataset<Row> read(int expectedRowCount, SparkConf sparkConf, SQLContext sql, SparkContext sc)
+ {
+ int coresPerExecutor = sparkConf.getInt("spark.executor.cores", 1);
+ int numExecutors = sparkConf.getInt("spark.dynamicAllocation.maxExecutors", sparkConf.getInt("spark.executor.instances", 1));
+ int numCores = coresPerExecutor * numExecutors;
+
+ Dataset<Row> df = sql.read().format("org.apache.cassandra.spark.sparksql.CassandraDataSource")
+ .option("sidecar_instances", "localhost,localhost2,localhost3")
+ .option("keyspace", "spark_test")
+ .option("table", "test")
+ .option("DC", "datacenter1")
+ .option("snapshotName", UUID.randomUUID().toString())
+ .option("createSnapshot", "true")
+ .option("defaultParallelism", sc.defaultParallelism())
+ .option("numCores", numCores)
+ .option("sizing", "default")
+ .load();
+
+ long count = df.count();
+ LOGGER.info("Found {} records", count);
+
+ if (count != expectedRowCount)
+ {
+ LOGGER.error("Expected {} records but found {} records", expectedRowCount, count);
+ return null;
+ }
+ return df;
+ }
+
+ private static void checkSmallDataFrameEquality(Dataset<Row> expected, Dataset<Row> actual)
+ {
+ if (actual == null)
+ {
+ throw new NullPointerException("actual dataframe is null");
+ }
+ if (!actual.exceptAll(expected).isEmpty())
+ {
+ throw new IllegalStateException("The content of the dataframes differs");
+ }
+ }
+
+ private static JavaRDD<Row> genDataset(JavaSparkContext sc, long records, Integer parallelism)
+ {
+ long recordsPerPartition = records / parallelism;
+ long remainder = records - (recordsPerPartition * parallelism);
+ List<Integer> seq = IntStream.range(0, parallelism).boxed().collect(Collectors.toList());
+ JavaRDD<Row> dataset = sc.parallelize(seq, parallelism).mapPartitionsWithIndex(
+ (Function2<Integer, Iterator<Integer>, Iterator<Row>>) (index, integerIterator) -> {
+ long firstRecordNumber = index * recordsPerPartition;
+ long recordsToGenerate = index.equals(parallelism) ? remainder : recordsPerPartition;
+ java.util.Iterator<Row> rows = LongStream.range(0, recordsToGenerate).mapToObj(offset -> {
+ long recordNumber = firstRecordNumber + offset;
+ String courseNameString = String.valueOf(recordNumber);
+ Integer courseNameStringLen = courseNameString.length();
+ Integer courseNameMultiplier = 1000 / courseNameStringLen;
+ byte[] courseName = dupStringAsBytes(courseNameString, courseNameMultiplier);
+ return RowFactory.create(recordNumber, courseName, recordNumber);
+ }).iterator();
+ return rows;
+ }, false);
+ return dataset;
+ }
+
+ private static byte[] dupStringAsBytes(String string, Integer times)
+ {
+ byte[] stringBytes = string.getBytes();
+ ByteBuffer buffer = ByteBuffer.allocate(stringBytes.length * times);
+ for (int time = 0; time < times; time++)
+ {
+ buffer.put(stringBytes);
+ }
+ return buffer.array();
+ }
+}
diff --git a/cassandra-analytics-core/build.gradle b/cassandra-analytics-core/build.gradle
new file mode 100644
index 0000000..77cae04
--- /dev/null
+++ b/cassandra-analytics-core/build.gradle
@@ -0,0 +1,179 @@
+/*
+ * 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.
+ */
+
+project(':cassandra-analytics-core') {
+ apply(plugin: 'java-library')
+ apply(plugin: 'jacoco')
+ apply(plugin: 'maven-publish')
+
+ publishing {
+ publications {
+ maven(MavenPublication) {
+ from components.java
+ groupId project.group
+ artifactId "${archivesBaseName}"
+ version version
+ }
+ }
+ }
+
+ configurations {
+ fourzero {
+ description = 'Bridge to Cassandra 4.0 with all of its dependencies'
+ }
+ all {
+ resolutionStrategy {
+ force("com.fasterxml.jackson.core:jackson-databind:${jacksonVersion}")
+ }
+ }
+ }
+
+ dependencies {
+ compileOnly(project(':cassandra-bridge'))
+
+ fourzero(project(path: ':cassandra-four-zero', configuration: 'shadow'))
+
+ // Including newer JNA which works with "Modern" Linux/GLIBC.
+ // End users can exclude this and include their own if necessary.
+ runtimeOnly(group: 'net.java.dev.jna', name: 'jna', version: "${jnaVersion}")
+ runtimeOnly(group: 'net.java.dev.jna', name: 'jna-platform', version: "${jnaVersion}")
+
+ // This dependency needs to be installed locally by building the Cassandra Sidecar project using the following
+ // commands:
+ // ./gradlew -Pversion=1.0.0-local :common:publishToMavenLocal :client:publishToMavenLocal :vertx-client:publishToMavenLocal :vertx-client-shaded:publishToMavenLocal
+ implementation(group: 'org.apache.cassandra.sidecar', name: 'vertx-client-all', version: "${sidecarVersion}") {
+ exclude(group: 'com.fasterxml.jackson.core', module: 'jackson-annotations') // use provided annotations from spark
+ }
+
+ if ("${scalaMajorVersion}" == "2.11")
+ {
+ implementation(group: 'org.scala-lang.modules', name: "scala-java8-compat_2.11", version: '1.0.1', transitive: false)
+ }
+
+ compileOnly(group: "${sparkGroupId}", name: "spark-core_${scalaMajorVersion}", version: "${project.rootProject.sparkVersion}")
+ compileOnly(group: "${sparkGroupId}", name: "spark-sql_${scalaMajorVersion}", version: "${project.rootProject.sparkVersion}")
+
+ testImplementation(group: 'com.google.guava', name: 'guava', version: '19.0')
+ testImplementation(group: 'junit', name: 'junit', version: "${project.rootProject.junitVersion}")
+ testImplementation(group: 'org.quicktheories', name: 'quicktheories', version: "${project.rootProject.quickTheoriesVersion}")
+ testImplementation(group: 'org.slf4j', name: 'slf4j-simple', version: '1.7.26')
+ testImplementation(group: 'org.mockito', name: 'mockito-core', version: "${project.rootProject.mockitoVersion}")
+ testImplementation(group: 'org.mockito', name: 'mockito-inline', version: "${project.rootProject.mockitoVersion}")
+ testImplementation(group: "${sparkGroupId}", name: "spark-core_${scalaMajorVersion}", version: "${project.rootProject.sparkVersion}")
+ testImplementation(group: "${sparkGroupId}", name: "spark-sql_${scalaMajorVersion}", version: "${project.rootProject.sparkVersion}")
+ testImplementation(group: 'org.hamcrest', name: 'hamcrest-all', version: '1.3')
+ testImplementation(group: 'org.openjdk.jmh', name: 'jmh-core', version: '1.33')
+ testImplementation(group: 'org.apache.httpcomponents', name: 'httpclient', version: '4.5.13')
+ testImplementation(group: 'com.github.luben', name: 'zstd-jni', version: '1.5.0-4')
+ testImplementation(group: 'org.xerial.snappy', name: 'snappy-java', version: '1.1.8.4')
+
+ testImplementation(group: 'com.github.luben', name: 'zstd-jni', version: '1.5.0-4')
+ testCompileOnly(group: "${sparkGroupId}", name: "spark-core_${scalaMajorVersion}", version: "${project.rootProject.sparkVersion}")
+ testCompileOnly(group: "${sparkGroupId}", name: "spark-sql_${scalaMajorVersion}", version: "${project.rootProject.sparkVersion}")
+
+ testRuntimeOnly(group: 'net.java.dev.jna', name: 'jna', version: "${jnaVersion}")
+ testRuntimeOnly(group: 'net.java.dev.jna', name: 'jna-platform', version: "${jnaVersion}")
+
+ testImplementation(project(':cassandra-bridge'))
+ }
+
+ def utils = ['org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext.class',
+ 'org/apache/cassandra/spark/bulkwriter/MockBulkWriterContext$CommitResultSupplier.class',
+ 'org/apache/cassandra/spark/bulkwriter/RingUtils.class',
+ 'org/apache/cassandra/spark/bulkwriter/TableSchemaTestCommon.class',
+ 'org/apache/cassandra/spark/bulkwriter/TableSchemaTestCommon$MockTableInfoProvider.class',
+ 'org/apache/cassandra/spark/bulkwriter/TableSchemaTestCommon$MockTableSchemaBuilder.class']
+
+ jar {
+ dependsOn(configurations.fourzero)
+
+ writeBuildVersion(version, project.projectDir)
+
+ exclude('META-INF/*')
+ from(project(':cassandra-bridge').sourceSets.main.output)
+ from(project(':cassandra-analytics-core').sourceSets.test.output) {
+ include(utils)
+ }
+ from(project(':cassandra-analytics-core').sourceSets.test.output) {
+ exclude(utils)
+ into('tests')
+ }
+ into('bridges') {
+ from(configurations.fourzero.singleFile)
+ }
+ }
+
+ test {
+ minHeapSize = '1024m'
+ maxHeapSize = '3072m'
+ maxParallelForks = Math.max(Runtime.runtime.availableProcessors() * 2, 8)
+ forkEvery = 1 // Enables different end-to-end test classes use Spark contexts with different configurations
+
+ // Make it so unit tests run on a JAr with Cassandra bridge implementations built in
+ dependsOn(tasks.jar)
+ classpath = project.sourceSets.test.output + configurations.testRuntimeClasspath + files(jar.archiveFile)
+ useJUnit()
+ }
+
+ /* Start: JaCoCo check */
+ jacoco {
+ toolVersion = '0.8.4'
+ }
+ task jacocoMerge(type: JacocoMerge) {
+ destinationFile = file("$buildDir/jacoco/mergedFile.exec")
+ executionData = files("$buildDir/jacoco/test.exec")
+ }
+ jacocoTestReport {
+ dependsOn(tasks.jacocoMerge)
+ executionData(tasks.jacocoMerge.executionData)
+ reports {
+ html.required = true
+ html.destination(file("$buildDir/reports/jacocoHtml"))
+ xml.required = true
+ xml.destination(file("$buildDir/reports/jacocoXml"))
+ }
+ }
+ jacocoTestCoverageVerification {
+ dependsOn(tasks.jacocoMerge)
+ executionData(tasks.jacocoMerge.executionData)
+ violationRules {
+ rule {
+ limit {
+ minimum = 0.5
+ }
+ }
+ }
+ }
+ check {
+ dependsOn(tasks.jacocoTestCoverageVerification)
+ finalizedBy(tasks.jacocoTestReport)
+ }
+ /* End: JaCoCo check */
+}
+
+private void writeBuildVersion(version, projectDir) {
+ def prop = new Properties()
+ prop.setProperty('build-version', "${version}")
+ def hash = System.getenv('GIT_HASH') ?: 'unknown'
+ prop.setProperty('build-rev', hash)
+ def propFile = new File("${projectDir}/src/main/resources/cassandra-analytics-build.properties")
+ propFile.createNewFile()
+ prop.store(propFile.newWriter(), null)
+ logger.info("Wrote to file={}, version={}, hash={}", propFile, version, hash)
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/CassandraBridgeFactory.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/CassandraBridgeFactory.java
new file mode 100644
index 0000000..7e2cd51
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/CassandraBridgeFactory.java
@@ -0,0 +1,135 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.bridge;
+
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStream;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.InvocationTargetException;
+import java.net.URL;
+import java.nio.file.Files;
+import java.util.Arrays;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+
+import com.google.common.base.Preconditions;
+import org.apache.commons.io.FileUtils;
+
+import org.jetbrains.annotations.NotNull;
+
+public final class CassandraBridgeFactory
+{
+ private static final Map<String, CassandraBridge> CASSANDRA_BRIDGES = new ConcurrentHashMap<>(CassandraVersion.values().length);
+
+ private CassandraBridgeFactory()
+ {
+ throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
+ }
+
+ @NotNull
+ public static CassandraVersion getCassandraVersion(@NotNull String version)
+ {
+ CassandraVersionFeatures features = CassandraVersionFeatures.cassandraVersionFeaturesFromCassandraVersion(version);
+ Preconditions.checkArgument(features != null, "Cassandra version " + version + " is not supported");
+ return getCassandraVersion(features);
+ }
+
+ @NotNull
+ public static CassandraVersion getCassandraVersion(@NotNull CassandraVersionFeatures features)
+ {
+ Optional<CassandraVersion> version = Arrays.stream(CassandraVersion.values())
+ .filter(value -> value.versionNumber() == features.getMajorVersion())
+ .findAny();
+ Preconditions.checkArgument(version.isPresent(), "Cassandra features " + features + " are not supported");
+ return version.get();
+ }
+
+ @NotNull
+ public static CassandraBridge get(@NotNull String version)
+ {
+ return get(getCassandraVersion(version));
+ }
+
+ @NotNull
+ public static CassandraBridge get(@NotNull CassandraVersionFeatures features)
+ {
+ return get(getCassandraVersion(features));
+ }
+
+ @NotNull
+ public static CassandraBridge get(@NotNull CassandraVersion version)
+ {
+ String jarBaseName = version.jarBaseName();
+ Preconditions.checkNotNull(jarBaseName, "Cassandra version " + version + " is not supported");
+ return CASSANDRA_BRIDGES.computeIfAbsent(jarBaseName, CassandraBridgeFactory::create);
+ }
+
+ /**
+ * Ensures that every supported Cassandra version has a corresponding Cassandra bridge implementation embedded
+ * into this library's binary as a separate JAR file in the {@code bridges} directory, fails fast otherwise
+ *
+ * @throws IllegalStateException If a Cassandra bridge implementation is missing for a supported Cassandra version
+ */
+ public static void validateBridges()
+ {
+ for (CassandraVersion version : CassandraVersion.implementedVersions())
+ {
+ String jarBaseName = version.jarBaseName();
+ String name = resourceName(jarBaseName);
+ URL locator = CassandraBridgeFactory.class.getResource(name);
+ if (locator == null)
+ {
+ throw new IllegalStateException("Missing Cassandra bridge implementation for version " + version);
+ }
+ }
+ }
+
+ @NotNull
+ private static String resourceName(@NotNull String label)
+ {
+ return "/bridges/" + label + ".jar";
+ }
+
+ @NotNull
+ @SuppressWarnings("unchecked")
+ private static CassandraBridge create(@NotNull String label)
+ {
+ try
+ {
+ String name = resourceName(label);
+ InputStream contents = CassandraBridgeFactory.class.getResourceAsStream(name);
+ File jar = Files.createTempFile(null, ".jar").toFile();
+ FileUtils.copyInputStreamToFile(contents, jar);
+
+ ClassLoader loader = new PostDelegationClassLoader(jar, Thread.currentThread().getContextClassLoader());
+ Class<CassandraBridge> bridge = (Class<CassandraBridge>) loader.loadClass("org.apache.cassandra.bridge.CassandraBridgeImplementation");
+ Constructor<CassandraBridge> constructor = bridge.getConstructor();
+ return constructor.newInstance();
+
+ }
+ catch (IOException | ClassNotFoundException | NoSuchMethodException | InstantiationException
+ | IllegalAccessException | InvocationTargetException exception)
+ {
+ throw new RuntimeException("Failed to create Cassandra bridge for label " + label, exception);
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/CassandraVersionFeatures.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/CassandraVersionFeatures.java
new file mode 100644
index 0000000..089f574
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/CassandraVersionFeatures.java
@@ -0,0 +1,198 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.bridge;
+
+import java.util.Objects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class CassandraVersionFeatures implements Comparable<CassandraVersionFeatures>
+{
+ private static final Pattern VERSION_PATTERN_3 = Pattern.compile("(?:.+-)?([0-9]+)\\.([0-9]+)\\.([0-9]+)([a-zA-Z0-9-]*)");
+ private static final Pattern VERSION_PATTERN_4 = Pattern.compile("(?:.+-)?([0-9]+)\\.([0-9]+)\\.([0-9]+)\\.([0-9]+)[a-zA-Z0-9-]*");
+ private static final Pattern VERSION_PATTERN_SNAPSHOT = Pattern.compile("(?:.+-)?([0-9]+)\\.([0-9]+)-(SNAPSHOT)$");
+
+ protected final int majorVersion;
+ protected final int minorVersion;
+ protected final String suffix;
+
+ public CassandraVersionFeatures(int majorVersion, int minorVersion, @Nullable String suffix)
+ {
+ this.majorVersion = majorVersion;
+ this.minorVersion = minorVersion;
+ this.suffix = suffix;
+ }
+
+ /**
+ * Given a cassandra version string (e.g. cassandra-2.0.14-v3) parse it into a CassandraVersionFeatures object
+ *
+ * @param cassandraVersion the string representing the cassandra version
+ * @return an instance representing the parsed values from the version string
+ */
+ public static CassandraVersionFeatures cassandraVersionFeaturesFromCassandraVersion(
+ @NotNull String cassandraVersion)
+ {
+ String versionCode = getCassandraVersionCode(cassandraVersion);
+ String minorVersionCode = getCassandraMinorVersionCode(cassandraVersion);
+ String versionSuffix = getCassandraVersionSuffix(cassandraVersion);
+
+ return new CassandraVersionFeatures(Integer.parseInt(versionCode), Integer.parseInt(minorVersionCode), versionSuffix);
+ }
+
+ // E.g if cassandra version = cassandra-1.2.11-v1, we return 12;
+ // or if cassandra version = cassandra-4.0-SNAPSHOT, we return 40
+ private static String getCassandraVersionCode(String cassandraVersion)
+ {
+ Matcher matcher = matchVersion(cassandraVersion);
+
+ return matcher.group(1) + matcher.group(2);
+ }
+
+ // E.g if cassandra version = cassandra-1.2.11-v1, we return 11;
+ // or if cassandra version = cassandra-4.0-SNAPSHOT, we return 0
+ private static String getCassandraMinorVersionCode(String cassandraVersion)
+ {
+ Matcher matcher = matchVersion(cassandraVersion);
+ if (matchesSnapshot(matcher.group(3)))
+ {
+ return "0";
+ }
+
+ return matcher.group(3);
+ }
+
+ // E.g if cassandra version = cassandra-1.2.11-v1, we return -v1;
+ // or if cassandra version = cassandra-1.2.11.2-tag, we return 2;
+ // or if cassandra version = cassandra-4.0-SNAPSHOT, we return SNAPSHOT
+ private static String getCassandraVersionSuffix(String cassandraVersion)
+ {
+ Matcher matcher = matchVersion(cassandraVersion);
+ if (matchesSnapshot(matcher.group(3)) || matchesSnapshot(matcher.group(4)))
+ {
+ return "SNAPSHOT";
+ }
+
+ return matcher.group(4);
+ }
+
+ private static boolean matchesSnapshot(String snapshot)
+ {
+ return "SNAPSHOT".equals(snapshot) || "-SNAPSHOT".equals(snapshot);
+ }
+
+ /**
+ * Returns a matched matcher using VERSION_PATTERN; throws if no match
+ */
+ private static Matcher matchVersion(String cassandraVersion)
+ {
+ Matcher matcher = VERSION_PATTERN_4.matcher(cassandraVersion);
+
+ if (!matcher.find())
+ {
+ matcher = VERSION_PATTERN_3.matcher(cassandraVersion);
+
+ if (!matcher.find())
+ {
+ matcher = VERSION_PATTERN_SNAPSHOT.matcher(cassandraVersion);
+ if (!matcher.find())
+ {
+ throw new RuntimeException("cassandraVersion does not match version pattern, pattern=" + VERSION_PATTERN_3
+ + ", version=" + cassandraVersion);
+ }
+ }
+ }
+ return matcher;
+ }
+
+ public int getMajorVersion()
+ {
+ return majorVersion;
+ }
+
+ public int getMinorVersion()
+ {
+ return minorVersion;
+ }
+
+ public String getSuffix()
+ {
+ return suffix;
+ }
+
+ @Override
+ public int compareTo(@NotNull CassandraVersionFeatures that)
+ {
+ int difference = this.getMajorVersion() - that.getMajorVersion();
+ if (difference != 0)
+ {
+ return difference;
+ }
+
+ difference = this.getMinorVersion() - that.getMinorVersion();
+ if (difference != 0)
+ {
+ return difference;
+ }
+
+ // Try to treat suffix as micro version
+ try
+ {
+ return Integer.parseInt(getSuffix()) - Integer.parseInt(that.getSuffix());
+ }
+ catch (NumberFormatException exception)
+ {
+ // One with smallest suffix will be smallest
+ int thisLength = getSuffix() != null ? getSuffix().length() : 0;
+ int thatLength = that.getSuffix() != null ? that.getSuffix().length() : 0;
+ return thisLength - thatLength;
+ }
+ }
+
+ @Override
+ public boolean equals(@Nullable Object other)
+ {
+ if (other == null || !(other instanceof CassandraVersionFeatures))
+ {
+ return false;
+ }
+
+ CassandraVersionFeatures that = (CassandraVersionFeatures) other;
+ return Objects.equals(this.majorVersion, that.majorVersion)
+ && Objects.equals(this.minorVersion, that.minorVersion)
+ && Objects.equals(this.suffix, that.suffix);
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return Objects.hash(majorVersion, minorVersion, suffix);
+ }
+
+ @Override
+ @NotNull
+ public String toString()
+ {
+ return String.format("CassandraVersionFeatures{majorVersion=%d, minorVersion=%d, suffix='%s'}",
+ majorVersion, minorVersion, suffix);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/PostDelegationClassLoader.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/PostDelegationClassLoader.java
new file mode 100644
index 0000000..c1e26a7
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/bridge/PostDelegationClassLoader.java
@@ -0,0 +1,71 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.bridge;
+
+import java.io.File;
+import java.net.MalformedURLException;
+import java.net.URL;
+import java.net.URLClassLoader;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * This custom implementation of a {@link ClassLoader} enables deferred execution-time loading of a particular version
+ * of class hierarchy from one of many embedded the {@code cassandra-all} library JARs. It first attempts to load any
+ * requested class from the extracted JAR, and resorts to using the parent class loader when the the class is not there.
+ * This behavior is opposite to the one of standard {@link URLClassLoader}, which invokes its parent class loader first.
+ */
+public class PostDelegationClassLoader extends URLClassLoader
+{
+ public PostDelegationClassLoader(@NotNull File jar, @Nullable ClassLoader parent) throws MalformedURLException
+ {
+ super(new URL[]{jar.toURI().toURL()}, parent);
+ }
+
+ @Override
+ @Nullable
+ protected synchronized Class<?> loadClass(@Nullable String name, boolean resolve) throws ClassNotFoundException
+ {
+ // First, check if the class has already been loaded
+ Class<?> type = findLoadedClass(name);
+ if (type == null)
+ {
+ try
+ {
+ type = findClass(name);
+ }
+ catch (ClassNotFoundException | SecurityException | LinkageError exception)
+ {
+ // ClassNotFoundException thrown if class not found
+ }
+ if (type == null)
+ {
+ // If not found, then invoke findClass in order to find the class
+ type = super.loadClass(name, false);
+ }
+ }
+ if (resolve)
+ {
+ resolveClass(type);
+ }
+ return type;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/ExecutorHolder.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/ExecutorHolder.java
new file mode 100644
index 0000000..dbb4df1
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/ExecutorHolder.java
@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.clients;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+/**
+ * A class that holds Executors used during the Cassandra Analytics execution
+ */
+public final class ExecutorHolder
+{
+ // TODO: Make the size configurable
+ public static final ExecutorService EXECUTOR_SERVICE = Executors.newFixedThreadPool(64,
+ new ThreadFactoryBuilder().setNameFormat("cassandra-analytics-%d")
+ .setDaemon(true)
+ .build());
+
+ private ExecutorHolder()
+ {
+ throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/Sidecar.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/Sidecar.java
new file mode 100644
index 0000000..71d8f43
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/Sidecar.java
@@ -0,0 +1,405 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.clients;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import o.a.c.sidecar.client.shaded.io.vertx.core.Vertx;
+import o.a.c.sidecar.client.shaded.io.vertx.core.VertxOptions;
+import org.apache.cassandra.sidecar.client.HttpClientConfig;
+import org.apache.cassandra.sidecar.client.SidecarClient;
+import org.apache.cassandra.sidecar.client.SidecarConfig;
+import org.apache.cassandra.sidecar.client.SidecarInstance;
+import org.apache.cassandra.sidecar.client.SidecarInstancesProvider;
+import org.apache.cassandra.sidecar.client.VertxHttpClient;
+import org.apache.cassandra.sidecar.client.VertxRequestExecutor;
+import org.apache.cassandra.sidecar.client.retry.ExponentialBackoffRetryPolicy;
+import org.apache.cassandra.sidecar.client.retry.RetryPolicy;
+import org.apache.cassandra.sidecar.common.NodeSettings;
+import org.apache.cassandra.spark.bulkwriter.BulkSparkConf;
+import org.apache.cassandra.spark.data.FileType;
+import org.apache.cassandra.spark.utils.BuildInfo;
+import org.apache.cassandra.spark.utils.MapUtils;
+
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_CHUNK_BUFFER_OVERRIDE;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_CHUNK_BUFFER_SIZE;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_MAX_BUFFER_OVERRIDE;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_MAX_BUFFER_SIZE;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_MAX_MILLIS_TO_SLEEP;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_MAX_POOL_SIZE;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_MAX_RETRIES;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_MILLIS_TO_SLEEP;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_SIDECAR_PORT;
+import static org.apache.cassandra.spark.utils.Properties.DEFAULT_TIMEOUT_SECONDS;
+
+/**
+ * A helper class that encapsulates configuration for the Spark Bulk Reader and Writer and helper methods to build the
+ * {@link SidecarClient}
+ */
+public final class Sidecar
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(Sidecar.class);
+
+ private Sidecar()
+ {
+ throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
+ }
+
+ public static SidecarClient from(SidecarInstancesProvider sidecarInstancesProvider,
+ ClientConfig config,
+ SslConfig sslConfig) throws IOException
+ {
+ Vertx vertx = Vertx.vertx(new VertxOptions().setUseDaemonThread(true).setWorkerPoolSize(config.maxPoolSize()));
+
+ HttpClientConfig.Builder<?> builder = new HttpClientConfig.Builder<>()
+ .ssl(false)
+ .timeoutMillis(TimeUnit.SECONDS.toMillis(config.timeoutSeconds()))
+ .idleTimeoutMillis((int) TimeUnit.SECONDS.toMillis(config.timeoutSeconds()))
+ .receiveBufferSize((int) config.chunkBufferSize())
+ .maxChunkSize((int) config.maxBufferSize())
+ .userAgent(BuildInfo.READER_USER_AGENT);
+
+ if (sslConfig != null)
+ {
+ builder = builder
+ .ssl(true)
+ .keyStoreInputStream(sslConfig.keyStoreInputStream())
+ .keyStorePassword(sslConfig.keyStorePassword())
+ .keyStoreType(sslConfig.keyStoreType())
+ .trustStoreInputStream(sslConfig.trustStoreInputStream())
+ .trustStorePassword(sslConfig.trustStorePassword())
+ .trustStoreType(sslConfig.trustStoreType());
+ }
+
+ HttpClientConfig httpClientConfig = builder.build();
+
+ SidecarConfig sidecarConfig = new SidecarConfig.Builder()
+ .maxRetries(config.maxRetries())
+ .retryDelayMillis(config.millisToSleep())
+ .maxRetryDelayMillis(config.maxMillisToSleep())
+ .build();
+
+ return buildClient(sidecarConfig, vertx, httpClientConfig, sidecarInstancesProvider);
+ }
+
+ public static SidecarClient from(SidecarInstancesProvider sidecarInstancesProvider, BulkSparkConf conf)
+ {
+ Vertx vertx = Vertx.vertx(new VertxOptions().setUseDaemonThread(true)
+ .setWorkerPoolSize(conf.getMaxHttpConnections()));
+
+ HttpClientConfig httpClientConfig = new HttpClientConfig.Builder<>()
+ .timeoutMillis(conf.getHttpResponseTimeoutMs())
+ .idleTimeoutMillis(conf.getHttpConnectionTimeoutMs())
+ .userAgent(BuildInfo.WRITER_USER_AGENT)
+ .keyStoreInputStream(conf.getKeyStore())
+ .keyStorePassword(conf.getKeyStorePassword())
+ .keyStoreType(conf.getKeyStoreTypeOrDefault())
+ .trustStoreInputStream(conf.getTrustStore())
+ .trustStorePassword(conf.getTrustStorePasswordOrDefault())
+ .trustStoreType(conf.getTrustStoreTypeOrDefault())
+ .ssl(conf.hasKeystoreAndKeystorePassword())
+ .build();
+
+ SidecarConfig sidecarConfig = new SidecarConfig.Builder()
+ .maxRetries(conf.getSidecarRequestRetries())
+ .retryDelayMillis(TimeUnit.SECONDS.toMillis(conf.getSidecarRequestRetryDelayInSeconds()))
+ .maxRetryDelayMillis(TimeUnit.SECONDS.toMillis(conf.getSidecarRequestMaxRetryDelayInSeconds()))
+ .build();
+
+ return buildClient(sidecarConfig, vertx, httpClientConfig, sidecarInstancesProvider);
+ }
+
+ public static SidecarClient buildClient(SidecarConfig sidecarConfig,
+ Vertx vertx,
+ HttpClientConfig httpClientConfig,
+ SidecarInstancesProvider clusterConfig)
+ {
+ RetryPolicy defaultRetryPolicy = new ExponentialBackoffRetryPolicy(sidecarConfig.maxRetries(),
+ sidecarConfig.retryDelayMillis(),
+ sidecarConfig.maxRetryDelayMillis());
+
+ VertxHttpClient vertxHttpClient = new VertxHttpClient(vertx, httpClientConfig);
+ VertxRequestExecutor requestExecutor = new VertxRequestExecutor(vertxHttpClient);
+ return new SidecarClient(clusterConfig, requestExecutor, sidecarConfig, defaultRetryPolicy);
+ }
+
+ public static List<NodeSettings> allNodeSettingsBlocking(BulkSparkConf conf,
+ SidecarClient client,
+ Set<? extends SidecarInstance> instances)
+ {
+ List<CompletableFuture<Void>> futures = new ArrayList<>();
+ List<NodeSettings> allNodeSettings = Collections.synchronizedList(new ArrayList<>());
+ for (SidecarInstance instance : instances)
+ {
+ futures.add(client.nodeSettings(instance)
+ .exceptionally(throwable -> {
+ LOGGER.warn(String.format("Failed to execute node settings on instance=%s",
+ instance), throwable);
+ return null;
+ })
+ .thenAccept(nodeSettings -> {
+ if (nodeSettings != null)
+ {
+ allNodeSettings.add(nodeSettings);
+ }
+ }));
+ }
+ try
+ {
+ CompletableFuture.allOf(futures.toArray(new CompletableFuture[0]))
+ .get(conf.getSidecarRequestMaxRetryDelayInSeconds(), TimeUnit.SECONDS);
+ }
+ catch (ExecutionException | InterruptedException exception)
+ {
+ throw new RuntimeException(exception);
+ }
+ catch (TimeoutException exception)
+ {
+ // Any futures that have already completed will have put their results into `allNodeSettings`
+ // at this point. Cancel any remaining futures and move on.
+ for (CompletableFuture<Void> future : futures)
+ {
+ future.cancel(true);
+ }
+ }
+ long successCount = allNodeSettings.size();
+ if (successCount == 0)
+ {
+ throw new RuntimeException(String.format("Unable to determine the node settings. 0/%d instances available.",
+ instances.size()));
+ }
+ else if (successCount < instances.size())
+ {
+ LOGGER.debug("{}/{} instances were used to determine the node settings",
+ successCount, instances.size());
+ }
+ return allNodeSettings;
+ }
+
+ public static final class ClientConfig
+ {
+ public static final String SIDECAR_PORT = "sidecar_port";
+ public static final String MAX_RETRIES_KEY = "maxRetries";
+ public static final String DEFAULT_MILLIS_TO_SLEEP_KEY = "defaultMillisToSleep";
+ public static final String MAX_MILLIS_TO_SLEEP_KEY = "maxMillisToSleep";
+ public static final String MAX_BUFFER_SIZE_BYTES_KEY = "maxBufferSizeBytes";
+ public static final String CHUNK_BUFFER_SIZE_BYTES_KEY = "chunkBufferSizeBytes";
+ public static final String MAX_POOL_SIZE_KEY = "maxPoolSize";
+ public static final String TIMEOUT_SECONDS_KEY = "timeoutSeconds";
+
+ private final int port;
+ private final int maxRetries;
+ private final int maxPoolSize;
+ private final int timeoutSeconds;
+ private final long millisToSleep;
+ private final long maxMillisToSleep;
+ private final long maxBufferSize;
+ private final long chunkSize;
+ private final Map<FileType, Long> maxBufferOverride;
+ private final Map<FileType, Long> chunkBufferOverride;
+
+ // CHECKSTYLE IGNORE: Constructor with many parameters
+ private ClientConfig(int port,
+ int maxRetries,
+ long millisToSleep,
+ long maxMillisToSleep,
+ long maxBufferSize,
+ long chunkSize,
+ int maxPoolSize,
+ int timeoutSeconds,
+ Map<FileType, Long> maxBufferOverride,
+ Map<FileType, Long> chunkBufferOverride)
+ {
+ this.port = port;
+ this.maxRetries = maxRetries;
+ this.millisToSleep = millisToSleep;
+ this.maxMillisToSleep = maxMillisToSleep;
+ this.maxBufferSize = maxBufferSize;
+ this.chunkSize = chunkSize;
+ this.maxPoolSize = maxPoolSize;
+ this.timeoutSeconds = timeoutSeconds;
+ this.maxBufferOverride = maxBufferOverride;
+ this.chunkBufferOverride = chunkBufferOverride;
+ }
+
+ public int port()
+ {
+ return port;
+ }
+
+ public int maxRetries()
+ {
+ return maxRetries;
+ }
+
+ public long millisToSleep()
+ {
+ return millisToSleep;
+ }
+
+ public long maxMillisToSleep()
+ {
+ return maxMillisToSleep;
+ }
+
+ public long maxBufferSize()
+ {
+ return maxBufferSize(FileType.DATA);
+ }
+
+ public long maxBufferSize(FileType fileType)
+ {
+ return maxBufferOverride.getOrDefault(fileType, maxBufferSize);
+ }
+
+ public Map<FileType, Long> maxBufferOverride()
+ {
+ return maxBufferOverride;
+ }
+
+ public long chunkBufferSize()
+ {
+ return chunkBufferSize(FileType.DATA);
+ }
+
+ public long chunkBufferSize(FileType fileType)
+ {
+ return chunkBufferOverride.getOrDefault(fileType, chunkSize);
+ }
+
+ public Map<FileType, Long> chunkBufferOverride()
+ {
+ return chunkBufferOverride;
+ }
+
+ public int maxPoolSize()
+ {
+ return maxPoolSize;
+ }
+
+ public int timeoutSeconds()
+ {
+ return timeoutSeconds;
+ }
+
+ public static ClientConfig create()
+ {
+ return ClientConfig.create(DEFAULT_SIDECAR_PORT, DEFAULT_MAX_RETRIES, DEFAULT_MILLIS_TO_SLEEP);
+ }
+
+ public static ClientConfig create(int port)
+ {
+ return ClientConfig.create(port, DEFAULT_MAX_RETRIES, DEFAULT_MILLIS_TO_SLEEP);
+ }
+
+ public static ClientConfig create(int port, int maxRetries, long millisToSleep)
+ {
+ return ClientConfig.create(port,
+ maxRetries,
+ millisToSleep,
+ DEFAULT_MAX_MILLIS_TO_SLEEP,
+ DEFAULT_MAX_BUFFER_SIZE,
+ DEFAULT_CHUNK_BUFFER_SIZE,
+ DEFAULT_MAX_POOL_SIZE,
+ DEFAULT_TIMEOUT_SECONDS,
+ DEFAULT_MAX_BUFFER_OVERRIDE,
+ DEFAULT_CHUNK_BUFFER_OVERRIDE);
+ }
+
+ public static ClientConfig create(Map<String, String> options)
+ {
+ return create(MapUtils.getInt(options, SIDECAR_PORT, DEFAULT_SIDECAR_PORT),
+ MapUtils.getInt(options, MAX_RETRIES_KEY, DEFAULT_MAX_RETRIES),
+ MapUtils.getLong(options, DEFAULT_MILLIS_TO_SLEEP_KEY, DEFAULT_MILLIS_TO_SLEEP),
+ MapUtils.getLong(options, MAX_MILLIS_TO_SLEEP_KEY, DEFAULT_MAX_MILLIS_TO_SLEEP),
+ MapUtils.getLong(options, MAX_BUFFER_SIZE_BYTES_KEY, DEFAULT_MAX_BUFFER_SIZE),
+ MapUtils.getLong(options, CHUNK_BUFFER_SIZE_BYTES_KEY, DEFAULT_CHUNK_BUFFER_SIZE),
+ MapUtils.getInt(options, MAX_POOL_SIZE_KEY, DEFAULT_MAX_POOL_SIZE),
+ MapUtils.getInt(options, TIMEOUT_SECONDS_KEY, DEFAULT_TIMEOUT_SECONDS),
+ buildMaxBufferOverride(options, DEFAULT_MAX_BUFFER_OVERRIDE),
+ buildChunkBufferOverride(options, DEFAULT_CHUNK_BUFFER_OVERRIDE)
+ );
+ }
+
+ public static Map<FileType, Long> buildMaxBufferOverride(Map<String, String> options,
+ Map<FileType, Long> defaultValue)
+ {
+ return buildOverrideMap(MAX_BUFFER_SIZE_BYTES_KEY, options, defaultValue);
+ }
+
+ public static Map<FileType, Long> buildChunkBufferOverride(Map<String, String> options,
+ Map<FileType, Long> defaultValue)
+ {
+ return buildOverrideMap(CHUNK_BUFFER_SIZE_BYTES_KEY, options, defaultValue);
+ }
+
+ private static Map<FileType, Long> buildOverrideMap(String keyPrefix,
+ Map<String, String> options,
+ Map<FileType, Long> defaultValue)
+ {
+ Map<FileType, Long> result = new HashMap<>(defaultValue);
+ for (FileType type : FileType.values())
+ {
+ // Override with DataSourceOptions if set, e.g. maxBufferSizeBytes_Index.db
+ String key = MapUtils.lowerCaseKey(String.format("%s_%s", keyPrefix, type.getFileSuffix()));
+ Optional.ofNullable(options.get(key)).map(Long::parseLong).ifPresent(s -> result.put(type, s));
+ }
+ return result;
+ }
+
+ // CHECKSTYLE IGNORE: Method with many parameters
+ public static ClientConfig create(int port,
+ int maxRetries,
+ long millisToSleep,
+ long maxMillisToSleep,
+ long maxBufferSizeBytes,
+ long chunkSizeBytes,
+ int maxPoolSize,
+ int timeoutSeconds,
+ Map<FileType, Long> maxBufferOverride,
+ Map<FileType, Long> chunkBufferOverride)
+ {
+ return new ClientConfig(port,
+ maxRetries,
+ millisToSleep,
+ maxMillisToSleep,
+ maxBufferSizeBytes,
+ chunkSizeBytes,
+ maxPoolSize,
+ timeoutSeconds,
+ maxBufferOverride,
+ chunkBufferOverride);
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SidecarInstanceImpl.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SidecarInstanceImpl.java
new file mode 100644
index 0000000..2e58be6
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SidecarInstanceImpl.java
@@ -0,0 +1,122 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.clients;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.util.Objects;
+
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+import org.apache.cassandra.sidecar.client.SidecarInstance;
+
+/**
+ * A simple implementation of the {@link SidecarInstance} interface
+ */
+public class SidecarInstanceImpl implements Serializable, SidecarInstance
+{
+ private static final long serialVersionUID = -8650006905764842232L;
+ private static final Logger LOGGER = LoggerFactory.getLogger(SidecarInstanceImpl.class);
+
+ private int port;
+ private String hostname;
+
+ /**
+ * Constructs a new Sidecar instance with the given {@code port} and {@code hostname}
+ *
+ * @param hostname the host name where Sidecar is running
+ * @param port the port where Sidecar is running
+ */
+ public SidecarInstanceImpl(String hostname, int port)
+ {
+ Preconditions.checkArgument(0 < port && port <= 65535,
+ "The Sidecar port number must be in the range 1-65535: %s", port);
+ this.port = port;
+ this.hostname = Objects.requireNonNull(hostname, "The Sidecar hostname must be non-null");
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public int port()
+ {
+ return port;
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public String hostname()
+ {
+ return hostname;
+ }
+
+ /**
+ * {@inheritDoc}
+ */
+ @Override
+ public String toString()
+ {
+ return String.format("SidecarInstanceImpl{hostname='%s', port=%d}", hostname, port);
+ }
+
+ // JDK Serialization
+
+ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException
+ {
+ LOGGER.warn("Falling back to JDK deserialization");
+ hostname = in.readUTF();
+ port = in.readInt();
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException, ClassNotFoundException
+ {
+ LOGGER.warn("Falling back to JDK serialization");
+ out.writeUTF(hostname);
+ out.writeInt(port);
+ }
+
+ // Kryo Serialization
+
+ public static class Serializer extends com.esotericsoftware.kryo.Serializer<SidecarInstanceImpl>
+ {
+ @Override
+ public void write(Kryo kryo, Output out, SidecarInstanceImpl sidecarInstance)
+ {
+ out.writeString(sidecarInstance.hostname);
+ out.writeInt(sidecarInstance.port);
+ }
+
+ @Override
+ public SidecarInstanceImpl read(Kryo kryo, Input input, Class<SidecarInstanceImpl> type)
+ {
+ return new SidecarInstanceImpl(input.readString(), input.readInt());
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SidecarStreamConsumerAdapter.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SidecarStreamConsumerAdapter.java
new file mode 100644
index 0000000..0e48f0b
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SidecarStreamConsumerAdapter.java
@@ -0,0 +1,107 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.clients;
+
+import java.nio.ByteBuffer;
+import java.util.Objects;
+
+import org.apache.cassandra.sidecar.client.StreamBuffer;
+import org.apache.cassandra.spark.utils.streaming.StreamConsumer;
+
+/**
+ * A delegate class that connects Sidecar's {@link org.apache.cassandra.sidecar.client.StreamConsumer} to the analytics
+ * {@link StreamConsumer}
+ */
+public class SidecarStreamConsumerAdapter implements org.apache.cassandra.sidecar.client.StreamConsumer
+{
+ private final StreamConsumer delegate;
+
+ public SidecarStreamConsumerAdapter(StreamConsumer delegate)
+ {
+ this.delegate = delegate;
+ }
+
+ @Override
+ public void onRead(org.apache.cassandra.sidecar.client.StreamBuffer streamBuffer)
+ {
+ delegate.onRead(SidecarStreamBufferWrapper.wrap(streamBuffer));
+ }
+
+ @Override
+ public void onComplete()
+ {
+ delegate.onEnd();
+ }
+
+ @Override
+ public void onError(Throwable throwable)
+ {
+ delegate.onError(throwable);
+ }
+
+ /**
+ * A {@link org.apache.cassandra.spark.utils.streaming.StreamBuffer} implementations that internally wraps
+ * Sidecar's {@link org.apache.cassandra.sidecar.client.StreamBuffer}
+ */
+ static final class SidecarStreamBufferWrapper implements org.apache.cassandra.spark.utils.streaming.StreamBuffer
+ {
+ public final org.apache.cassandra.sidecar.client.StreamBuffer buffer;
+
+ private SidecarStreamBufferWrapper(StreamBuffer buffer)
+ {
+ this.buffer = Objects.requireNonNull(buffer, "the buffer parameter must be non-null");
+ }
+
+ public static SidecarStreamBufferWrapper wrap(StreamBuffer buffer)
+ {
+ return new SidecarStreamBufferWrapper(buffer);
+ }
+
+ @Override
+ public void getBytes(int index, ByteBuffer destination, int length)
+ {
+ buffer.copyBytes(index, destination, length);
+ }
+
+ @Override
+ public void getBytes(int index, byte[] destination, int destinationIndex, int length)
+ {
+ buffer.copyBytes(index, destination, destinationIndex, length);
+ }
+
+ @Override
+ public byte getByte(int index)
+ {
+ return buffer.getByte(index);
+ }
+
+ @Override
+ public int readableBytes()
+ {
+ return buffer.readableBytes();
+ }
+
+ @Override
+ public void release()
+ {
+ buffer.release();
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SslConfig.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SslConfig.java
new file mode 100644
index 0000000..2fd72b9
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/clients/SslConfig.java
@@ -0,0 +1,486 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.clients;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.Base64;
+import java.util.Map;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+import org.apache.cassandra.spark.utils.MapUtils;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Encapsulates SSL configuration for Sidecar
+ */
+public class SslConfig implements Serializable
+{
+ private static final long serialVersionUID = -3844712192096436932L;
+ private static final Logger LOGGER = LoggerFactory.getLogger(SslConfig.class);
+ public static final String KEYSTORE_PATH = "KEYSTORE_PATH";
+ public static final String KEYSTORE_BASE64_ENCODED = "KEYSTORE_BASE64_ENCODED";
+ public static final String KEYSTORE_PASSWORD = "KEYSTORE_PASSWORD";
+ public static final String KEYSTORE_TYPE = "KEYSTORE_TYPE";
+ public static final String DEFAULT_KEYSTORE_TYPE = "PKCS12";
+ public static final String TRUSTSTORE_PATH = "TRUSTSTORE_PATH";
+ public static final String TRUSTSTORE_BASE64_ENCODED = "TRUSTSTORE_BASE64_ENCODED";
+ public static final String TRUSTSTORE_PASSWORD = "TRUSTSTORE_PASSWORD";
+ public static final String TRUSTSTORE_TYPE = "TRUSTSTORE_TYPE";
+ public static final String DEFAULT_TRUSTSTORE_TYPE = "JKS";
+ protected String keyStorePath;
+ protected String base64EncodedKeyStore;
+ protected String keyStorePassword;
+ protected String keyStoreType;
+ protected String trustStorePath;
+ protected String base64EncodedTrustStore;
+ protected String trustStorePassword;
+ protected String trustStoreType;
+
+ protected SslConfig(Builder<?> builder)
+ {
+ keyStorePath = builder.keyStorePath;
+ base64EncodedKeyStore = builder.base64EncodedKeyStore;
+ keyStorePassword = builder.keyStorePassword;
+ keyStoreType = builder.keyStoreType;
+ trustStorePath = builder.trustStorePath;
+ base64EncodedTrustStore = builder.base64EncodedTrustStore;
+ trustStorePassword = builder.trustStorePassword;
+ trustStoreType = builder.trustStoreType;
+ }
+
+ /**
+ * @return an input stream to the keystore source
+ * @throws IOException when an IO exception occurs
+ */
+ public InputStream keyStoreInputStream() throws IOException
+ {
+ if (keyStorePath != null)
+ {
+ return Files.newInputStream(Paths.get(keyStorePath));
+ }
+ else if (base64EncodedKeyStore != null)
+ {
+ return new ByteArrayInputStream(Base64.getDecoder().decode(base64EncodedKeyStore));
+ }
+ else
+ {
+ // It should never reach here
+ throw new RuntimeException("keyStorePath or encodedKeyStore must be provided");
+ }
+ }
+
+ /**
+ * @return the path to the key store file
+ */
+ public String keyStorePath()
+ {
+ return keyStorePath;
+ }
+
+ /**
+ * @return the encoded key store
+ */
+ public String base64EncodedKeyStore()
+ {
+ return base64EncodedKeyStore;
+ }
+
+ /**
+ * @return the key store password
+ */
+ public String keyStorePassword()
+ {
+ return keyStorePassword;
+ }
+
+ /**
+ * @return the key store type
+ */
+ public String keyStoreType()
+ {
+ return keyStoreType != null ? keyStoreType : DEFAULT_KEYSTORE_TYPE;
+ }
+
+ /**
+ * @return an input stream to the truststore source
+ * @throws IOException when an IO exception occurs
+ */
+ public InputStream trustStoreInputStream() throws IOException
+ {
+ if (trustStorePath != null)
+ {
+ return Files.newInputStream(Paths.get(trustStorePath));
+ }
+ else if (base64EncodedTrustStore != null)
+ {
+ return new ByteArrayInputStream(Base64.getDecoder().decode(base64EncodedTrustStore));
+ }
+ return null;
+ }
+
+ /**
+ * @return the path to the trust store
+ */
+ public String trustStorePath()
+ {
+ return trustStorePath;
+ }
+
+ /**
+ * @return the encoded trust store
+ */
+ public String base64EncodedTrustStore()
+ {
+ return base64EncodedTrustStore;
+ }
+
+ /**
+ * @return the trust store password
+ */
+ public String trustStorePassword()
+ {
+ return trustStorePassword;
+ }
+
+ /**
+ * @return the trust store type
+ */
+ public String trustStoreType()
+ {
+ return trustStoreType != null ? trustStoreType : DEFAULT_TRUSTSTORE_TYPE;
+ }
+
+ // JDK Serialization
+
+ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException
+ {
+ LOGGER.warn("Falling back to JDK deserialization");
+ this.keyStorePath = readNullableString(in);
+ this.base64EncodedKeyStore = readNullableString(in);
+ this.keyStorePassword = readNullableString(in);
+ this.keyStoreType = readNullableString(in);
+ this.trustStorePath = readNullableString(in);
+ this.base64EncodedTrustStore = readNullableString(in);
+ this.trustStorePassword = readNullableString(in);
+ this.trustStoreType = readNullableString(in);
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException
+ {
+ LOGGER.warn("Falling back to JDK serialization");
+ writeNullableString(out, keyStorePath);
+ writeNullableString(out, base64EncodedKeyStore);
+ writeNullableString(out, keyStorePassword);
+ writeNullableString(out, keyStoreType);
+ writeNullableString(out, trustStorePath);
+ writeNullableString(out, base64EncodedTrustStore);
+ writeNullableString(out, trustStorePassword);
+ writeNullableString(out, trustStoreType);
+ }
+
+ private String readNullableString(ObjectInputStream in) throws IOException
+ {
+ return in.readBoolean() ? in.readUTF() : null;
+ }
+
+ private void writeNullableString(ObjectOutputStream out, String string) throws IOException
+ {
+ if (string != null)
+ {
+ out.writeBoolean(true);
+ out.writeUTF(string);
+ }
+ else
+ {
+ out.writeBoolean(false);
+ }
+ }
+
+ // Kryo
+
+ public static class Serializer extends com.esotericsoftware.kryo.Serializer<SslConfig>
+ {
+ public SslConfig read(Kryo kryo, Input in, Class type)
+ {
+ return new Builder<>()
+ .keyStorePath(in.readString())
+ .base64EncodedKeyStore(in.readString())
+ .keyStorePassword(in.readString())
+ .keyStoreType(in.readString())
+ .trustStorePath(in.readString())
+ .base64EncodedTrustStore(in.readString())
+ .trustStorePassword(in.readString())
+ .trustStoreType(in.readString())
+ .build();
+ }
+
+ public void write(Kryo kryo, Output out, SslConfig config)
+ {
+ out.writeString(config.keyStorePath);
+ out.writeString(config.base64EncodedKeyStore);
+ out.writeString(config.keyStorePassword);
+ out.writeString(config.keyStoreType);
+ out.writeString(config.trustStorePath);
+ out.writeString(config.base64EncodedTrustStore);
+ out.writeString(config.trustStorePassword);
+ out.writeString(config.trustStoreType);
+ }
+ }
+
+ @Nullable
+ public static SslConfig create(Map<String, String> options)
+ {
+ String keyStorePath = MapUtils.getOrDefault(options, KEYSTORE_PATH, null);
+ String encodedKeyStore = MapUtils.getOrDefault(options, KEYSTORE_BASE64_ENCODED, null);
+ String keyStorePassword = MapUtils.getOrDefault(options, KEYSTORE_PASSWORD, null);
+ String keyStoreType = MapUtils.getOrDefault(options, KEYSTORE_TYPE, null);
+ String trustStorePath = MapUtils.getOrDefault(options, TRUSTSTORE_PATH, null);
+ String encodedTrustStore = MapUtils.getOrDefault(options, TRUSTSTORE_BASE64_ENCODED, null);
+ String trustStorePassword = MapUtils.getOrDefault(options, TRUSTSTORE_PASSWORD, null);
+ String trustStoreType = MapUtils.getOrDefault(options, TRUSTSTORE_TYPE, null);
+
+ // If any of the values are provided we try to create a valid SecretsConfig object
+ if (keyStorePath != null
+ || encodedKeyStore != null
+ || keyStorePassword != null
+ || keyStoreType != null
+ || trustStorePath != null
+ || encodedTrustStore != null
+ || trustStorePassword != null
+ || trustStoreType != null)
+ {
+ Builder<?> validatedConfig = new Builder<>()
+ .keyStorePath(keyStorePath)
+ .base64EncodedKeyStore(encodedKeyStore)
+ .keyStorePassword(keyStorePassword)
+ .keyStoreType(keyStoreType)
+ .trustStorePath(trustStorePath)
+ .base64EncodedTrustStore(encodedTrustStore)
+ .trustStorePassword(trustStorePassword)
+ .trustStoreType(trustStoreType)
+ .validate();
+ LOGGER.info("Valid SSL configuration");
+ return validatedConfig.build();
+ }
+
+ LOGGER.warn("No SSL configured");
+ return null;
+ }
+
+ /**
+ * {@code SslConfig} builder static inner class
+ *
+ * @param <T> itself
+ */
+ public static class Builder<T extends SslConfig.Builder<T>>
+ {
+ protected String keyStorePath;
+ protected String base64EncodedKeyStore;
+ protected String keyStorePassword;
+ protected String keyStoreType;
+ protected String trustStorePath;
+ protected String base64EncodedTrustStore;
+ protected String trustStorePassword;
+ protected String trustStoreType;
+
+ /**
+ * @return a reference to itself
+ */
+ @SuppressWarnings("unchecked")
+ protected T self()
+ {
+ return (T) this;
+ }
+
+ /**
+ * Sets the {@code keyStorePath} and returns a reference to this Builder enabling method chaining
+ *
+ * @param keyStorePath the {@code keyStorePath} to set
+ * @return a reference to this Builder
+ */
+ public T keyStorePath(String keyStorePath)
+ {
+ this.keyStorePath = keyStorePath;
+ return self();
+ }
+
+ /**
+ * Sets the {@code base64EncodedKeyStore} and returns a reference to this Builder enabling method chaining
+ *
+ * @param base64EncodedKeyStore the {@code base64EncodedKeyStore} to set
+ * @return a reference to this Builder
+ */
+ public T base64EncodedKeyStore(String base64EncodedKeyStore)
+ {
+ this.base64EncodedKeyStore = base64EncodedKeyStore;
+ return self();
+ }
+
+ /**
+ * Sets the {@code keyStorePassword} and returns a reference to this Builder enabling method chaining
+ *
+ * @param keyStorePassword the {@code keyStorePassword} to set
+ * @return a reference to this Builder
+ */
+ public T keyStorePassword(String keyStorePassword)
+ {
+ this.keyStorePassword = keyStorePassword;
+ return self();
+ }
+
+ /**
+ * Sets the {@code keyStoreType} and returns a reference to this Builder enabling method chaining
+ *
+ * @param keyStoreType the {@code keyStoreType} to set
+ * @return a reference to this Builder
+ */
+ public T keyStoreType(String keyStoreType)
+ {
+ this.keyStoreType = keyStoreType;
+ return self();
+ }
+
+ /**
+ * Sets the {@code trustStorePath} and returns a reference to this Builder enabling method chaining
+ *
+ * @param trustStorePath the {@code trustStorePath} to set
+ * @return a reference to this Builder
+ */
+ public T trustStorePath(String trustStorePath)
+ {
+ this.trustStorePath = trustStorePath;
+ return self();
+ }
+
+ /**
+ * Sets the {@code base64EncodedTrustStore} and returns a reference to this Builder enabling method chaining
+ *
+ * @param base64EncodedTrustStore the {@code base64EncodedTrustStore} to set
+ * @return a reference to this Builder
+ */
+ public T base64EncodedTrustStore(String base64EncodedTrustStore)
+ {
+ this.base64EncodedTrustStore = base64EncodedTrustStore;
+ return self();
+ }
+
+ /**
+ * Sets the {@code trustStorePassword} and returns a reference to this Builder enabling method chaining
+ *
+ * @param trustStorePassword the {@code trustStorePassword} to set
+ * @return a reference to this Builder
+ */
+ public T trustStorePassword(String trustStorePassword)
+ {
+ this.trustStorePassword = trustStorePassword;
+ return self();
+ }
+
+ /**
+ * Sets the {@code trustStoreType} and returns a reference to this Builder enabling method chaining
+ *
+ * @param trustStoreType the {@code trustStoreType} to set
+ * @return a reference to this Builder
+ */
+ public T trustStoreType(String trustStoreType)
+ {
+ this.trustStoreType = trustStoreType;
+ return self();
+ }
+
+ /**
+ * Returns a {@code SslConfig} built from the parameters previously set
+ *
+ * @return a {@code SslConfig} built with parameters of this {@code SslConfig.Builder}
+ */
+ public SslConfig build()
+ {
+ return new SslConfig(this);
+ }
+
+ protected T validate()
+ {
+ if (keyStorePath != null && base64EncodedKeyStore != null)
+ {
+ throw new IllegalArgumentException(
+ String.format("Both '%s' and '%s' options were provided. Only one of the options can be provided",
+ KEYSTORE_PATH, KEYSTORE_BASE64_ENCODED));
+ }
+
+ if (keyStorePassword != null)
+ {
+ // When the keystore password is provided, either the key store path or
+ // the encoded key store must be provided
+ if (keyStorePath == null && base64EncodedKeyStore == null)
+ {
+ throw new IllegalArgumentException(
+ String.format("One of the '%s' or '%s' options must be provided when the '%s' option is provided",
+ KEYSTORE_PATH, KEYSTORE_BASE64_ENCODED, KEYSTORE_PASSWORD));
+ }
+ }
+ else
+ {
+
+ throw new IllegalArgumentException(
+ String.format("The '%s' option must be provided when either the '%s' or '%s' options are provided",
+ KEYSTORE_PASSWORD, KEYSTORE_PATH, KEYSTORE_BASE64_ENCODED));
+ }
+
+ if (trustStorePassword != null)
+ {
+ // Only one of trust store path or encoded trust store must be provided, not both
+ if (trustStorePath != null && base64EncodedTrustStore != null)
+ {
+ throw new IllegalArgumentException(
+ String.format("Both '%s' and '%s' options were provided. Only one of the options can be provided",
+ TRUSTSTORE_PATH, TRUSTSTORE_BASE64_ENCODED));
+ }
+
+ // When the trust store password is provided, either the trust store
+ // path or the encoded trust store must be provided
+ if (trustStorePath == null && base64EncodedTrustStore == null)
+ {
+ throw new IllegalArgumentException(
+ String.format("One of the '%s' or '%s' options must be provided when the '%s' option is provided",
+ TRUSTSTORE_PATH, TRUSTSTORE_BASE64_ENCODED, TRUSTSTORE_PASSWORD));
+ }
+ }
+ else if (trustStorePath != null || base64EncodedTrustStore != null)
+ {
+ throw new IllegalArgumentException(
+ String.format("The '%s' option must be provided when either the '%s' or '%s' options are provided",
+ TRUSTSTORE_PASSWORD, TRUSTSTORE_PATH, TRUSTSTORE_BASE64_ENCODED));
+ }
+ return self();
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/KryoRegister.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/KryoRegister.java
new file mode 100644
index 0000000..d00b705
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/KryoRegister.java
@@ -0,0 +1,108 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.Serializer;
+import org.apache.cassandra.bridge.BigNumberConfigImpl;
+import org.apache.cassandra.bridge.CassandraBridgeFactory;
+import org.apache.cassandra.bridge.CassandraVersion;
+import org.apache.cassandra.clients.SidecarInstanceImpl;
+import org.apache.cassandra.clients.SslConfig;
+import org.apache.cassandra.spark.data.CassandraDataLayer;
+import org.apache.cassandra.spark.data.LocalDataLayer;
+import org.apache.cassandra.spark.data.ReplicationFactor;
+import org.apache.cassandra.spark.data.partitioner.CassandraInstance;
+import org.apache.cassandra.spark.data.partitioner.CassandraRing;
+import org.apache.cassandra.spark.data.partitioner.TokenPartitioner;
+import org.apache.spark.SparkConf;
+import org.apache.spark.serializer.KryoRegistrator;
+import org.jetbrains.annotations.NotNull;
+
+/**
+ * Helper class to register classes for Kryo serialization
+ */
+public class KryoRegister implements KryoRegistrator
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(KryoRegister.class);
+ private static final String SPARK_SERIALIZER = "spark.serializer";
+ private static final String SPARK_REGISTRATORS = "spark.kryo.registrator";
+ private static final Map<Class<?>, Serializer<?>> KRYO_SERIALIZERS = Collections.synchronizedMap(new LinkedHashMap<>());
+
+ static
+ {
+ // Cassandra-version-agnostic Kryo serializers
+ KRYO_SERIALIZERS.put(LocalDataLayer.class, new LocalDataLayer.Serializer());
+ KRYO_SERIALIZERS.put(CassandraInstance.class, new CassandraInstance.Serializer());
+ KRYO_SERIALIZERS.put(ReplicationFactor.class, new ReplicationFactor.Serializer());
+ KRYO_SERIALIZERS.put(CassandraRing.class, new CassandraRing.Serializer());
+ KRYO_SERIALIZERS.put(TokenPartitioner.class, new TokenPartitioner.Serializer());
+ KRYO_SERIALIZERS.put(CassandraDataLayer.class, new CassandraDataLayer.Serializer());
+ KRYO_SERIALIZERS.put(BigNumberConfigImpl.class, new BigNumberConfigImpl.Serializer());
+ KRYO_SERIALIZERS.put(SidecarInstanceImpl.class, new SidecarInstanceImpl.Serializer());
+ KRYO_SERIALIZERS.put(SslConfig.class, new SslConfig.Serializer());
+ }
+
+ public static <T> void addSerializer(@NotNull Class<T> type, @NotNull Serializer<T> serializer)
+ {
+ LOGGER.info("Registering custom Kryo serializer type={}", type.getName());
+ KRYO_SERIALIZERS.put(type, serializer);
+ }
+
+ @Override
+ public void registerClasses(@NotNull Kryo kryo)
+ {
+ LOGGER.info("Initializing KryoRegister");
+
+ // TODO: Implicitly defaulting to Cassandra version 4.0 is a part of a previously published API.
+ // We might want to persist the version of Cassandra into the Spark configuration instead.
+ CassandraBridgeFactory.get(CassandraVersion.FOURZERO).kryoRegister(kryo);
+
+ KRYO_SERIALIZERS.forEach(kryo::register);
+ }
+
+ public static void setup(@NotNull SparkConf configuration)
+ {
+ // Use KryoSerializer
+ LOGGER.info("Setting up Kryo");
+ configuration.set(SPARK_SERIALIZER, "org.apache.spark.serializer.KryoSerializer");
+
+ // Add KryoRegister to SparkConf serialization if not already there
+ Set<String> registratorsSet = Arrays.stream(configuration.get(SPARK_REGISTRATORS, "").split(","))
+ .filter(string -> string != null && !string.isEmpty())
+ .collect(Collectors.toSet());
+ registratorsSet.add(KryoRegister.class.getName());
+ String registratorsString = String.join(",", registratorsSet);
+ LOGGER.info("Setting kryo registrators: " + registratorsString);
+ configuration.set(SPARK_REGISTRATORS, registratorsString);
+
+ configuration.registerKryoClasses(new Class<?>[]{KryoRegister.class});
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkSparkConf.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkSparkConf.java
new file mode 100644
index 0000000..292282f
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkSparkConf.java
@@ -0,0 +1,464 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.ByteArrayInputStream;
+import java.io.FileInputStream;
+import java.io.FileNotFoundException;
+import java.io.InputStream;
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.bridge.CassandraBridgeFactory;
+import org.apache.cassandra.clients.SidecarInstanceImpl;
+import org.apache.cassandra.sidecar.client.SidecarInstance;
+import org.apache.cassandra.spark.bulkwriter.token.ConsistencyLevel;
+import org.apache.cassandra.spark.bulkwriter.util.SbwKryoRegistrator;
+import org.apache.cassandra.spark.utils.BuildInfo;
+import org.apache.cassandra.spark.utils.MapUtils;
+import org.apache.spark.SparkConf;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+@SuppressWarnings("WeakerAccess")
+public class BulkSparkConf implements Serializable
+{
+ private static final long serialVersionUID = -5060973521517656241L;
+ private static final Logger LOGGER = LoggerFactory.getLogger(BulkSparkConf.class);
+
+ public static final String JDK11_OPTIONS = " -Djdk.attach.allowAttachSelf=true"
+ + " --add-exports java.base/jdk.internal.misc=ALL-UNNAMED"
+ + " --add-exports java.base/jdk.internal.ref=ALL-UNNAMED"
+ + " --add-exports java.base/sun.nio.ch=ALL-UNNAMED"
+ + " --add-exports java.management.rmi/com.sun.jmx.remote.internal.rmi=ALL-UNNAMED"
+ + " --add-exports java.rmi/sun.rmi.registry=ALL-UNNAMED"
+ + " --add-exports java.rmi/sun.rmi.server=ALL-UNNAMED"
+ + " --add-exports java.sql/java.sql=ALL-UNNAMED"
+ + " --add-opens java.base/java.lang.module=ALL-UNNAMED"
+ + " --add-opens java.base/jdk.internal.loader=ALL-UNNAMED"
+ + " --add-opens java.base/jdk.internal.ref=ALL-UNNAMED"
+ + " --add-opens java.base/jdk.internal.reflect=ALL-UNNAMED"
+ + " --add-opens java.base/jdk.internal.math=ALL-UNNAMED"
+ + " --add-opens java.base/jdk.internal.module=ALL-UNNAMED"
+ + " --add-opens java.base/jdk.internal.util.jar=ALL-UNNAMED"
+ + " --add-opens jdk.management/com.sun.management.internal=ALL-UNNAMED";
+
+ public static final int DEFAULT_NUM_SPLITS = -1;
+ public static final int DEFAULT_HTTP_CONNECTION_TIMEOUT = 100_000;
+ public static final int DEFAULT_HTTP_RESPONSE_TIMEOUT = 100_000;
+ public static final int DEFAULT_HTTP_MAX_CONNECTIONS = 200;
+ public static final int DEFAULT_SIDECAR_PORT = 9043;
+ public static final int DEFAULT_SIDECAR_REQUEST_RETRIES = 10;
+ public static final long DEFAULT_SIDECAR_REQUEST_RETRY_DELAY_SECONDS = 1L;
+ public static final long DEFAULT_SIDECAR_REQUEST_MAX_RETRY_DELAY_SECONDS = 60L;
+ public static final int DEFAULT_COMMIT_BATCH_SIZE = 10_000;
+ public static final int DEFAULT_RING_RETRY_COUNT = 3;
+ public static final RowBufferMode DEFAULT_ROW_BUFFER_MODE = RowBufferMode.UNBUFFERRED;
+ public static final int DEFAULT_BATCH_SIZE_IN_ROWS = 1_000_000;
+
+ // NOTE: All Cassandra Analytics setting names must start with "spark" in order to not be ignored by Spark,
+ // and must not start with "spark.cassandra" so as to not conflict with Spark Cassandra Connector
+ // which will throw a configuration exception for each setting with that prefix it does not recognize
+ public static final String SETTING_PREFIX = "spark.cassandra_analytics.";
+
+ public static final String HTTP_MAX_CONNECTIONS = SETTING_PREFIX + "request.max_connections";
+ public static final String HTTP_RESPONSE_TIMEOUT = SETTING_PREFIX + "request.response_timeout";
+ public static final String HTTP_CONNECTION_TIMEOUT = SETTING_PREFIX + "request.connection_timeout";
+ public static final String SIDECAR_PORT = SETTING_PREFIX + "ports.sidecar";
+ public static final String SIDECAR_REQUEST_RETRIES = SETTING_PREFIX + "sidecar.request.retries";
+ public static final String SIDECAR_REQUEST_RETRY_DELAY_SECONDS = SETTING_PREFIX + "sidecar.request.retries.delay.seconds";
+ public static final String SIDECAR_REQUEST_MAX_RETRY_DELAY_SECONDS = SETTING_PREFIX + "sidecar.request.retries.max.delay.seconds";
+ public static final String SKIP_CLEAN = SETTING_PREFIX + "job.skip_clean";
+ public static final String USE_OPENSSL = SETTING_PREFIX + "use_openssl";
+ public static final String RING_RETRY_COUNT = SETTING_PREFIX + "ring_retry_count";
+
+ public final Set<? extends SidecarInstance> sidecarInstances;
+ public final String keyspace;
+ public final String table;
+ public final ConsistencyLevel.CL consistencyLevel;
+ public final String localDC;
+ public final Integer numberSplits;
+ public final RowBufferMode rowBufferMode;
+ public final Integer sstableDataSizeInMB;
+ public final Integer sstableBatchSize;
+ public final int commitBatchSize;
+ public final boolean skipExtendedVerify;
+ public final WriteMode writeMode;
+ protected final String keystorePassword;
+ protected final String keystorePath;
+ protected final String keystoreBase64Encoded;
+ protected final String keystoreType;
+ protected final String truststorePassword;
+ protected final String truststorePath;
+ protected final String truststoreBase64Encoded;
+ protected final String truststoreType;
+ protected final SparkConf conf;
+ public final boolean validateSSTables;
+ public final int commitThreadsPerInstance;
+ protected final int sidecarPort;
+ protected boolean useOpenSsl;
+ protected int ringRetryCount;
+
+ public BulkSparkConf(SparkConf conf, Map<String, String> options)
+ {
+ this.conf = conf;
+ this.sidecarPort = MapUtils.getInt(options, WriterOptions.SIDECAR_PORT.name(), getInt(SIDECAR_PORT, DEFAULT_SIDECAR_PORT), "sidecar port");
+ this.sidecarInstances = buildSidecarInstances(options, sidecarPort);
+ this.keyspace = MapUtils.getOrThrow(options, WriterOptions.KEYSPACE.name());
+ this.table = MapUtils.getOrThrow(options, WriterOptions.TABLE.name());
+ this.validateSSTables = MapUtils.getBoolean(options, WriterOptions.VALIDATE_SSTABLES.name(), true, "validate SSTables");
+ this.skipExtendedVerify = MapUtils.getBoolean(options, WriterOptions.SKIP_EXTENDED_VERIFY.name(), true,
+ "skip extended verification of SSTables by Cassandra");
+ this.consistencyLevel = ConsistencyLevel.CL.valueOf(MapUtils.getOrDefault(options, WriterOptions.BULK_WRITER_CL.name(), "EACH_QUORUM"));
+ this.localDC = MapUtils.getOrDefault(options, WriterOptions.LOCAL_DC.name(), null);
+ this.numberSplits = MapUtils.getInt(options, WriterOptions.NUMBER_SPLITS.name(), DEFAULT_NUM_SPLITS, "number of splits");
+ this.rowBufferMode = MapUtils.getEnumOption(options, WriterOptions.ROW_BUFFER_MODE.name(), DEFAULT_ROW_BUFFER_MODE, "row bufferring mode");
+ this.sstableDataSizeInMB = MapUtils.getInt(options, WriterOptions.SSTABLE_DATA_SIZE_IN_MB.name(), 160, "sstable data size in MB");
+ this.sstableBatchSize = MapUtils.getInt(options, WriterOptions.BATCH_SIZE.name(), 1_000_000, "sstable batch size");
+ this.commitBatchSize = MapUtils.getInt(options, WriterOptions.COMMIT_BATCH_SIZE.name(), DEFAULT_COMMIT_BATCH_SIZE, "commit batch size");
+ this.commitThreadsPerInstance = MapUtils.getInt(options, WriterOptions.COMMIT_THREADS_PER_INSTANCE.name(), 2, "commit threads per instance");
+ this.keystorePassword = MapUtils.getOrDefault(options, WriterOptions.KEYSTORE_PASSWORD.name(), null);
+ this.keystorePath = MapUtils.getOrDefault(options, WriterOptions.KEYSTORE_PATH.name(), null);
+ this.keystoreBase64Encoded = MapUtils.getOrDefault(options, WriterOptions.KEYSTORE_BASE64_ENCODED.name(), null);
+ this.keystoreType = MapUtils.getOrDefault(options, WriterOptions.KEYSTORE_TYPE.name(), "PKCS12");
+ this.truststorePassword = MapUtils.getOrDefault(options, WriterOptions.TRUSTSTORE_PASSWORD.name(), null);
+ this.truststorePath = MapUtils.getOrDefault(options, WriterOptions.TRUSTSTORE_PATH.name(), null);
+ this.truststoreBase64Encoded = MapUtils.getOrDefault(options, WriterOptions.TRUSTSTORE_BASE64_ENCODED.name(), null);
+ this.truststoreType = MapUtils.getOrDefault(options, WriterOptions.TRUSTSTORE_TYPE.name(), null);
+ this.writeMode = MapUtils.getEnumOption(options, WriterOptions.WRITE_MODE.name(), WriteMode.INSERT, "write mode");
+ // For backwards-compatibility with port settings, use writer option if available,
+ // else fall back to props, and then default if neither specified
+ this.useOpenSsl = getBoolean(USE_OPENSSL, true);
+ this.ringRetryCount = getInt(RING_RETRY_COUNT, DEFAULT_RING_RETRY_COUNT);
+ validateEnvironment();
+ }
+
+ protected Set<? extends SidecarInstance> buildSidecarInstances(Map<String, String> options, int sidecarPort)
+ {
+ String sidecarInstances = MapUtils.getOrThrow(options, WriterOptions.SIDECAR_INSTANCES.name(), "sidecar_instances");
+ return Arrays.stream(sidecarInstances.split(","))
+ .map(hostname -> new SidecarInstanceImpl(hostname, sidecarPort))
+ .collect(Collectors.toSet());
+ }
+
+ protected void validateEnvironment() throws RuntimeException
+ {
+ Preconditions.checkNotNull(keyspace);
+ Preconditions.checkNotNull(table);
+ Preconditions.checkArgument(getHttpResponseTimeoutMs() > 0, HTTP_RESPONSE_TIMEOUT + " must be > 0");
+ validateSslConfiguration();
+ validateTableWriterSettings();
+ CassandraBridgeFactory.validateBridges();
+ }
+
+ protected void validateTableWriterSettings()
+ {
+ boolean batchSizeIsZero = sstableBatchSize == 0;
+ if (rowBufferMode == RowBufferMode.BUFFERRED
+ && (!batchSizeIsZero && sstableBatchSize != DEFAULT_BATCH_SIZE_IN_ROWS))
+ {
+ LOGGER.warn("BATCH_SIZE is set to a non-zero, non-default value ({}) but ROW_BUFFER_MODE is set to BUFFERRED."
+ + " Ignoring BATCH_SIZE.", sstableBatchSize);
+ }
+ Preconditions.checkArgument(rowBufferMode == RowBufferMode.UNBUFFERRED && !batchSizeIsZero,
+ "If writing in sorted order (ROW_BUFFER_MODE is UNBUFFERRED) then BATCH_SIZE "
+ + "should be non zero, but it was set to 0 in writer options");
+ }
+
+ /*
+ * This method Will throw if the SSL configuration is incorrect (PATH provided w/o password, for example)
+ */
+ protected void validateSslConfiguration()
+ {
+ if (getKeyStorePassword() != null)
+ {
+ // If the mTLS keystore password is provided, we validate that the path or a base64 keystore is provided
+ if (getKeyStorePath() == null && getKeystoreBase64Encoded() == null)
+ {
+ throw new NullPointerException("Keystore password was set. "
+ + "But both keystore path and base64 encoded string are not set. "
+ + "Please either set option " + WriterOptions.KEYSTORE_PATH
+ + " or option " + WriterOptions.KEYSTORE_BASE64_ENCODED);
+ }
+ }
+
+ // Check to make sure if either trust store password or trust store path are specified, both are specified
+ if (getConfiguredTrustStorePassword() != null)
+ {
+ Preconditions.checkArgument(getTruststoreBase64Encoded() != null || getTrustStorePath() != null,
+ "Trust Store password was provided, but both truststore path and base64 encoded string are missing. "
+ + "Please provide either option " + WriterOptions.TRUSTSTORE_PATH
+ + " or option " + WriterOptions.TRUSTSTORE_BASE64_ENCODED);
+ }
+ else
+ {
+ if (getTrustStorePath() != null || getTruststoreBase64Encoded() != null)
+ {
+ Preconditions.checkNotNull(getTruststoreBase64Encoded(),
+ "Trust Store Path was provided, but password is missing. "
+ + "Please provide option " + WriterOptions.TRUSTSTORE_PASSWORD);
+ Preconditions.checkNotNull(getTrustStorePath(),
+ "Trust Store Base64 encoded was provided, but password is missing."
+ + "Please provide option " + WriterOptions.TRUSTSTORE_PASSWORD);
+ }
+ }
+ }
+
+ public int getSidecarPort()
+ {
+ return sidecarPort;
+ }
+
+ protected String getTrustStorePath()
+ {
+ return truststorePath;
+ }
+
+ protected String getTruststoreBase64Encoded()
+ {
+ return truststoreBase64Encoded;
+ }
+
+ public String getTrustStoreTypeOrDefault()
+ {
+ return truststoreType != null ? truststoreType : "PKCS12";
+ }
+
+ protected String getKeyStorePath()
+ {
+ return keystorePath;
+ }
+
+ protected String getKeystoreBase64Encoded()
+ {
+ return keystoreBase64Encoded;
+ }
+
+ public InputStream getKeyStore()
+ {
+ return getKeyStorePath() != null
+ ? getKeyStoreFromPath(getKeyStorePath())
+ : getKeyStoreFromBase64EncodedString(getKeystoreBase64Encoded());
+ }
+
+ @Nullable
+ public InputStream getTrustStore()
+ {
+ return getTrustStorePath() != null
+ ? getKeyStoreFromPath(getTrustStorePath())
+ : getKeyStoreFromBase64EncodedString(getTruststoreBase64Encoded());
+ }
+
+ protected InputStream getKeyStoreFromPath(String keyStorePath)
+ {
+ if (keyStorePath != null)
+ {
+ try
+ {
+ return new FileInputStream(keyStorePath);
+ }
+ catch (FileNotFoundException exception)
+ {
+ throw new RuntimeException("Could not load keystore at path '" + keyStorePath + "'", exception);
+ }
+ }
+ return null;
+ }
+
+ protected InputStream getKeyStoreFromBase64EncodedString(String keyStoreBase64Encoded)
+ {
+ if (keyStoreBase64Encoded != null)
+ {
+ return new ByteArrayInputStream(Base64.getDecoder().decode(keyStoreBase64Encoded));
+ }
+ return null;
+ }
+
+ public String getTrustStorePasswordOrDefault()
+ {
+ return truststorePassword != null ? truststorePassword : "password";
+ }
+
+ public String getKeyStoreTypeOrDefault()
+ {
+ return keystoreType != null ? keystoreType : "PKCS12";
+ }
+
+ public String getKeyStorePassword()
+ {
+ return keystorePassword;
+ }
+
+ public String getConfiguredTrustStorePassword()
+ {
+ return truststorePassword;
+ }
+
+ public String getConfiguredKeyStorePassword()
+ {
+ return keystorePassword;
+ }
+
+ public int getSidecarRequestRetries()
+ {
+ return getInt(SIDECAR_REQUEST_RETRIES, DEFAULT_SIDECAR_REQUEST_RETRIES);
+ }
+
+ public long getSidecarRequestRetryDelayInSeconds()
+ {
+ return getLong(SIDECAR_REQUEST_RETRY_DELAY_SECONDS, DEFAULT_SIDECAR_REQUEST_RETRY_DELAY_SECONDS);
+ }
+
+ public long getSidecarRequestMaxRetryDelayInSeconds()
+ {
+ return getLong(SIDECAR_REQUEST_MAX_RETRY_DELAY_SECONDS, DEFAULT_SIDECAR_REQUEST_MAX_RETRY_DELAY_SECONDS);
+ }
+
+ public int getHttpConnectionTimeoutMs()
+ {
+ return getInt(HTTP_CONNECTION_TIMEOUT, DEFAULT_HTTP_CONNECTION_TIMEOUT);
+ }
+
+ public int getHttpResponseTimeoutMs()
+ {
+ return getInt(HTTP_RESPONSE_TIMEOUT, DEFAULT_HTTP_RESPONSE_TIMEOUT);
+ }
+
+ public int getMaxHttpConnections()
+ {
+ return getInt(HTTP_MAX_CONNECTIONS, DEFAULT_HTTP_MAX_CONNECTIONS);
+ }
+
+ public boolean getSkipClean()
+ {
+ return getBoolean(SKIP_CLEAN, false);
+ }
+
+ public Integer getCores()
+ {
+ int coresPerExecutor = conf.getInt("spark.executor.cores", 1);
+ int numExecutors = conf.getInt("spark.dynamicAllocation.maxExecutors", conf.getInt("spark.executor.instances", 1));
+ return coresPerExecutor * numExecutors;
+ }
+
+ protected int getInt(String settingName, int defaultValue)
+ {
+ String finalSetting = getSettingNameOrDeprecatedName(settingName);
+ return conf.getInt(finalSetting, defaultValue);
+ }
+
+ protected long getLong(String settingName, long defaultValue)
+ {
+ String finalSetting = getSettingNameOrDeprecatedName(settingName);
+ return conf.getLong(finalSetting, defaultValue);
+ }
+
+ protected boolean getBoolean(String settingName, boolean defaultValue)
+ {
+ String finalSetting = getSettingNameOrDeprecatedName(settingName);
+ return conf.getBoolean(finalSetting, defaultValue);
+ }
+
+ protected String getSettingNameOrDeprecatedName(String settingName)
+ {
+ if (!conf.contains(settingName))
+ {
+ String settingSuffix = settingName.startsWith(SETTING_PREFIX)
+ ? settingName.substring(SETTING_PREFIX.length())
+ : settingName;
+ for (String settingPrefix : getDeprecatedSettingPrefixes())
+ {
+ String deprecatedSetting = settingPrefix + settingSuffix;
+ if (conf.contains(deprecatedSetting))
+ {
+ LOGGER.warn("Found deprecated setting '{}'. Please use {} in the future.",
+ deprecatedSetting, settingName);
+ return deprecatedSetting;
+ }
+ }
+ }
+ return settingName;
+ }
+
+ @NotNull
+ protected List<String> getDeprecatedSettingPrefixes()
+ {
+ return Collections.emptyList();
+ }
+
+ /**
+ * The SBW utilizes Cassandra libraries to generate SSTables. Under JDK11, this library needs additional JVM options
+ * to be set for the executors for some backward-compatibility reasons. This method will add the appropriate
+ * JVM options. Additionally, we set up the SBW KryoRegistrator here rather than requiring the end-user
+ * to call 2 static methods to set up the appropriate settings for the job.
+ *
+ * @param conf the Spark Configuration to set up
+ * @param addKryoRegistrator passs true if your application hasn't separately added
+ * the bulk-specific Kryo registrator, false if you have set it up
+ * separately (see the usage docs for more details)
+ */
+ public static void setupSparkConf(SparkConf conf, boolean addKryoRegistrator)
+ {
+ String previousOptions = conf.get("spark.executor.extraJavaOptions", "");
+ if (BuildInfo.isAtLeastJava11(BuildInfo.javaSpecificationVersion()))
+ {
+ conf.set("spark.executor.extraJavaOptions", previousOptions + JDK11_OPTIONS);
+ }
+
+ if (addKryoRegistrator)
+ {
+ // Use `SbwKryoRegistrator.setupKryoRegistrator(conf);` to add the Spark Bulk Writer's custom
+ // KryoRegistrator. This needs to happen before the SparkSession is built in order for it to work properly.
+ // This utility method will add the necessary configuration, but it may be overwritten later if your own
+ // code resets `spark.kryo.registrator` after this is called.
+ SbwKryoRegistrator.setupKryoRegistrator(conf);
+ }
+ }
+
+ protected SparkConf getConf()
+ {
+ return conf;
+ }
+
+ public boolean getUseOpenSsl()
+ {
+ return useOpenSsl;
+ }
+
+ public int getRingRetryCount()
+ {
+ return ringRetryCount;
+ }
+
+ public boolean hasKeystoreAndKeystorePassword()
+ {
+ return keystorePassword != null && (keystorePath != null || keystoreBase64Encoded != null);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java
new file mode 100644
index 0000000..1531738
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriteValidator.java
@@ -0,0 +1,169 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.math.BigInteger;
+import java.util.AbstractMap;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Range;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.bulkwriter.token.ReplicaAwareFailureHandler;
+
+public class BulkWriteValidator implements AutoCloseable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(BulkWriteValidator.class);
+
+ private final ReplicaAwareFailureHandler<RingInstance> failureHandler;
+ private final CassandraRingMonitor monitor;
+ private final JobInfo job;
+ private String phase = "Initializing";
+
+ private final ClusterInfo cluster;
+
+ public BulkWriteValidator(BulkWriterContext bulkWriterContext,
+ Consumer<CancelJobEvent> cancelJobFunc) throws Exception
+ {
+ cluster = bulkWriterContext.cluster();
+ job = bulkWriterContext.job();
+ failureHandler = new ReplicaAwareFailureHandler<>(cluster.getRing(true));
+ monitor = new CassandraRingMonitor(cluster, cancelJobFunc, 1000, TimeUnit.MILLISECONDS);
+ }
+
+ public void setPhase(String phase)
+ {
+ this.phase = phase;
+ }
+
+ public String getPhase()
+ {
+ return phase;
+ }
+
+ public void validateInitialEnvironment()
+ {
+ validateCLOrFail();
+ }
+
+ public void close()
+ {
+ monitor.stop();
+ }
+
+ private void validateCLOrFail()
+ {
+ updateInstanceAvailability();
+ validateClOrFail(failureHandler, LOGGER, phase, job);
+ }
+
+ public static void validateClOrFail(ReplicaAwareFailureHandler<RingInstance> failureHandler,
+ Logger logger,
+ String phase,
+ JobInfo job)
+ {
+ Collection<AbstractMap.SimpleEntry<Range<BigInteger>, Multimap<RingInstance, String>>> failedRanges =
+ failureHandler.getFailedEntries(job.getConsistencyLevel(), job.getLocalDC());
+ if (failedRanges.isEmpty())
+ {
+ logger.info("Succeeded {} with {}", phase, job.getConsistencyLevel());
+ }
+ else
+ {
+ String message = String.format("Failed to load %s ranges with %s for job %s in phase %s",
+ failedRanges.size(), job.getConsistencyLevel(), job.getId(), phase);
+ logger.error(message);
+ failedRanges.forEach(failedRange -> failedRange.getValue().keySet().forEach(instance ->
+ logger.error("Failed {} for {} on {}", phase, failedRange.getKey(), instance.toString())));
+ throw new RuntimeException(message);
+ }
+ }
+
+ public static void updateFailureHandler(CommitResult commitResult,
+ String phase,
+ ReplicaAwareFailureHandler<RingInstance> failureHandler)
+ {
+ LOGGER.debug("Commit Result: {}", commitResult);
+ commitResult.failures.forEach((uuid, err) -> {
+ LOGGER.warn("[{}]: {} failed on {} with message {}",
+ uuid, phase, commitResult.instance.getNodeName(), err.errMsg);
+ failureHandler.addFailure(err.tokenRange, commitResult.instance, err.errMsg);
+ });
+ }
+
+ private void updateInstanceAvailability()
+ {
+ cluster.refreshClusterInfo();
+ Map<RingInstance, InstanceAvailability> availability = cluster.getInstanceAvailability();
+ availability.forEach(this::checkInstance);
+ }
+
+ private void checkInstance(RingInstance instance, InstanceAvailability availability)
+ {
+ throwIfInvalidState(instance, availability);
+ updateFailureHandler(instance, availability);
+ }
+
+ private void throwIfInvalidState(RingInstance instance, InstanceAvailability availability)
+ {
+ if (availability == InstanceAvailability.INVALID_STATE)
+ {
+ // If we find any nodes in a totally invalid state, just throw as we can't continue
+ String message = String.format("Instance (%s) is in an invalid state (%s) during import. "
+ + "Please rerun import once topology changes are complete.",
+ instance.getNodeName(), cluster.getInstanceState(instance));
+ throw new RuntimeException(message);
+ }
+ }
+
+ private void updateFailureHandler(RingInstance instX, InstanceAvailability availability)
+ {
+ if (availability != InstanceAvailability.AVAILABLE)
+ {
+ addFailedInstance(instX, availability.getMessage());
+ }
+ }
+
+ private void addFailedInstance(RingInstance instance, String reason)
+ {
+ Collection<Range<BigInteger>> failedRanges = cluster.getRing(true).getTokenRanges(instance);
+ failedRanges.forEach(failedRange -> {
+ String nodeDisplayName = instance.getNodeName();
+ String message = String.format("%s %s", nodeDisplayName, reason);
+ LOGGER.warn("{} failed in phase {} on {} because {}", failedRange, phase, nodeDisplayName, message);
+ failureHandler.addFailure(failedRange, instance, message);
+ });
+ }
+
+ public void failIfRingChanged()
+ {
+ if (monitor.getRingChanged())
+ {
+ throw new RuntimeException(String.format("Ring changed during %s stage of import. "
+ + "Please rerun import once topology changes are complete.",
+ phase));
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriterContext.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriterContext.java
new file mode 100644
index 0000000..e3b2421
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/BulkWriterContext.java
@@ -0,0 +1,37 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+
+public interface BulkWriterContext extends Serializable
+{
+ ClusterInfo cluster();
+
+ JobInfo job();
+
+ SchemaInfo schema();
+
+ DataTransferApi transfer();
+
+ // NOTE: This interface intentionally does *not* implement AutoClosable as Spark can close Broadcast variables
+ // that implement AutoClosable while they are still in use, causing the underlying object to become unusable
+ void shutdown();
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CancelJobEvent.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CancelJobEvent.java
new file mode 100644
index 0000000..ab3119b
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CancelJobEvent.java
@@ -0,0 +1,38 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public class CancelJobEvent
+{
+ public final Throwable exception;
+ public final String reason;
+
+ public CancelJobEvent(String reason)
+ {
+ this.reason = reason;
+ this.exception = null;
+ }
+
+ public CancelJobEvent(String reason, Throwable throwable)
+ {
+ this.reason = reason;
+ this.exception = throwable;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkSourceRelation.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkSourceRelation.java
new file mode 100644
index 0000000..15d4371
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkSourceRelation.java
@@ -0,0 +1,185 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.spark.api.java.JavaPairRDD;
+import org.apache.spark.api.java.JavaSparkContext;
+import org.apache.spark.broadcast.Broadcast;
+import org.apache.spark.sql.Dataset;
+import org.apache.spark.sql.Row;
+import org.apache.spark.sql.SQLContext;
+import org.apache.spark.sql.sources.BaseRelation;
+import org.apache.spark.sql.sources.InsertableRelation;
+import org.apache.spark.sql.types.StructType;
+import org.jetbrains.annotations.NotNull;
+import scala.Tuple2;
+import scala.collection.JavaConverters;
+import scala.util.control.NonFatal$;
+
+public class CassandraBulkSourceRelation extends BaseRelation implements InsertableRelation
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(CassandraBulkSourceRelation.class);
+
+ private final BulkWriterContext writerContext;
+ private final SQLContext sqlContext;
+ private final JavaSparkContext sparkContext;
+ private final Broadcast<BulkWriterContext> broadcastContext;
+ private final BulkWriteValidator writeValidator;
+
+ @SuppressWarnings("RedundantTypeArguments")
+ public CassandraBulkSourceRelation(BulkWriterContext writerContext, SQLContext sqlContext) throws Exception
+ {
+ this.writerContext = writerContext;
+ this.sqlContext = sqlContext;
+ this.sparkContext = JavaSparkContext.fromSparkContext(sqlContext.sparkContext());
+ this.broadcastContext = sparkContext.<BulkWriterContext>broadcast(writerContext);
+ this.writeValidator = new BulkWriteValidator(writerContext, this::cancelJob);
+ }
+
+ @Override
+ @NotNull
+ public SQLContext sqlContext()
+ {
+ return sqlContext;
+ }
+
+ /**
+ * @return An empty {@link StructType}, as this is a writer only, so schema is not applicable
+ */
+ @Override
+ @NotNull
+ public StructType schema()
+ {
+ LOGGER.warn("This instance is used as writer, a schema is not supported");
+ return new StructType();
+ }
+
+ /**
+ * @return {@code 0} size as not applicable use by the planner in the writer-only use case
+ */
+ @Override
+ public long sizeInBytes()
+ {
+ LOGGER.warn("This instance is used as writer, sizeInBytes is not supported");
+ return 0L;
+ }
+
+ @Override
+ public void insert(@NotNull Dataset<Row> data, boolean overwrite)
+ {
+ if (overwrite)
+ {
+ throw new LoadNotSupportedException("Overwriting existing data needs TRUNCATE on Cassandra, which is not supported");
+ }
+ writerContext.cluster().checkBulkWriterIsEnabledOrThrow();
+ Tokenizer tokenizer = new Tokenizer(writerContext);
+ TableSchema tableSchema = writerContext.schema().getTableSchema();
+ JavaPairRDD<DecoratedKey, Object[]> sortedRDD = data.toJavaRDD()
+ .map(Row::toSeq)
+ .map(seq -> JavaConverters.seqAsJavaListConverter(seq).asJava().toArray())
+ .map(tableSchema::normalize)
+ .keyBy(tokenizer::getDecoratedKey)
+ .repartitionAndSortWithinPartitions(broadcastContext.getValue().job().getTokenPartitioner());
+ persist(sortedRDD);
+ }
+
+ private void cancelJob(@NotNull CancelJobEvent cancelJobEvent)
+ {
+ if (cancelJobEvent.exception != null)
+ {
+ LOGGER.error("An unrecoverable error occurred during {} stage of import while validating the current cluster state; cancelling job",
+ writeValidator.getPhase(), cancelJobEvent.exception);
+ }
+ else
+ {
+ LOGGER.error("Job was canceled due to '{}' during {} stage of import; please rerun import once topology changes are complete",
+ cancelJobEvent.reason, writeValidator.getPhase());
+ }
+ sparkContext.cancelJobGroup(writerContext.job().getId().toString());
+ }
+
+ @SuppressWarnings("RedundantCast")
+ private void persist(@NotNull JavaPairRDD<DecoratedKey, Object[]> sortedRDD)
+ {
+ writeValidator.setPhase("Environment Validation");
+ writeValidator.validateInitialEnvironment();
+ writeValidator.setPhase("UploadAndCommit");
+
+ try
+ {
+ sortedRDD.foreachPartition(new WriteIterator(broadcastContext)::call);
+ writeValidator.failIfRingChanged();
+ }
+ catch (Throwable throwable)
+ {
+ LOGGER.error("Bulk Write Failed", throwable);
+ throw new RuntimeException("Bulk Write to Cassandra has failed", throwable);
+ }
+ finally
+ {
+ writeValidator.close(); // Uses the MgrClient, so needs to stop first
+ try
+ {
+ writerContext.shutdown();
+ sqlContext().sparkContext().clearJobGroup();
+ }
+ catch (Exception ignored)
+ {
+ // We've made our best effort to close the Bulk Writer context
+ }
+ try
+ {
+ broadcastContext.unpersist(false);
+ }
+ catch (Throwable throwable)
+ {
+ if (NonFatal$.MODULE$.apply(throwable))
+ {
+ LOGGER.error("Uncaught exception in thread {} attempting to unpersist broadcast variable",
+ Thread.currentThread().getName(), throwable);
+ }
+ else
+ {
+ throw throwable;
+ }
+ }
+ }
+ }
+
+ private static class WriteIterator implements Serializable
+ {
+ private final Broadcast<BulkWriterContext> broadcastContext;
+
+ WriteIterator(Broadcast<BulkWriterContext> broadcastContext)
+ {
+ this.broadcastContext = broadcastContext;
+ }
+
+ public void call(java.util.Iterator<Tuple2<DecoratedKey, Object[]>> iterator)
+ {
+ new RecordWriter(broadcastContext.getValue()).write(iterator);
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkWriterContext.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkWriterContext.java
new file mode 100644
index 0000000..4813b1c
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraBulkWriterContext.java
@@ -0,0 +1,194 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.util.Map;
+import java.util.Set;
+
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.esotericsoftware.kryo.Kryo;
+import com.esotericsoftware.kryo.KryoSerializable;
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.Output;
+import io.netty.channel.EventLoopGroup;
+import org.apache.cassandra.bridge.CassandraBridge;
+import org.apache.cassandra.bridge.CassandraBridgeFactory;
+import org.apache.cassandra.spark.bulkwriter.token.CassandraRing;
+import org.apache.cassandra.spark.data.CqlTable;
+import org.apache.cassandra.spark.data.ReplicationFactor;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+import org.apache.cassandra.spark.utils.CqlUtils;
+import org.apache.cassandra.spark.utils.ScalaFunctions;
+import org.apache.spark.SparkContext;
+import org.apache.spark.sql.types.StructType;
+import org.apache.spark.util.ShutdownHookManager;
+import org.jetbrains.annotations.NotNull;
+
+// CHECKSTYLE IGNORE: This class cannot be declared as final, because consumers should be able to extend it
+public class CassandraBulkWriterContext implements BulkWriterContext, KryoSerializable
+{
+ private static final long serialVersionUID = 8241993502687688783L;
+ private static final Logger LOGGER = LoggerFactory.getLogger(CassandraBulkWriterContext.class);
+
+ @NotNull
+ private final BulkSparkConf conf;
+ private final JobInfo jobInfo;
+ private final DataTransferApi dataTransferApi;
+ private final CassandraClusterInfo clusterInfo;
+ private final SchemaInfo schemaInfo;
+
+ static
+ {
+ configureRelocatedNetty();
+ }
+
+ private CassandraBulkWriterContext(@NotNull BulkSparkConf conf,
+ @NotNull StructType dfSchema,
+ SparkContext sparkContext)
+ {
+ this.conf = conf;
+ clusterInfo = new CassandraClusterInfo(conf);
+ CassandraRing<RingInstance> ring = clusterInfo.getRing(true);
+ jobInfo = new CassandraJobInfo(conf,
+ new TokenPartitioner(ring, conf.numberSplits, sparkContext.defaultParallelism(), conf.getCores()));
+ Preconditions.checkArgument(!conf.consistencyLevel.isLocal()
+ || (conf.localDC != null && ring.getReplicationFactor().getOptions().containsKey(conf.localDC)),
+ String.format("Keyspace %s is not replicated on datacenter %s",
+ conf.keyspace, conf.localDC));
+
+ dataTransferApi = new SidecarDataTransferApi(clusterInfo.getCassandraContext().getSidecarClient(), jobInfo, conf);
+
+ String keyspace = conf.keyspace;
+ String table = conf.table;
+
+ String keyspaceSchema = clusterInfo.getKeyspaceSchema(true);
+ CassandraBridge bridge = CassandraBridgeFactory.get(clusterInfo.getLowestCassandraVersion());
+ Partitioner partitioner = clusterInfo.getPartitioner();
+ String tableSchema = CqlUtils.extractTableSchema(keyspaceSchema, keyspace, table);
+ Set<String> udts = CqlUtils.extractUdts(keyspaceSchema, keyspace);
+ ReplicationFactor replicationFactor = CqlUtils.extractReplicationFactor(keyspaceSchema, keyspace);
+ int indexCount = CqlUtils.extractIndexCount(keyspaceSchema, keyspace, table);
+ CqlTable cqlTable = bridge.buildSchema(tableSchema, keyspace, replicationFactor, partitioner, udts, null, indexCount);
+
+ TableInfoProvider tableInfoProvider = new CqlTableInfoProvider(tableSchema, cqlTable);
+ schemaInfo = new CassandraSchemaInfo(new TableSchema(dfSchema, tableInfoProvider, conf.writeMode));
+ }
+
+ public static BulkWriterContext fromOptions(@NotNull SparkContext sparkContext,
+ @NotNull Map<String, String> strOptions,
+ @NotNull StructType dfSchema)
+ {
+ Preconditions.checkNotNull(dfSchema);
+
+ BulkSparkConf conf = new BulkSparkConf(sparkContext.getConf(), strOptions);
+ CassandraBulkWriterContext bulkWriterContext = new CassandraBulkWriterContext(conf, dfSchema, sparkContext);
+ ShutdownHookManager.addShutdownHook(org.apache.spark.util.ShutdownHookManager.TEMP_DIR_SHUTDOWN_PRIORITY(),
+ ScalaFunctions.wrapLambda(bulkWriterContext::shutdown));
+ bulkWriterContext.dialHome(sparkContext.version());
+ return bulkWriterContext;
+ }
+
+ @Override
+ public void shutdown()
+ {
+ LOGGER.info("Shutting down {}", this);
+ synchronized (this)
+ {
+ if (clusterInfo != null)
+ {
+ clusterInfo.close();
+ }
+ }
+ }
+
+ // If using the shaded JAR, Configure our shaded Netty so it can find the correct native libraries
+ private static void configureRelocatedNetty()
+ {
+ if (EventLoopGroup.class.getName().startsWith("analytics"))
+ {
+ System.setProperty("analytics.io.netty.packagePrefix", "analytics");
+ }
+ }
+
+ /**
+ * Use the implementation of the KryoSerializable interface as a detection device to make sure the Spark Bulk
+ * Writer's KryoRegistrator is properly in place.
+ *
+ * If this class is serialized by Kryo, it means we're <b>not</b> set up correctly, and therefore we log and fail.
+ * This failure will occur early in the job and be very clear, so users can quickly fix their code and get up and
+ * running again, rather than having a random NullPointerException further down the line.
+ */
+ public static final String KRYO_REGISTRATION_WARNING =
+ "Spark Bulk Writer Kryo Registrator (SbwKryoRegistrator) was not registered with Spark - "
+ + "please see the README.md file for more details on how to register the Spark Bulk Writer.";
+
+ @Override
+ public void write(Kryo kryo, Output output)
+ {
+ failIfKryoNotRegistered();
+ }
+
+ @Override
+ public void read(Kryo kryo, Input input)
+ {
+ failIfKryoNotRegistered();
+ }
+
+ private void failIfKryoNotRegistered()
+ {
+ LOGGER.error(KRYO_REGISTRATION_WARNING);
+ throw new RuntimeException(KRYO_REGISTRATION_WARNING);
+ }
+
+ protected void dialHome(String sparkVersion)
+ {
+ LOGGER.info("Dial home. clientConfig={}, sparkVersion={}", conf, sparkVersion);
+ }
+
+ @Override
+ @NotNull
+ public ClusterInfo cluster()
+ {
+ return clusterInfo;
+ }
+
+ @Override
+ @NotNull
+ public JobInfo job()
+ {
+ return jobInfo;
+ }
+
+ @Override
+ public SchemaInfo schema()
+ {
+ return schemaInfo;
+ }
+
+ @Override
+ @NotNull
+ public DataTransferApi transfer()
+ {
+ return dataTransferApi;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java
new file mode 100644
index 0000000..e84957f
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraClusterInfo.java
@@ -0,0 +1,499 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Closeable;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
+
+import com.google.common.annotations.VisibleForTesting;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.bridge.CassandraVersionFeatures;
+import org.apache.cassandra.clients.Sidecar;
+import org.apache.cassandra.sidecar.client.SidecarInstance;
+import org.apache.cassandra.sidecar.client.SidecarInstanceImpl;
+import org.apache.cassandra.sidecar.common.NodeSettings;
+import org.apache.cassandra.sidecar.common.data.GossipInfoResponse;
+import org.apache.cassandra.sidecar.common.data.RingEntry;
+import org.apache.cassandra.sidecar.common.data.RingResponse;
+import org.apache.cassandra.sidecar.common.data.SchemaResponse;
+import org.apache.cassandra.sidecar.common.data.TimeSkewResponse;
+import org.apache.cassandra.spark.bulkwriter.token.CassandraRing;
+import org.apache.cassandra.spark.common.client.InstanceState;
+import org.apache.cassandra.spark.common.client.InstanceStatus;
+import org.apache.cassandra.spark.data.ReplicationFactor;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+import org.apache.cassandra.spark.utils.CqlUtils;
+import org.jetbrains.annotations.NotNull;
+
+public class CassandraClusterInfo implements ClusterInfo, Closeable
+{
+ private static final long serialVersionUID = -6944818863462956767L;
+ private static final Logger LOGGER = LoggerFactory.getLogger(CassandraClusterInfo.class);
+
+ protected final BulkSparkConf conf;
+ protected String cassandraVersion;
+ protected Partitioner partitioner;
+ protected transient CassandraRing<RingInstance> ring;
+ protected transient String keyspaceSchema;
+ protected transient volatile RingResponse ringResponse;
+ protected transient GossipInfoResponse gossipInfo;
+ protected transient CassandraContext cassandraContext;
+ protected transient NodeSettings nodeSettings;
+
+ public CassandraClusterInfo(BulkSparkConf conf)
+ {
+ this.conf = conf;
+ this.cassandraContext = buildCassandraContext();
+ }
+
+ @Override
+ public void checkBulkWriterIsEnabledOrThrow()
+ {
+ // DO NOTHING
+ }
+
+ public String getVersion()
+ {
+ return CassandraClusterInfo.class.getPackage().getImplementationVersion();
+ }
+
+ @Override
+ public Map<RingInstance, InstanceAvailability> getInstanceAvailability()
+ {
+ RingResponse ringResponse = getRingResponse();
+ Map<RingInstance, InstanceAvailability> result = ringResponse
+ .stream()
+ .collect(Collectors.toMap(CassandraClusterInfo::getCasInstanceMethodsImpl,
+ this::determineInstanceAvailability));
+ if (LOGGER.isDebugEnabled())
+ {
+ result.forEach((instance, availability) ->
+ LOGGER.debug("Instance {} has availability {}", instance, availability));
+ }
+ return result;
+ }
+
+ @Override
+ public boolean instanceIsAvailable(RingInstance ringInstance)
+ {
+ return instanceIsUp(ringInstance.getRingInstance())
+ && instanceIsNormal(ringInstance.getRingInstance())
+ && !instanceIsBlocked(ringInstance);
+ }
+
+ @Override
+ public InstanceState getInstanceState(RingInstance ringInstance)
+ {
+ return InstanceState.valueOf(ringInstance.getRingInstance().state().toUpperCase());
+ }
+
+ public CassandraContext getCassandraContext()
+ {
+ CassandraContext currentCassandraContext = cassandraContext;
+ if (currentCassandraContext != null)
+ {
+ return currentCassandraContext;
+ }
+
+ synchronized (this)
+ {
+ if (cassandraContext == null)
+ {
+ cassandraContext = buildCassandraContext();
+ }
+ return cassandraContext;
+ }
+ }
+
+ /**
+ * Gets a Cassandra Context
+ *
+ * NOTE: The caller of this method is required to call `shutdown` on the returned CassandraContext instance
+ *
+ * @return an instance of CassandraContext based on the configuration settings
+ */
+ protected CassandraContext buildCassandraContext()
+ {
+ return buildCassandraContext(conf);
+ }
+
+ private static CassandraContext buildCassandraContext(BulkSparkConf conf)
+ {
+ return CassandraContext.create(conf);
+ }
+
+ @Override
+ public void close()
+ {
+ synchronized (this)
+ {
+ LOGGER.info("Closing {}", this);
+ getCassandraContext().close();
+ }
+ }
+
+ @Override
+ public Partitioner getPartitioner()
+ {
+ Partitioner currentPartitioner = partitioner;
+ if (currentPartitioner != null)
+ {
+ return currentPartitioner;
+ }
+
+ synchronized (this)
+ {
+ if (partitioner == null)
+ {
+ try
+ {
+ String partitionerString;
+ NodeSettings currentNodeSettings = nodeSettings;
+ if (currentNodeSettings != null)
+ {
+ partitionerString = currentNodeSettings.partitioner();
+ }
+ else
+ {
+ partitionerString = getCassandraContext().getSidecarClient().nodeSettings().get().partitioner();
+ }
+ partitioner = Partitioner.from(partitionerString);
+ }
+ catch (ExecutionException | InterruptedException exception)
+ {
+ throw new RuntimeException("Unable to retrieve partitioner information", exception);
+ }
+ }
+ return partitioner;
+ }
+ }
+
+ @Override
+ public TimeSkewResponse getTimeSkew(List<RingInstance> replicas)
+ {
+ try
+ {
+ List<SidecarInstance> instances = replicas
+ .stream()
+ .map(replica -> new SidecarInstanceImpl(replica.getNodeName(), conf.getSidecarPort()))
+ .collect(Collectors.toList());
+ return getCassandraContext().getSidecarClient().timeSkew(instances).get();
+ }
+ catch (InterruptedException | ExecutionException exception)
+ {
+ throw new RuntimeException(exception);
+ }
+ }
+
+ @Override
+ public void refreshClusterInfo()
+ {
+ synchronized (this)
+ {
+ // Set backing stores to null and let them lazy-load on the next call
+ ringResponse = null;
+ gossipInfo = null;
+ keyspaceSchema = null;
+ getCassandraContext().refreshClusterConfig();
+ }
+ }
+
+ protected String getCurrentKeyspaceSchema() throws Exception
+ {
+ SchemaResponse schemaResponse = getCassandraContext().getSidecarClient().schema(conf.keyspace).get();
+ return schemaResponse.schema();
+ }
+
+ @NotNull
+ protected CassandraRing<RingInstance> getCurrentRing() throws Exception
+ {
+ RingResponse ringResponse = getCurrentRingResponse();
+ List<RingInstance> instances = getSerializableInstances(ringResponse);
+ ReplicationFactor replicationFactor = getReplicationFactor();
+ return new CassandraRing<>(getPartitioner(), conf.keyspace, replicationFactor, instances);
+ }
+
+ @NotNull
+ protected ReplicationFactor getReplicationFactor()
+ {
+ String keyspaceSchema = getKeyspaceSchema(true);
+ if (keyspaceSchema == null)
+ {
+ throw new RuntimeException(String.format("Could not keyspace schema information for keyspace %s",
+ conf.keyspace));
+ }
+ return CqlUtils.extractReplicationFactor(keyspaceSchema, conf.keyspace);
+ }
+
+ @Override
+ public String getKeyspaceSchema(boolean cached)
+ {
+ String currentKeyspaceSchema = keyspaceSchema;
+ if (cached && currentKeyspaceSchema != null)
+ {
+ return currentKeyspaceSchema;
+ }
+
+ synchronized (this)
+ {
+ if (!cached || keyspaceSchema == null)
+ {
+ try
+ {
+ keyspaceSchema = getCurrentKeyspaceSchema();
+ }
+ catch (Exception exception)
+ {
+ throw new RuntimeException("Unable to initialize schema information for keyspace " + conf.keyspace,
+ exception);
+ }
+ }
+ return keyspaceSchema;
+ }
+ }
+
+ @Override
+ public CassandraRing<RingInstance> getRing(boolean cached)
+ {
+ CassandraRing<RingInstance> currentRing = ring;
+ if (cached && currentRing != null)
+ {
+ return currentRing;
+ }
+
+ synchronized (this)
+ {
+ if (!cached || ring == null)
+ {
+ try
+ {
+ ring = getCurrentRing();
+ }
+ catch (Exception exception)
+ {
+ throw new RuntimeException("Unable to initialize ring information", exception);
+ }
+ }
+ return ring;
+ }
+ }
+
+ @Override
+ public String getLowestCassandraVersion()
+ {
+ String currentCassandraVersion = cassandraVersion;
+ if (currentCassandraVersion != null)
+ {
+ return currentCassandraVersion;
+ }
+
+ synchronized (this)
+ {
+ if (cassandraVersion == null)
+ {
+ String versionFromFeature = getVersionFromFeature();
+ if (versionFromFeature != null)
+ {
+ // Forcing writer to use a particular version
+ cassandraVersion = versionFromFeature;
+ }
+ else
+ {
+ cassandraVersion = getVersionFromSidecar();
+ }
+ }
+ }
+ return cassandraVersion;
+ }
+
+ public String getVersionFromFeature()
+ {
+ return null;
+ }
+
+ public String getVersionFromSidecar()
+ {
+ LOGGER.info("Getting Cassandra versions from all nodes");
+ List<NodeSettings> allNodeSettings = Sidecar.allNodeSettingsBlocking(conf,
+ cassandraContext.getSidecarClient(),
+ cassandraContext.clusterConfig);
+ return getLowestVersion(allNodeSettings);
+ }
+
+ protected RingResponse getRingResponse()
+ {
+ RingResponse currentRingResponse = ringResponse;
+ if (currentRingResponse != null)
+ {
+ return currentRingResponse;
+ }
+
+ synchronized (this)
+ {
+ if (ringResponse == null)
+ {
+ try
+ {
+ ringResponse = getCurrentRingResponse();
+ }
+ catch (Exception exception)
+ {
+ LOGGER.error("Failed to load Cassandra ring", exception);
+ throw new RuntimeException(exception);
+ }
+ }
+ return ringResponse;
+ }
+ }
+
+ private RingResponse getCurrentRingResponse() throws Exception
+ {
+ return getCassandraContext().getSidecarClient().ring(conf.keyspace).get();
+ }
+
+ private static List<RingInstance> getSerializableInstances(RingResponse ringResponse)
+ {
+ return ringResponse.stream()
+ .map(RingInstance::new)
+ .collect(Collectors.toList());
+ }
+
+ private static RingInstance getCasInstanceMethodsImpl(RingEntry ringEntry)
+ {
+ return new RingInstance(ringEntry);
+ }
+
+ protected GossipInfoResponse getGossipInfo(boolean forceRefresh)
+ {
+ GossipInfoResponse currentGossipInfo = gossipInfo;
+ if (!forceRefresh && currentGossipInfo != null)
+ {
+ return currentGossipInfo;
+ }
+
+ synchronized (this)
+ {
+ if (forceRefresh || gossipInfo == null)
+ {
+ try
+ {
+ gossipInfo = cassandraContext.getSidecarClient().gossipInfo().get(conf.getHttpResponseTimeoutMs(),
+ TimeUnit.MILLISECONDS);
+ }
+ catch (ExecutionException | InterruptedException exception)
+ {
+ LOGGER.error("Failed to retrieve gossip information");
+ throw new RuntimeException("Failed to retrieve gossip information", exception);
+ }
+ catch (TimeoutException exception)
+ {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Failed to retrieve gossip information", exception);
+ }
+ }
+ return gossipInfo;
+ }
+ }
+
+ private InstanceAvailability determineInstanceAvailability(RingEntry ringEntry)
+ {
+ if (!instanceIsUp(ringEntry))
+ {
+ return InstanceAvailability.UNAVAILABLE_DOWN;
+ }
+ else if (instanceIsJoining(ringEntry) && isReplacement(ringEntry))
+ {
+ return InstanceAvailability.UNAVAILABLE_REPLACEMENT;
+ }
+ else if (instanceIsBlocked(getCasInstanceMethodsImpl(ringEntry)))
+ {
+ return InstanceAvailability.UNAVAILABLE_BLOCKED;
+ }
+ else if (instanceIsNormal(ringEntry))
+ {
+ return InstanceAvailability.AVAILABLE;
+ }
+ else
+ {
+ // If it's not one of the above, it's inherently INVALID
+ return InstanceAvailability.INVALID_STATE;
+ }
+ }
+
+ @VisibleForTesting
+ public String getLowestVersion(List<NodeSettings> allNodeSettings)
+ {
+ nodeSettings = allNodeSettings
+ .stream()
+ .filter(settings -> !settings.releaseVersion().equalsIgnoreCase("unknown"))
+ .min(Comparator.comparing(settings ->
+ CassandraVersionFeatures.cassandraVersionFeaturesFromCassandraVersion(settings.releaseVersion())))
+ .orElseThrow(() -> new RuntimeException("No valid Cassandra Versions were returned from Cassandra Sidecar"));
+ return nodeSettings.releaseVersion();
+ }
+
+ protected boolean instanceIsBlocked(RingInstance ignored)
+ {
+ return false;
+ }
+
+ protected boolean instanceIsNormal(RingEntry ringEntry)
+ {
+ return InstanceState.NORMAL.name().equalsIgnoreCase(ringEntry.state());
+ }
+
+ protected boolean instanceIsUp(RingEntry ringEntry)
+ {
+ return InstanceStatus.UP.name().equalsIgnoreCase(ringEntry.status());
+ }
+
+ protected boolean instanceIsJoining(RingEntry ringEntry)
+ {
+ return InstanceState.JOINING.name().equalsIgnoreCase(ringEntry.state());
+ }
+
+ protected boolean isReplacement(RingEntry ringEntry)
+ {
+ GossipInfoResponse gossipInfo = getGossipInfo(true);
+ LOGGER.debug("Gossip info={}", gossipInfo);
+
+ GossipInfoResponse.GossipInfo hostInfo = gossipInfo.get(ringEntry.address());
+ if (hostInfo != null)
+ {
+ LOGGER.info("Found hostinfo: {}", hostInfo);
+ String hostStatus = hostInfo.status();
+ if (hostStatus != null)
+ {
+ // If status has gone to NORMAL, we can't determine here if this was a host replacement or not.
+ // CassandraRingManager will handle detecting the ring change if it's gone NORMAL after the job starts.
+ return hostStatus.startsWith("BOOT_REPLACE,") || hostStatus.equals("NORMAL");
+ }
+ }
+ return false;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraContext.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraContext.java
new file mode 100644
index 0000000..e5b6853
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraContext.java
@@ -0,0 +1,99 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Closeable;
+import java.util.ArrayList;
+import java.util.Set;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.clients.Sidecar;
+import org.apache.cassandra.sidecar.client.SidecarClient;
+import org.apache.cassandra.sidecar.client.SidecarInstance;
+import org.apache.cassandra.sidecar.client.SimpleSidecarInstancesProvider;
+import org.jetbrains.annotations.NotNull;
+
+public class CassandraContext implements Closeable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(CassandraContext.class);
+ @NotNull
+ protected Set<? extends SidecarInstance> clusterConfig;
+ private final BulkSparkConf conf;
+ private final transient SidecarClient sidecarClient;
+
+ protected CassandraContext(BulkSparkConf conf)
+ {
+ this.conf = conf;
+ this.clusterConfig = createClusterConfig();
+ this.sidecarClient = initializeSidecarClient(conf);
+ LOGGER.debug("[{}] Created Cassandra Context", Thread.currentThread().getName());
+ }
+
+ public static CassandraContext create(BulkSparkConf conf)
+ {
+ return new CassandraContext(conf);
+ }
+
+ public Set<? extends SidecarInstance> getCluster()
+ {
+ return clusterConfig;
+ }
+
+ public void refreshClusterConfig()
+ {
+ // DO NOTHING
+ }
+
+ @Override
+ public void close()
+ {
+ try
+ {
+ sidecarClient.close();
+ LOGGER.debug("[{}] Closed Cassandra Context", Thread.currentThread().getName());
+ }
+ catch (Throwable throwable)
+ {
+ LOGGER.error("Could not shut down CassandraContext.", throwable);
+ }
+ }
+
+ protected SidecarClient initializeSidecarClient(BulkSparkConf conf)
+ {
+ return Sidecar.from(new SimpleSidecarInstancesProvider(new ArrayList<>(clusterConfig)), conf);
+ }
+
+ protected Set<? extends SidecarInstance> createClusterConfig()
+ {
+ return conf.sidecarInstances;
+ }
+
+ public SidecarClient getSidecarClient()
+ {
+ return sidecarClient;
+ }
+
+ protected BulkSparkConf conf()
+ {
+ return conf;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraJobInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraJobInfo.java
new file mode 100644
index 0000000..3151d6e
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraJobInfo.java
@@ -0,0 +1,119 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.util.UUID;
+
+import org.apache.cassandra.spark.bulkwriter.token.ConsistencyLevel;
+import org.jetbrains.annotations.NotNull;
+
+public class CassandraJobInfo implements JobInfo
+{
+ private static final long serialVersionUID = 6140098484732683759L;
+ private final BulkSparkConf conf;
+ @NotNull
+ private final UUID jobId = UUID.randomUUID();
+ private final TokenPartitioner tokenPartitioner;
+
+ CassandraJobInfo(BulkSparkConf conf, TokenPartitioner tokenPartitioner)
+ {
+ this.conf = conf;
+ this.tokenPartitioner = tokenPartitioner;
+ }
+
+ @Override
+ public ConsistencyLevel getConsistencyLevel()
+ {
+ return conf.consistencyLevel;
+ }
+
+ @Override
+ public String getLocalDC()
+ {
+ return conf.localDC;
+ }
+
+ @Override
+ public RowBufferMode getRowBufferMode()
+ {
+ return conf.rowBufferMode;
+ }
+
+ @Override
+ public int getSstableDataSizeInMB()
+ {
+ return conf.sstableDataSizeInMB;
+ }
+
+ @Override
+ public int getSstableBatchSize()
+ {
+ return conf.sstableBatchSize;
+ }
+
+ @Override
+ public int getCommitBatchSize()
+ {
+ return conf.commitBatchSize;
+ }
+
+ @Override
+ public boolean validateSSTables()
+ {
+ return conf.validateSSTables;
+ }
+
+ @Override
+ public boolean skipExtendedVerify()
+ {
+ return conf.skipExtendedVerify;
+ }
+
+ @Override
+ public boolean getSkipClean()
+ {
+ return conf.getSkipClean();
+ }
+
+ @Override
+ public int getCommitThreadsPerInstance()
+ {
+ return conf.commitThreadsPerInstance;
+ }
+
+ @Override
+ public UUID getId()
+ {
+ return jobId;
+ }
+
+ @Override
+ public TokenPartitioner getTokenPartitioner()
+ {
+ return tokenPartitioner;
+ }
+
+ @Override
+ @NotNull
+ public String getFullTableName()
+ {
+ return conf.keyspace + "." + conf.table;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraRingMonitor.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraRingMonitor.java
new file mode 100644
index 0000000..ead2cc3
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraRingMonitor.java
@@ -0,0 +1,135 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.math.BigInteger;
+import java.util.Collection;
+import java.util.concurrent.Callable;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Range;
+import org.apache.commons.collections.CollectionUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.bulkwriter.token.CassandraRing;
+import org.apache.cassandra.spark.bulkwriter.util.ThreadUtil;
+import org.apache.cassandra.spark.common.client.ClientException;
+
+public class CassandraRingMonitor
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(CassandraRingMonitor.class);
+
+ private CassandraRing<RingInstance> initialRing;
+ private volatile boolean ringChanged = false;
+ private Callable<CassandraRing<RingInstance>> ringSupplier;
+ private Consumer<CancelJobEvent> changeFunc;
+ private final ScheduledExecutorService executorService;
+
+ public CassandraRingMonitor(ClusterInfo clusterInfo,
+ Consumer<CancelJobEvent> changeFunc,
+ int checkInterval,
+ TimeUnit checkIntervalUnit) throws Exception
+ {
+ this(() -> clusterInfo.getRing(false),
+ changeFunc,
+ checkInterval,
+ checkIntervalUnit,
+ Executors.newSingleThreadScheduledExecutor(ThreadUtil.threadFactory("Cassandra Ring Monitor")));
+ }
+
+ @VisibleForTesting
+ public CassandraRingMonitor(Callable<CassandraRing<RingInstance>> ringSupplier,
+ Consumer<CancelJobEvent> changeFunc,
+ int checkInterval,
+ TimeUnit checkIntervalUnit,
+ ScheduledExecutorService executorService) throws Exception
+ {
+ this.ringSupplier = ringSupplier;
+ this.initialRing = ringSupplier.call();
+ this.changeFunc = changeFunc;
+ this.executorService = executorService;
+ executorService.scheduleAtFixedRate(this::checkRingChanged, 0, checkInterval, checkIntervalUnit);
+ }
+
+ /**
+ * Reads the ring changed value from the monitor. This actual ring
+ * is retrieved and compared every second - this is just the backing
+ * store for the result.
+ *
+ * @return if the ring has changed or not since this instance was created
+ */
+ public boolean getRingChanged()
+ {
+ return ringChanged;
+ }
+
+ public void stop()
+ {
+ executorService.shutdown();
+ try
+ {
+ executorService.awaitTermination(2, TimeUnit.SECONDS);
+ }
+ catch (InterruptedException exception)
+ {
+ // Do nothing as we're just waiting for the ring call to complete
+ }
+ }
+
+ private void checkRingChanged()
+ {
+ LOGGER.debug("Checking ring for changes");
+ try
+ {
+ CassandraRing<RingInstance> currentRing = ringSupplier.call();
+ if (checkRingChanged(currentRing))
+ {
+ ringChanged = true;
+ changeFunc.accept(new CancelJobEvent("Ring changed during bulk load"));
+ executorService.shutdownNow();
+ }
+ }
+ catch (ClientException exception)
+ {
+ LOGGER.warn("Could not retrieve current ring. Will retry momentarily. Continuing bulk load.", exception);
+ }
+ catch (Exception exception)
+ {
+ String message = "Error while attempting to determine if ring changed. Failing job";
+ LOGGER.error(message, exception);
+ changeFunc.accept(new CancelJobEvent(message, exception));
+ executorService.shutdownNow();
+ }
+ }
+
+ private boolean checkRingChanged(CassandraRing<RingInstance> currentRing)
+ {
+ Collection<Range<BigInteger>> currentTokenRanges = currentRing.getTokenRanges().values();
+ Collection<Range<BigInteger>> startingTokenRanges = initialRing.getTokenRanges().values();
+ LOGGER.debug("Current token range: {}", currentTokenRanges);
+ LOGGER.debug("Initial token range: {}", startingTokenRanges);
+ return !CollectionUtils.isEqualCollection(currentTokenRanges, startingTokenRanges);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraSchemaInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraSchemaInfo.java
new file mode 100644
index 0000000..d55b49b
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CassandraSchemaInfo.java
@@ -0,0 +1,37 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public class CassandraSchemaInfo implements SchemaInfo
+{
+ private static final long serialVersionUID = -2327383232935001862L;
+ private final TableSchema tableSchema;
+
+ public CassandraSchemaInfo(TableSchema tableSchema)
+ {
+ this.tableSchema = tableSchema;
+ }
+
+ @Override
+ public TableSchema getTableSchema()
+ {
+ return tableSchema;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java
new file mode 100644
index 0000000..5b65a55
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/ClusterInfo.java
@@ -0,0 +1,52 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.util.List;
+import java.util.Map;
+
+import org.apache.cassandra.sidecar.common.data.TimeSkewResponse;
+import org.apache.cassandra.spark.bulkwriter.token.CassandraRing;
+import org.apache.cassandra.spark.common.client.InstanceState;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+
+public interface ClusterInfo extends Serializable
+{
+ void refreshClusterInfo();
+
+ CassandraRing<RingInstance> getRing(boolean cached);
+
+ String getLowestCassandraVersion();
+
+ Map<RingInstance, InstanceAvailability> getInstanceAvailability();
+
+ boolean instanceIsAvailable(RingInstance ringInstance);
+
+ InstanceState getInstanceState(RingInstance instance);
+
+ Partitioner getPartitioner();
+
+ void checkBulkWriterIsEnabledOrThrow();
+
+ TimeSkewResponse getTimeSkew(List<RingInstance> replicas);
+
+ String getKeyspaceSchema(boolean cached);
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitCoordinator.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitCoordinator.java
new file mode 100644
index 0000000..2892a12
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitCoordinator.java
@@ -0,0 +1,223 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.math.BigInteger;
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import com.google.common.collect.Lists;
+import com.google.common.collect.Range;
+import com.google.common.util.concurrent.AbstractFuture;
+import com.google.common.util.concurrent.FutureCallback;
+import com.google.common.util.concurrent.Futures;
+import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.bulkwriter.util.ThreadUtil;
+import org.jetbrains.annotations.Nullable;
+
+public final class CommitCoordinator extends AbstractFuture<List<CommitResult>> implements AutoCloseable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(CommitCoordinator.class);
+
+ private final HashMap<RingInstance, ListeningExecutorService> executors = new HashMap<>();
+ private final List<StreamResult> successfulUploads;
+ private final DataTransferApi transferApi;
+ private final ClusterInfo cluster;
+ private final JobInfo job;
+ private ListenableFuture<List<CommitResult>> allCommits;
+ private final String jobSufix;
+
+ public static CommitCoordinator commit(BulkWriterContext bulkWriterContext, StreamResult[] uploadResults)
+ {
+ CommitCoordinator coordinator = new CommitCoordinator(bulkWriterContext, uploadResults);
+ coordinator.commit();
+ return coordinator;
+ }
+
+ private CommitCoordinator(BulkWriterContext writerContext, StreamResult[] uploadResults)
+ {
+ this.transferApi = writerContext.transfer();
+ this.cluster = writerContext.cluster();
+ this.job = writerContext.job();
+ this.jobSufix = "-" + job.getId();
+ successfulUploads = Arrays.stream(uploadResults)
+ .filter(result -> !result.passed.isEmpty())
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public boolean cancel(boolean mayInterruptIfRunning)
+ {
+ close();
+ return allCommits == null || allCommits.cancel(mayInterruptIfRunning);
+ }
+
+ void commit()
+ {
+ // We may have already committed - we should never get here if we did, but if somehow we do we should
+ // simply return the commit results we already collected
+ if (successfulUploads.size() > 0 && successfulUploads.stream()
+ .allMatch(result -> result.commitResults != null
+ && result.commitResults.size() > 0))
+ {
+ List<CommitResult> collect = successfulUploads.stream()
+ .flatMap(streamResult -> streamResult.commitResults.stream())
+ .collect(Collectors.toList());
+ set(collect);
+ return;
+ }
+ // First, group commits by instance so we can multi-commit
+ Map<RingInstance, Map<String, Range<BigInteger>>> resultsByInstance = getResultsByInstance(successfulUploads);
+ List<ListenableFuture<CommitResult>> commitFutures = resultsByInstance.entrySet().stream()
+ .flatMap(entry -> commit(executors, entry.getKey(), entry.getValue()))
+ .collect(Collectors.toList());
+ // Create an aggregate ListenableFuture around the list of futures containing the results of the commit calls.
+ // We'll fail fast if any of those errMsg (note that an errMsg here means an unexpected exception,
+ // not a failure response from CassandraManager).
+ // The callback on the aggregate listener sets the return value for this AbstractFuture
+ // so callers can make blocking calls to CommitCoordinator::get.
+ allCommits = Futures.allAsList(commitFutures);
+ Futures.addCallback(allCommits,
+ new FutureCallback<List<CommitResult>>()
+ {
+ public void onSuccess(@Nullable List<CommitResult> result)
+ {
+ set(result);
+ }
+
+ public void onFailure(Throwable throwable)
+ {
+ setException(throwable);
+ }
+ },
+ Runnable::run);
+ }
+
+ private Stream<ListenableFuture<CommitResult>> commit(Map<RingInstance, ListeningExecutorService> executors,
+ RingInstance instance,
+ Map<String, Range<BigInteger>> uploadRanges)
+ {
+ if (cluster.instanceIsAvailable(instance))
+ {
+ ListeningExecutorService executorService = executors.computeIfAbsent(instance,
+ key -> MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(
+ job.getCommitThreadsPerInstance(),
+ ThreadUtil.threadFactory("commit-sstable-" + key.getNodeName()))));
+ List<String> allUuids = new ArrayList<>(uploadRanges.keySet());
+ LOGGER.info("Committing UUIDs={}, Ranges={}, instance={}",
+ allUuids, uploadRanges.values(), instance.getNodeName());
+ List<List<String>> batches = Lists.partition(allUuids, job.getCommitBatchSize());
+ return batches.stream().map(uuids -> {
+ String migrationId = UUID.randomUUID().toString();
+ return executorService.submit(() -> {
+ CommitResult commitResult = new CommitResult(migrationId, instance, uploadRanges);
+ try
+ {
+ DataTransferApi.RemoteCommitResult result = transferApi.commitSSTables(instance, migrationId, uuids);
+ if (result.isSuccess)
+ {
+ LOGGER.info("[{}]: Commit succeeded on {} for {}",
+ migrationId, instance.getNodeName(), uploadRanges);
+ }
+ else
+ {
+ LOGGER.error("[{}]: Commit failed: uploadRanges: {}, failedUuids: {}, stdErr: {}",
+ migrationId, uploadRanges.entrySet(), result.failedUuids, result.stdErr);
+ if (result.failedUuids.size() > 0)
+ {
+ addFailures(result.failedUuids, uploadRanges, commitResult, result.stdErr);
+ }
+ else
+ {
+ addFailures(uploadRanges, commitResult, result.stdErr);
+ }
+ }
+ }
+ catch (Throwable throwable)
+ {
+ addFailures(uploadRanges, commitResult, throwable.toString());
+ // On errMsg, refresh cluster information so we get the latest block list and status information to react accordingly
+ cluster.refreshClusterInfo();
+ }
+ return commitResult;
+ });
+ });
+ }
+ else
+ {
+ String migrationId = UUID.randomUUID().toString();
+ CommitResult commitResult = new CommitResult(migrationId, instance, uploadRanges);
+
+ String message = String.format("Instance %s is not available.", instance.getNodeName());
+ addFailures(uploadRanges, commitResult, message);
+ return Stream.of(Futures.immediateFuture(commitResult));
+ }
+ }
+
+ private void addFailures(List<String> failedRanges,
+ Map<String, Range<BigInteger>> uploadRanges,
+ CommitResult commitResult,
+ String error)
+ {
+ failedRanges.forEach(uuid -> {
+ String shortUuid = uuid.replace(jobSufix, "");
+ commitResult.addFailedCommit(shortUuid, uploadRanges.get(shortUuid), error != null ? error : "Unknown Commit Failure");
+ });
+ }
+
+ private void addFailures(Map<String, Range<BigInteger>> failedRanges, CommitResult commitResult, String message)
+ {
+ failedRanges.forEach((key, value) -> commitResult.addFailedCommit(key, value, message));
+ LOGGER.debug("Added failures to commitResult by Range: {}", commitResult);
+ }
+
+ private Map<RingInstance, Map<String, Range<BigInteger>>> getResultsByInstance(List<StreamResult> successfulUploads)
+ {
+ return successfulUploads
+ .stream()
+ .flatMap(upload -> upload.passed
+ .stream()
+ .map(instance -> new AbstractMap.SimpleEntry<>(instance,
+ new AbstractMap.SimpleEntry<>(upload.sessionID, upload.tokenRange))))
+ .collect(Collectors.groupingBy(AbstractMap.SimpleEntry::getKey,
+ Collectors.toMap(instance -> instance.getValue().getKey(),
+ instance -> instance.getValue().getValue())));
+ }
+
+ @Override
+ public void close()
+ {
+ executors.values().forEach(ExecutorService::shutdownNow);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitError.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitError.java
new file mode 100644
index 0000000..62dcc02
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitError.java
@@ -0,0 +1,48 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.math.BigInteger;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Range;
+
+import org.jetbrains.annotations.NotNull;
+
+public class CommitError implements Serializable
+{
+ protected final Range<BigInteger> tokenRange;
+ protected final String errMsg;
+
+ public CommitError(@NotNull Range<BigInteger> tokenRange, @NotNull String errMsg)
+ {
+ Preconditions.checkNotNull(tokenRange, "CommitError created without a token range");
+ Preconditions.checkNotNull(errMsg, "CommitError created without an error message");
+ this.tokenRange = tokenRange;
+ this.errMsg = errMsg;
+ }
+
+ @Override
+ public String toString()
+ {
+ return String.format("CommitError{tokenRange=%s, errMsg='%s'}", tokenRange, errMsg);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitResult.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitResult.java
new file mode 100644
index 0000000..6020d05
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CommitResult.java
@@ -0,0 +1,66 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.math.BigInteger;
+import java.util.HashMap;
+import java.util.Map;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Range;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.jetbrains.annotations.NotNull;
+
+public class CommitResult implements Serializable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(StreamResult.class);
+
+ public final String migrationId;
+ protected RingInstance instance;
+ public final Map<String, Range<BigInteger>> passed;
+ public final Map<String, CommitError> failures;
+
+ public CommitResult(String migrationId, RingInstance instance, Map<String, Range<BigInteger>> commitRanges)
+ {
+ this.migrationId = migrationId;
+ this.instance = instance;
+ this.passed = new HashMap<>(commitRanges);
+ this.failures = new HashMap<>();
+ }
+
+ public void addFailedCommit(String uuid, @NotNull Range<BigInteger> tokenRange, @NotNull String error)
+ {
+ Preconditions.checkNotNull(uuid, "Adding failed commit with missing UUID");
+ Preconditions.checkNotNull(tokenRange, "Adding failed commit with missing token range");
+ Preconditions.checkNotNull(error, "Adding failed commit with missing error message");
+ LOGGER.error("[{}]: Failed to commit {} on {}: {}", uuid, tokenRange, instance, error);
+ passed.remove(uuid);
+ failures.put(uuid, new CommitError(tokenRange, error));
+ }
+
+ @Override
+ public String toString()
+ {
+ return String.format("CommitResult{migrationId='%s', instance=%s, passed=%s, failures=%s}", migrationId, instance, passed, failures);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CqlTableInfoProvider.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CqlTableInfoProvider.java
new file mode 100644
index 0000000..e512d27
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/CqlTableInfoProvider.java
@@ -0,0 +1,175 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.util.List;
+import java.util.Map;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+
+import org.apache.cassandra.spark.common.schema.ColumnType;
+import org.apache.cassandra.spark.common.schema.ColumnTypes;
+import org.apache.cassandra.spark.data.CqlField;
+import org.apache.cassandra.spark.data.CqlTable;
+import org.jetbrains.annotations.NotNull;
+
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.ASCII;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.BIGINT;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.BLOB;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.BOOLEAN;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.DOUBLE;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.INT;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.TEXT;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.TIMESTAMP;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.TIMEUUID;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.UUID;
+import static org.apache.cassandra.spark.bulkwriter.SqlToCqlTypeConverter.VARCHAR;
+
+/**
+ * An implementation of the {@link TableInfoProvider} interface that leverages the {@link CqlTable} to
+ * provide table information
+ */
+public class CqlTableInfoProvider implements TableInfoProvider
+{
+ // CHECKSTYLE IGNORE: Long line
+ private static final Pattern DEPRECATED_OPTIONS = Pattern.compile("(?:\\s+AND)?\\s+(?:dclocal_)?read_repair_chance\\s*=\\s*[+\\-.\\dE]+", Pattern.CASE_INSENSITIVE);
+ private static final Pattern LEADING_CONJUNCTION = Pattern.compile("(?<=\\bWITH)\\s+AND\\b", Pattern.CASE_INSENSITIVE);
+
+ private final String createTableStatement;
+ private final CqlTable cqlTable;
+
+ /**
+ * Mapping from CQL Type Names to Java types
+ */
+ private static final Map<String, ColumnType<?>> DATA_TYPES = ImmutableMap.<String, ColumnType<?>>builder()
+ .put(BIGINT, ColumnTypes.LONG)
+ .put(BLOB, ColumnTypes.BYTES)
+ .put(DOUBLE, ColumnTypes.DOUBLE)
+ .put(INT, ColumnTypes.INT)
+ .put(BOOLEAN, ColumnTypes.BOOLEAN)
+ .put(TEXT, ColumnTypes.STRING)
+ .put(TIMESTAMP, ColumnTypes.LONG)
+ .put(UUID, ColumnTypes.UUID)
+ .put(VARCHAR, ColumnTypes.STRING)
+ .put(ASCII, ColumnTypes.STRING)
+ .put(TIMEUUID, ColumnTypes.UUID)
+ .build();
+
+ public CqlTableInfoProvider(String createTableStatement, CqlTable cqlTable)
+ {
+ this.createTableStatement = createTableStatement;
+ this.cqlTable = cqlTable;
+ }
+
+ @Override
+ public CqlField.CqlType getColumnType(String columnName)
+ {
+ CqlField cqlField = cqlTable.column(columnName);
+ if (cqlField == null)
+ {
+ throw new IllegalArgumentException("Unknown fields in data frame => " + columnName);
+ }
+ else
+ {
+ return cqlField.type();
+ }
+ }
+
+ @Override
+ public List<ColumnType<?>> getPartitionKeyTypes()
+ {
+ return cqlTable.partitionKeys().stream()
+ .map(cqlField -> DATA_TYPES.get(cqlField.type().cqlName().toLowerCase()))
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public boolean columnExists(String columnName)
+ {
+ return cqlTable.column(columnName) != null;
+ }
+
+ @Override
+ public List<String> getPartitionKeyColumnNames()
+ {
+ return cqlTable.partitionKeys().stream()
+ .map(CqlField::name)
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public String getCreateStatement()
+ {
+ return removeDeprecatedOptions(createTableStatement);
+ }
+
+ @Override
+ public List<String> getPrimaryKeyColumnNames()
+ {
+ return cqlTable.primaryKey().stream()
+ .map(CqlField::name)
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public String getName()
+ {
+ return cqlTable.table();
+ }
+
+ @Override
+ public String getKeyspaceName()
+ {
+ return cqlTable.keyspace();
+ }
+
+ @Override
+ public boolean hasSecondaryIndex()
+ {
+ return cqlTable.indexCount() > 0;
+ }
+
+ @Override
+ public List<String> getColumnNames()
+ {
+ return cqlTable.columns().stream()
+ .map(CqlField::name)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Removes table options {@code read_repair_chance} and {@code dclocal_read_repair_chance}
+ * that were deprecated in Cassandra 4.0 from the table creation CQL statement
+ *
+ * @param cql table creation CQL statement
+ * @return table creation CQL statement with deprecated table options removed
+ */
+ @VisibleForTesting
+ @NotNull
+ static String removeDeprecatedOptions(@NotNull String cql)
+ {
+ cql = DEPRECATED_OPTIONS.matcher(cql).replaceAll("");
+ cql = LEADING_CONJUNCTION.matcher(cql).replaceAll("");
+ return cql;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DataTransferApi.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DataTransferApi.java
new file mode 100644
index 0000000..c0786a2
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DataTransferApi.java
@@ -0,0 +1,63 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.nio.file.Path;
+import java.util.List;
+
+import org.apache.cassandra.spark.common.MD5Hash;
+import org.apache.cassandra.spark.common.client.ClientException;
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+import org.jetbrains.annotations.Nullable;
+
+public interface DataTransferApi extends Serializable
+{
+ class RemoteCommitResult
+ {
+ public final boolean isSuccess;
+ public final List<String> failedUuids;
+ public final List<String> successfulUuids;
+ public final String stdErr;
+
+ public RemoteCommitResult(boolean isSuccess,
+ @Nullable List<String> failedUuids,
+ List<String> successfulUuids,
+ @Nullable String stdErr)
+ {
+ this.isSuccess = isSuccess;
+ this.failedUuids = failedUuids;
+ this.successfulUuids = successfulUuids;
+ this.stdErr = stdErr;
+ }
+ }
+
+ RemoteCommitResult commitSSTables(CassandraInstance instance,
+ String migrationId,
+ List<String> uuids) throws ClientException;
+
+ void cleanUploadSession(CassandraInstance instance, String sessionID, String jobID) throws ClientException;
+
+ void uploadSSTableComponent(Path componentFile,
+ int ssTableIdx,
+ CassandraInstance instance,
+ String sessionID,
+ MD5Hash fileHash) throws ClientException;
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DecoratedKey.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DecoratedKey.java
new file mode 100644
index 0000000..8d6e903
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/DecoratedKey.java
@@ -0,0 +1,92 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.math.BigInteger;
+import java.nio.ByteBuffer;
+
+import org.apache.cassandra.spark.bulkwriter.util.FastByteOperations;
+import org.apache.cassandra.spark.utils.ByteBufferUtils;
+import org.jetbrains.annotations.NotNull;
+
+public class DecoratedKey implements Comparable<DecoratedKey>, Serializable
+{
+ @NotNull private final BigInteger token;
+ @NotNull private final ByteBuffer key;
+
+ DecoratedKey(@NotNull BigInteger token, @NotNull ByteBuffer key)
+ {
+ this.token = token;
+ this.key = key;
+ }
+
+ @NotNull
+ public BigInteger getToken()
+ {
+ return token;
+ }
+
+ @NotNull
+ public ByteBuffer getKey()
+ {
+ return key;
+ }
+
+ @Override
+ public int compareTo(@NotNull DecoratedKey that)
+ {
+ int cmp = token.compareTo(that.token);
+ if (cmp != 0)
+ {
+ return cmp;
+ }
+ return FastByteOperations.compareUnsigned(key, that.key);
+ }
+
+ @Override
+ public boolean equals(Object other)
+ {
+ if (this == other)
+ {
+ return true;
+ }
+ if (other == null || this.getClass() != other.getClass())
+ {
+ return false;
+ }
+
+ DecoratedKey that = (DecoratedKey) other;
+ return this.token.equals(that.token) && this.key.equals(that.key);
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return token.hashCode() ^ key.hashCode();
+ }
+
+ @Override
+ public String toString()
+ {
+ String keystring = ByteBufferUtils.toHexString(getKey());
+ return "SBW-DecoratedKey(" + getToken() + ", " + keystring + ")";
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/InstanceAvailability.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/InstanceAvailability.java
new file mode 100644
index 0000000..2aff548
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/InstanceAvailability.java
@@ -0,0 +1,53 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public enum InstanceAvailability
+{
+ AVAILABLE("is available"),
+ UNAVAILABLE_DOWN("is down"),
+ /**
+ * This node is listed in the cluster block list
+ */
+ UNAVAILABLE_BLOCKED("is blocked"),
+ /**
+ * This node is JOINING, and it's gossip state is BOOT_REPLACE, which means it's a host replacement.
+ * We can support SBW jobs with host replacements ongoing, but we treat them as UNAVAILABLE so the
+ * job's success or failure depends on its consistency level requirements.
+ */
+ UNAVAILABLE_REPLACEMENT("is a host replacement"),
+ /**
+ * INVALID_STATE is true when a node is in any state but NORMAL, except for host replacements (see above).
+ * The SBW doesn't currently support hosts in non-NORMAL states.
+ */
+ INVALID_STATE("is in an invalid state");
+
+ private final String message;
+
+ InstanceAvailability(String message)
+ {
+ this.message = message;
+ }
+
+ public String getMessage()
+ {
+ return message;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/JobInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/JobInfo.java
new file mode 100644
index 0000000..945012b
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/JobInfo.java
@@ -0,0 +1,54 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.util.UUID;
+
+import org.apache.cassandra.spark.bulkwriter.token.ConsistencyLevel;
+
+public interface JobInfo extends Serializable
+{
+ // Job Information API - should this really just move back to Config? Here to try to reduce the violations of the Law of Demeter more than anything else
+ ConsistencyLevel getConsistencyLevel();
+
+ String getLocalDC();
+
+ RowBufferMode getRowBufferMode();
+
+ int getSstableDataSizeInMB();
+
+ int getSstableBatchSize();
+
+ int getCommitBatchSize();
+
+ int getCommitThreadsPerInstance();
+
+ UUID getId();
+
+ TokenPartitioner getTokenPartitioner();
+
+ boolean validateSSTables();
+
+ boolean skipExtendedVerify();
+
+ String getFullTableName();
+ boolean getSkipClean();
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/LoadNotSupportedException.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/LoadNotSupportedException.java
new file mode 100644
index 0000000..37cd600
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/LoadNotSupportedException.java
@@ -0,0 +1,28 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public class LoadNotSupportedException extends UnsupportedOperationException
+{
+ public LoadNotSupportedException(String message)
+ {
+ super(message);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RecordWriter.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RecordWriter.java
new file mode 100644
index 0000000..a5aaf8d
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RecordWriter.java
@@ -0,0 +1,205 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.math.BigInteger;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Iterator;
+import java.util.List;
+import java.util.UUID;
+import java.util.function.BiFunction;
+import java.util.function.Supplier;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Range;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.sidecar.common.data.TimeSkewResponse;
+import org.apache.spark.InterruptibleIterator;
+import org.apache.spark.TaskContext;
+import scala.Tuple2;
+
+import static org.apache.cassandra.spark.utils.ScalaConversionUtils.asScalaIterator;
+
+@SuppressWarnings({"ConstantConditions"})
+public class RecordWriter implements Serializable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(RecordWriter.class);
+
+ private final BulkWriterContext writerContext;
+ private Supplier<TaskContext> taskContextSupplier;
+ private final BiFunction<BulkWriterContext, Path, SSTableWriter> tableWriterSupplier;
+ private SSTableWriter sstableWriter = null;
+ private int batchNumber = 0;
+ private int batchSize = 0;
+
+ public RecordWriter(BulkWriterContext writerContext)
+ {
+ this(writerContext, TaskContext::get, SSTableWriter::new);
+ }
+
+ @VisibleForTesting
+ RecordWriter(BulkWriterContext writerContext,
+ Supplier<TaskContext> taskContextSupplier,
+ BiFunction<BulkWriterContext, Path, SSTableWriter> tableWriterSupplier)
+ {
+ this.writerContext = writerContext;
+ this.taskContextSupplier = taskContextSupplier;
+ this.tableWriterSupplier = tableWriterSupplier;
+ }
+
+ private Range<BigInteger> getTokenRange(TaskContext taskContext)
+ {
+ return writerContext.job().getTokenPartitioner().getTokenRange(taskContext.partitionId());
+ }
+
+ private String getStreamId(TaskContext taskContext)
+ {
+ return String.format("%d-%s", taskContext.partitionId(), UUID.randomUUID());
+ }
+
+ public StreamResult write(Iterator<Tuple2<DecoratedKey, Object[]>> sourceIterator)
+ {
+ TaskContext taskContext = taskContextSupplier.get();
+ LOGGER.info("[{}]: Processing Bulk Writer partition", taskContext.partitionId());
+ scala.collection.Iterator<scala.Tuple2<DecoratedKey, Object[]>> dataIterator =
+ new InterruptibleIterator<>(taskContext, asScalaIterator(sourceIterator));
+ StreamSession streamSession = createStreamSession(taskContext);
+ validateAcceptableTimeSkewOrThrow(streamSession.replicas);
+ int partitionId = taskContext.partitionId();
+ Range<BigInteger> range = getTokenRange(taskContext);
+ JobInfo job = writerContext.job();
+ Path baseDir = Paths.get(System.getProperty("java.io.tmpdir"),
+ job.getId().toString(),
+ Integer.toString(taskContext.stageAttemptNumber()),
+ Integer.toString(taskContext.attemptNumber()),
+ Integer.toString(partitionId));
+ try
+ {
+ while (dataIterator.hasNext())
+ {
+ maybeCreateTableWriter(partitionId, baseDir);
+ writeRow(dataIterator, partitionId, range);
+ checkBatchSize(streamSession, partitionId, job);
+ }
+
+ if (batchSize != 0)
+ {
+ finalizeSSTable(streamSession, partitionId, sstableWriter, batchNumber, batchSize);
+ }
+
+ LOGGER.info("[{}] Done with all writers and waiting for stream to complete", partitionId);
+ return streamSession.close();
+ }
+ catch (Exception exception)
+ {
+ throw new RuntimeException(exception);
+ }
+ }
+
+ private void validateAcceptableTimeSkewOrThrow(List<RingInstance> replicas)
+ {
+ TimeSkewResponse timeSkewResponse = writerContext.cluster().getTimeSkew(replicas);
+ Instant localNow = Instant.now();
+ Instant remoteNow = Instant.ofEpochMilli(timeSkewResponse.currentTime);
+ Duration range = Duration.ofMinutes(timeSkewResponse.allowableSkewInMinutes);
+ if (localNow.isBefore(remoteNow.minus(range)) || localNow.isAfter(remoteNow.plus(range)))
+ {
+ String message = String.format("Time skew between Spark and Cassandra is too large. "
+ + "Allowable skew is %d minutes. "
+ + "Spark executor time is %s, Cassandra instance time is %s",
+ timeSkewResponse.allowableSkewInMinutes, localNow, remoteNow);
+ throw new UnsupportedOperationException(message);
+ }
+ }
+
+ public void writeRow(scala.collection.Iterator<Tuple2<DecoratedKey, Object[]>> dataIterator,
+ int partitionId,
+ Range<BigInteger> range) throws IOException
+ {
+ Tuple2<DecoratedKey, Object[]> tuple = dataIterator.next();
+ DecoratedKey key = tuple._1();
+ BigInteger token = key.getToken();
+ Preconditions.checkState(range.contains(token),
+ String.format("Received Token %s outside of expected range %s", token, range));
+ try
+ {
+ sstableWriter.addRow(token, tuple._2());
+ }
+ catch (RuntimeException exception)
+ {
+ String message = String.format("[%s]: Failed to write data to SSTable: SBW DecoratedKey was %s",
+ partitionId, key);
+ LOGGER.error(message, exception);
+ throw exception;
+ }
+ }
+
+ public void checkBatchSize(StreamSession streamSession, int partitionId, JobInfo job) throws IOException
+ {
+ batchSize++;
+ if (batchSize > job.getSstableBatchSize())
+ {
+ finalizeSSTable(streamSession, partitionId, sstableWriter, batchNumber, batchSize);
+
+ sstableWriter = null;
+ batchSize = 0;
+
+ }
+ }
+
+ public void maybeCreateTableWriter(int partitionId, Path baseDir) throws IOException
+ {
+ if (sstableWriter == null)
+ {
+ Path outDir = Paths.get(baseDir.toString(), Integer.toString(++batchNumber));
+ Files.createDirectories(outDir);
+
+ sstableWriter = tableWriterSupplier.apply(writerContext, outDir);
+
+ LOGGER.info("[{}][{}] Created new SSTable writer", partitionId, batchNumber);
+ }
+ }
+
+ private void finalizeSSTable(StreamSession streamSession,
+ int partitionId,
+ SSTableWriter sstableWriter,
+ int batchNumber,
+ int batchSize) throws IOException
+ {
+ LOGGER.info("[{}][{}] Closing writer and scheduling SStable stream with {} rows",
+ partitionId, batchNumber, batchSize);
+ sstableWriter.close(writerContext, partitionId);
+ streamSession.scheduleStream(sstableWriter);
+ }
+
+ private StreamSession createStreamSession(TaskContext taskContext)
+ {
+ return new StreamSession(writerContext, getStreamId(taskContext), getTokenRange(taskContext));
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java
new file mode 100644
index 0000000..75a2896
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RingInstance.java
@@ -0,0 +1,144 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.util.Objects;
+
+import org.apache.cassandra.sidecar.common.data.RingEntry;
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+import org.jetbrains.annotations.Nullable;
+
+public class RingInstance implements CassandraInstance, Serializable
+{
+ private static final long serialVersionUID = 4399143234683369652L;
+ private RingEntry ringEntry;
+
+ public RingInstance(RingEntry ringEntry)
+ {
+ this.ringEntry = ringEntry;
+ }
+
+ @Override
+ public String getToken()
+ {
+ return ringEntry.token();
+ }
+
+ @Override
+ public String getNodeName()
+ {
+ return ringEntry.fqdn();
+ }
+
+ @Override
+ public String getDataCenter()
+ {
+ return ringEntry.datacenter();
+ }
+
+ /**
+ * Custom equality that compares the token, fully qualified domain name, the port, and the datacenter
+ *
+ * @param other the other instance
+ * @return true if both instances are equal, false otherwise
+ */
+ @Override
+ public boolean equals(@Nullable Object other)
+ {
+ if (other == null || !(other instanceof RingInstance))
+ {
+ return false;
+ }
+
+ RingInstance that = (RingInstance) other;
+ return Objects.equals(this.ringEntry.token(), that.ringEntry.token())
+ && Objects.equals(this.ringEntry.fqdn(), that.ringEntry.fqdn())
+ && this.ringEntry.port() == that.ringEntry.port()
+ && Objects.equals(this.ringEntry.datacenter(), that.ringEntry.datacenter());
+ }
+
+ /**
+ * Custom hashCode that compares the token, fully qualified domain name, the port, and the datacenter
+ *
+ * @return The hashcode of this instance based on the important fields
+ */
+ @Override
+ public int hashCode()
+ {
+ return Objects.hash(ringEntry.token(), ringEntry.fqdn(), ringEntry.port(), ringEntry.datacenter());
+ }
+
+ @Override
+ public String toString()
+ {
+ return ringEntry.toString();
+ }
+
+ public RingEntry getRingInstance()
+ {
+ return ringEntry;
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException
+ {
+ out.writeUTF(ringEntry.address());
+ out.writeInt(ringEntry.port());
+ out.writeUTF(ringEntry.datacenter());
+ out.writeUTF(ringEntry.load());
+ out.writeUTF(ringEntry.token());
+ out.writeUTF(ringEntry.fqdn());
+ out.writeUTF(ringEntry.rack());
+ out.writeUTF(ringEntry.hostId());
+ out.writeUTF(ringEntry.status());
+ out.writeUTF(ringEntry.state());
+ out.writeUTF(ringEntry.owns());
+ }
+
+ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException
+ {
+ String address = in.readUTF();
+ int port = in.readInt();
+ String datacenter = in.readUTF();
+ String load = in.readUTF();
+ String token = in.readUTF();
+ String fqdn = in.readUTF();
+ String rack = in.readUTF();
+ String hostId = in.readUTF();
+ String status = in.readUTF();
+ String state = in.readUTF();
+ String owns = in.readUTF();
+ ringEntry = new RingEntry.Builder().datacenter(datacenter)
+ .address(address)
+ .port(port)
+ .rack(rack)
+ .status(status)
+ .state(state)
+ .load(load)
+ .owns(owns)
+ .token(token)
+ .fqdn(fqdn)
+ .hostId(hostId)
+ .build();
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RowBufferMode.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RowBufferMode.java
new file mode 100644
index 0000000..d54ce3f
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/RowBufferMode.java
@@ -0,0 +1,26 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public enum RowBufferMode
+{
+ BUFFERRED,
+ UNBUFFERRED
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SSTableWriter.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SSTableWriter.java
new file mode 100644
index 0000000..bf982ef
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SSTableWriter.java
@@ -0,0 +1,188 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.math.BigInteger;
+import java.nio.file.DirectoryStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.security.MessageDigest;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Range;
+import org.apache.commons.codec.digest.DigestUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.bridge.CassandraBridgeFactory;
+import org.apache.cassandra.bridge.CassandraVersion;
+import org.apache.cassandra.bridge.CassandraVersionFeatures;
+import org.apache.cassandra.spark.common.MD5Hash;
+import org.apache.cassandra.spark.common.SSTables;
+import org.apache.cassandra.spark.data.DataLayer;
+import org.apache.cassandra.spark.data.LocalDataLayer;
+import org.apache.cassandra.spark.reader.StreamScanner;
+import org.jetbrains.annotations.NotNull;
+
+@SuppressWarnings("WeakerAccess")
+public class SSTableWriter
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(SSTableWriter.class);
+
+ public static final String CASSANDRA_VERSION_PREFIX = "cassandra-";
+
+ private final Path outDir;
+ private final org.apache.cassandra.bridge.SSTableWriter cqlSSTableWriter;
+ private BigInteger minToken = null;
+ private BigInteger maxToken = null;
+ private final Map<Path, MD5Hash> fileHashes = new HashMap<>();
+
+ public SSTableWriter(org.apache.cassandra.bridge.SSTableWriter tableWriter, Path outDir)
+ {
+ cqlSSTableWriter = tableWriter;
+ this.outDir = outDir;
+ }
+
+ public SSTableWriter(BulkWriterContext writerContext, Path outDir)
+ {
+ this.outDir = outDir;
+
+ String lowestCassandraVersion = writerContext.cluster().getLowestCassandraVersion();
+ String packageVersion = getPackageVersion(lowestCassandraVersion);
+ LOGGER.info("Running with version " + packageVersion);
+
+ TableSchema tableSchema = writerContext.schema().getTableSchema();
+ boolean sorted = writerContext.job().getRowBufferMode() == RowBufferMode.UNBUFFERRED;
+ this.cqlSSTableWriter = SSTableWriterFactory.getSSTableWriter(
+ CassandraVersionFeatures.cassandraVersionFeaturesFromCassandraVersion(packageVersion),
+ this.outDir.toString(),
+ writerContext.cluster().getPartitioner().toString(),
+ tableSchema.createStatement,
+ tableSchema.modificationStatement,
+ sorted,
+ writerContext.job().getSstableDataSizeInMB());
+ }
+
+ @NotNull
+ public String getPackageVersion(String lowestCassandraVersion)
+ {
+ return CASSANDRA_VERSION_PREFIX + lowestCassandraVersion;
+ }
+
+ public void addRow(BigInteger token, Object[] values) throws IOException
+ {
+ if (minToken == null)
+ {
+ minToken = token;
+ }
+ maxToken = token;
+
+ cqlSSTableWriter.addRow(values);
+ }
+
+ public void close(BulkWriterContext writerContext, int partitionId) throws IOException
+ {
+ cqlSSTableWriter.close();
+ for (Path dataFile : getDataFileStream())
+ {
+ // NOTE: We calculate file hashes before re-reading so that we know what we hashed
+ // is what we validated. Then we send these along with the files and the
+ // receiving end re-hashes the files to make sure they still match.
+ fileHashes.putAll(calculateFileHashes(dataFile));
+ }
+ validateSSTables(writerContext, partitionId);
+ }
+
+ @VisibleForTesting
+ public void validateSSTables(@NotNull BulkWriterContext writerContext, int partitionId)
+ {
+ // NOTE: If this current implementation of SS-tables' validation proves to be a performance issue,
+ // we will need to modify LocalDataLayer to allow scanning and compaction of single data file,
+ // and then validate all of them in parallel threads
+ try
+ {
+ CassandraVersion version = CassandraBridgeFactory.getCassandraVersion(writerContext.cluster().getLowestCassandraVersion());
+ String keyspace = writerContext.cluster().getRing(true).getKeyspace();
+ String schema = writerContext.schema().getTableSchema().createStatement;
+ String directory = getOutDir().toString();
+ DataLayer layer = new LocalDataLayer(version, keyspace, schema, directory);
+ try (StreamScanner scanner = layer.openCompactionScanner(partitionId, Collections.emptyList(), null))
+ {
+ while (scanner.hasNext())
+ {
+ scanner.advanceToNextColumn();
+ }
+ }
+ }
+ catch (IOException exception)
+ {
+ LOGGER.error("[{}]: Unexpected exception while validating SSTables {}", partitionId, getOutDir());
+ throw new RuntimeException(exception);
+ }
+ }
+
+ private DirectoryStream<Path> getDataFileStream() throws IOException
+ {
+ return Files.newDirectoryStream(getOutDir(), "*Data.db");
+ }
+
+ private Map<Path, MD5Hash> calculateFileHashes(Path dataFile) throws IOException
+ {
+ Map<Path, MD5Hash> fileHashes = new HashMap<>();
+ try (DirectoryStream<Path> filesToHash =
+ Files.newDirectoryStream(dataFile.getParent(), SSTables.getSSTableBaseName(dataFile) + "*"))
+ {
+ for (Path path: filesToHash)
+ {
+ fileHashes.put(path, calculateFileHash(path));
+ }
+ }
+ return fileHashes;
+ }
+
+ private MD5Hash calculateFileHash(Path path) throws IOException
+ {
+ try (InputStream is = Files.newInputStream(path))
+ {
+ MessageDigest computedMd5 = DigestUtils.updateDigest(DigestUtils.getMd5Digest(), is);
+ return MD5Hash.fromDigest(computedMd5);
+ }
+ }
+
+ public Range<BigInteger> getTokenRange()
+ {
+ return Range.closed(minToken, maxToken);
+ }
+
+ public Path getOutDir()
+ {
+ return outDir;
+ }
+
+ public Map<Path, MD5Hash> getFileHashes()
+ {
+ return fileHashes;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SSTableWriterFactory.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SSTableWriterFactory.java
new file mode 100644
index 0000000..832e370
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SSTableWriterFactory.java
@@ -0,0 +1,50 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import org.apache.cassandra.bridge.CassandraBridge;
+import org.apache.cassandra.bridge.CassandraBridgeFactory;
+import org.apache.cassandra.bridge.CassandraVersionFeatures;
+import org.apache.cassandra.bridge.SSTableWriter;
+
+public final class SSTableWriterFactory
+{
+ private SSTableWriterFactory()
+ {
+ throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
+ }
+
+ public static SSTableWriter getSSTableWriter(CassandraVersionFeatures serverVersion,
+ String inDirectory,
+ String partitioner,
+ String createStatement,
+ String insertStatement,
+ boolean sorted,
+ int bufferSizeMB)
+ {
+ CassandraBridge cassandraBridge = CassandraBridgeFactory.get(serverVersion);
+ return cassandraBridge.getSSTableWriter(inDirectory,
+ partitioner,
+ createStatement,
+ insertStatement,
+ sorted,
+ bufferSizeMB);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SchemaInfo.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SchemaInfo.java
new file mode 100644
index 0000000..ca95618
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SchemaInfo.java
@@ -0,0 +1,27 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+
+public interface SchemaInfo extends Serializable
+{
+ TableSchema getTableSchema();
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SidecarDataTransferApi.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SidecarDataTransferApi.java
new file mode 100644
index 0000000..4d5b25b
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SidecarDataTransferApi.java
@@ -0,0 +1,153 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.nio.file.Path;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.ExecutionException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.sidecar.client.SidecarClient;
+import org.apache.cassandra.sidecar.client.SidecarInstanceImpl;
+import org.apache.cassandra.sidecar.client.request.ImportSSTableRequest;
+import org.apache.cassandra.sidecar.common.data.SSTableImportResponse;
+import org.apache.cassandra.spark.common.MD5Hash;
+import org.apache.cassandra.spark.common.client.ClientException;
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+
+/**
+ * A {@link DataTransferApi} implementation that interacts with Cassandra Sidecar
+ */
+public class SidecarDataTransferApi implements DataTransferApi
+{
+ private static final long serialVersionUID = 2563347232666882754L;
+ private static final Logger LOGGER = LoggerFactory.getLogger(SidecarDataTransferApi.class);
+ private static final String SSTABLE_NAME_SEPARATOR = "-";
+ private static final int SSTABLE_GENERATION_REVERSE_OFFSET = 3;
+
+ private final transient SidecarClient sidecarClient;
+ private final JobInfo job;
+ private final BulkSparkConf conf;
+
+ public SidecarDataTransferApi(SidecarClient sidecarClient, JobInfo job, BulkSparkConf conf)
+ {
+ this.sidecarClient = sidecarClient;
+ this.job = job;
+ this.conf = conf;
+ }
+
+ @Override
+ public void uploadSSTableComponent(Path componentFile,
+ int ssTableIdx,
+ CassandraInstance instance,
+ String sessionID,
+ MD5Hash fileHash) throws ClientException
+ {
+ String componentName = updateComponentName(componentFile, ssTableIdx);
+ String uploadId = getUploadId(sessionID, job.getId().toString());
+ try
+ {
+ sidecarClient.uploadSSTableRequest(toSidecarInstance(instance), conf.keyspace, conf.table, uploadId,
+ componentName, fileHash.toString(),
+ componentFile.toAbsolutePath().toString()).get();
+ }
+ catch (ExecutionException | InterruptedException exception)
+ {
+ LOGGER.warn("Failed to upload file={}, keyspace={}, table={}, uploadId={}, componentName={}, instance={}",
+ componentFile, conf.keyspace, conf.table, uploadId, componentName, instance);
+ throw new ClientException(
+ String.format("Failed to upload file=%s into keyspace=%s, table=%s, componentName=%s with uploadId=%s to instance=%s",
+ componentFile, conf.keyspace, conf.table, componentName, uploadId, instance), exception);
+ }
+ }
+
+ @Override
+ public RemoteCommitResult commitSSTables(CassandraInstance instance,
+ String migrationId,
+ List<String> uuids) throws ClientException
+ {
+ if (uuids.size() != 1)
+ {
+ throw new UnsupportedOperationException("Only a single UUID is supported, you provided " + uuids.size());
+ }
+ String uploadId = getUploadId(uuids.get(0), job.getId().toString());
+ ImportSSTableRequest.ImportOptions importOptions = new ImportSSTableRequest.ImportOptions();
+
+ if (job.validateSSTables())
+ {
+ importOptions.verifySSTables(true)
+ .extendedVerify(!job.skipExtendedVerify());
+ }
+
+ try
+ {
+ SSTableImportResponse response =
+ sidecarClient.importSSTableRequest(toSidecarInstance(instance), conf.keyspace, conf.table, uploadId, importOptions).get();
+ if (response.success())
+ {
+ return new RemoteCommitResult(response.success(), Collections.emptyList(), Collections.singletonList(uploadId), null);
+ }
+ return new RemoteCommitResult(response.success(), Collections.singletonList(uploadId), Collections.emptyList(), null);
+ }
+ catch (ExecutionException | InterruptedException exception)
+ {
+ throw new ClientException("Failed to commit sorted string tables", exception);
+ }
+ }
+
+ @Override
+ public void cleanUploadSession(CassandraInstance instance, String sessionID, String jobID) throws ClientException
+ {
+ String uploadId = getUploadId(sessionID, jobID);
+ try
+ {
+ sidecarClient.cleanUploadSession(toSidecarInstance(instance), uploadId).get();
+ }
+ catch (ExecutionException | InterruptedException exception)
+ {
+ LOGGER.warn("Failed to clean upload uploadId={}, instance={}", uploadId, instance);
+ throw new ClientException("Failed to clean the upload session with ID " + uploadId, exception);
+ }
+ }
+
+ protected String updateComponentName(Path componentFile, int ssTableIdx)
+ {
+ String[] ssTableNameParts = componentFile.getFileName().toString().split(SSTABLE_NAME_SEPARATOR);
+ if (ssTableNameParts.length < SSTABLE_GENERATION_REVERSE_OFFSET)
+ {
+ throw new IllegalArgumentException("Invalid component file name: " + componentFile.getFileName());
+ }
+ ssTableNameParts[ssTableNameParts.length - SSTABLE_GENERATION_REVERSE_OFFSET] = Integer.toString(ssTableIdx);
+ return String.join(SSTABLE_NAME_SEPARATOR, ssTableNameParts);
+ }
+
+ protected String getUploadId(String sessionID, String jobId)
+ {
+ return sessionID + "-" + jobId;
+ }
+
+ private SidecarInstanceImpl toSidecarInstance(CassandraInstance instance)
+ {
+ return new SidecarInstanceImpl(instance.getNodeName(), conf.getSidecarPort());
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SqlToCqlTypeConverter.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SqlToCqlTypeConverter.java
new file mode 100644
index 0000000..03772bd
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/SqlToCqlTypeConverter.java
@@ -0,0 +1,726 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import java.net.InetAddress;
+import java.nio.ByteBuffer;
+import java.sql.Timestamp;
+import java.util.ArrayList;
+import java.util.Date;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+
+import com.google.common.base.Preconditions;
+import com.google.common.net.InetAddresses;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.data.CqlField;
+import org.apache.cassandra.spark.utils.UUIDs;
+import scala.Tuple2;
+
+import static org.apache.cassandra.spark.utils.ScalaConversionUtils.asJavaIterable;
+
+@SuppressWarnings("unchecked")
+public final class SqlToCqlTypeConverter implements Serializable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(SqlToCqlTypeConverter.class);
+
+ public static final String ASCII = "ascii";
+ public static final String BIGINT = "bigint";
+ public static final String BLOB = "blob";
+ public static final String BOOLEAN = "boolean";
+ public static final String COUNTER = "counter";
+ public static final String CUSTOM = "custom";
+ public static final String DATE = "date";
+ public static final String DECIMAL = "decimal";
+ public static final String DOUBLE = "double";
+ public static final String FLOAT = "float";
+ public static final String FROZEN = "frozen";
+ public static final String INET = "inet";
+ public static final String INT = "int";
+ public static final String LIST = "list";
+ public static final String MAP = "map";
+ public static final String SET = "set";
+ public static final String SMALLINT = "smallint";
+ public static final String TEXT = "text";
+ public static final String TIME = "time";
+ public static final String UUID = "uuid";
+ public static final String TIMESTAMP = "timestamp";
+ public static final String TIMEUUID = "timeuuid";
+ public static final String TINYINT = "tinyint";
+ public static final String TUPLE = "tuple";
+ public static final String UDT = "udt";
+ public static final String VARCHAR = "varchar";
+ public static final String VARINT = "varint";
+
+ private static final NoOp<Object> NO_OP_CONVERTER = new NoOp<>();
+ private static final LongConverter LONG_CONVERTER = new LongConverter();
+ private static final BytesConverter BYTES_CONVERTER = new BytesConverter();
+ private static final BigDecimalConverter BIG_DECIMAL_CONVERTER = new BigDecimalConverter();
+ private static final IntegerConverter INTEGER_CONVERTER = new IntegerConverter();
+ private static final TimestampConverter TIMESTAMP_CONVERTER = new TimestampConverter();
+ private static final TimeConverter TIME_CONVERTER = new TimeConverter();
+ private static final UUIDConverter UUID_CONVERTER = new UUIDConverter();
+ private static final BigIntegerConverter BIG_INTEGER_CONVERTER = new BigIntegerConverter();
+ private static final TimeUUIDConverter TIME_UUID_CONVERTER = new TimeUUIDConverter();
+ private static final InetAddressConverter INET_ADDRESS_CONVERTER = new InetAddressConverter();
+ private static final DateConverter DATE_CONVERTER = new DateConverter();
+
+ private SqlToCqlTypeConverter()
+ {
+ throw new IllegalStateException(getClass() + " is static utility class and shall not be instantiated");
+ }
+
+ /**
+ * Method to get appropriate converter for a given CQL datatype
+ *
+ * @param cqlType the Cassandra data type from which a converter is needed
+ * @return a type converter that knows how to handle the appropriate type
+ */
+ public static Converter<?> getConverter(CqlField.CqlType cqlType)
+ {
+ String cqlName = cqlType.name().toLowerCase();
+ switch (cqlName)
+ {
+ case CUSTOM:
+ return determineCustomConvert((CqlField.CqlCustom) cqlType);
+ case ASCII:
+ return NO_OP_CONVERTER;
+ case BIGINT:
+ return LONG_CONVERTER;
+ case BLOB:
+ return BYTES_CONVERTER;
+ case BOOLEAN:
+ return NO_OP_CONVERTER;
+ case COUNTER:
+ return NO_OP_CONVERTER;
+ case DECIMAL:
+ return BIG_DECIMAL_CONVERTER;
+ case DOUBLE:
+ return NO_OP_CONVERTER;
+ case FLOAT:
+ return NO_OP_CONVERTER;
+ case FROZEN:
+ assert cqlType instanceof CqlField.CqlFrozen;
+ CqlField.CqlFrozen frozen = (CqlField.CqlFrozen) cqlType;
+ return getConverter(frozen.inner());
+ case INT:
+ return INTEGER_CONVERTER;
+ case TEXT:
+ return NO_OP_CONVERTER;
+ case TIMESTAMP:
+ return TIMESTAMP_CONVERTER;
+ case TIME:
+ return TIME_CONVERTER;
+ case UUID:
+ return UUID_CONVERTER;
+ case VARCHAR:
+ return NO_OP_CONVERTER;
+ case VARINT:
+ return BIG_INTEGER_CONVERTER;
+ case TIMEUUID:
+ return TIME_UUID_CONVERTER;
+ case INET:
+ return INET_ADDRESS_CONVERTER;
+ case DATE:
+ return DATE_CONVERTER;
+ case SMALLINT:
+ return NO_OP_CONVERTER;
+ case TINYINT:
+ return NO_OP_CONVERTER;
+ case LIST:
+ return new ListConverter<>((CqlField.CqlCollection) cqlType);
+ case MAP:
+ assert cqlType instanceof CqlField.CqlMap;
+ return new MapConverter<>((CqlField.CqlMap) cqlType);
+ case SET:
+ return new SetConverter<>((CqlField.CqlCollection) cqlType);
+ case UDT:
+ return NO_OP_CONVERTER;
+ case TUPLE:
+ return NO_OP_CONVERTER;
+ default:
+ LOGGER.warn("Unable to match type={}. Defaulting to NoOp Converter", cqlName);
+ return NO_OP_CONVERTER;
+ }
+ }
+
+ private static Converter<?> determineCustomConvert(CqlField.CqlCustom customType)
+ {
+ Preconditions.checkArgument(customType.name().equalsIgnoreCase(CUSTOM), "Non-custom types are not supported");
+ if (customType.customTypeClassName().equalsIgnoreCase("org.apache.cassandra.db.marshal.DateType"))
+ {
+ return TIMESTAMP_CONVERTER;
+ }
+ else
+ {
+ return NO_OP_CONVERTER;
+ }
+ }
+
+ abstract static class Converter<T> implements Serializable
+ {
+ public abstract T convertInternal(Object object) throws RuntimeException;
+
+ public T convert(Object object)
+ {
+ return convertInternal(object);
+ }
+ }
+
+ private abstract static class NullableConverter<T> extends Converter<T>
+ {
+ @Override
+ public T convert(Object object)
+ {
+ return object != null ? super.convert(object) : null;
+ }
+ }
+
+ static class NoOp<T> extends NullableConverter<T>
+ {
+ @Override
+ public T convertInternal(Object object)
+ {
+ return (T) object;
+ }
+
+ @Override
+ public String toString()
+ {
+ return "NoOp";
+ }
+ }
+
+ static class LongConverter extends NullableConverter<Long>
+ {
+ @Override
+ public Long convertInternal(Object object)
+ {
+ if (object instanceof Long
+ || object instanceof Integer
+ || object instanceof Short
+ || object instanceof Byte)
+ {
+ return ((Number) object).longValue();
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for LONG from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Long";
+ }
+ }
+
+ static class BytesConverter extends NullableConverter<ByteBuffer>
+ {
+ @Override
+ public ByteBuffer convertInternal(Object object)
+ {
+ if (object instanceof ByteBuffer)
+ {
+ return (ByteBuffer) object;
+ }
+ else if (object instanceof byte[])
+ {
+ return ByteBuffer.wrap((byte[]) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for BYTES from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Bytes";
+ }
+ }
+
+ static class BigDecimalConverter extends NullableConverter<BigDecimal>
+ {
+ @Override
+ public BigDecimal convertInternal(Object object)
+ {
+ if (object instanceof BigDecimal)
+ {
+ return (BigDecimal) object;
+ }
+ else if (object instanceof scala.math.BigDecimal)
+ {
+ return ((scala.math.BigDecimal) object).bigDecimal();
+ }
+ else if (object instanceof String)
+ {
+ return new BigDecimal((String) object);
+ }
+ else if (object instanceof Double)
+ {
+ return BigDecimal.valueOf((Double) object);
+ }
+ else if (object instanceof Float)
+ {
+ return BigDecimal.valueOf((Float) object);
+ }
+ else if (object instanceof Long)
+ {
+ return BigDecimal.valueOf((Long) object);
+ }
+ else if (object instanceof Integer)
+ {
+ return BigDecimal.valueOf((Integer) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for DECIMAL from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Decimal";
+ }
+ }
+
+ static class InetAddressConverter extends NullableConverter<InetAddress>
+ {
+ @Override
+ @SuppressWarnings("UnstableApiUsage")
+ public InetAddress convertInternal(Object object)
+ {
+ if (object instanceof InetAddress)
+ {
+ return (InetAddress) object;
+ }
+ else if (object instanceof String)
+ {
+ return InetAddresses.forString((String) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for INET from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Inet";
+ }
+ }
+
+ static class IntegerConverter extends NullableConverter<Integer>
+ {
+ @Override
+ public Integer convertInternal(Object object)
+ {
+ if (object instanceof Integer
+ || object instanceof Short
+ || object instanceof Byte)
+ {
+ return ((Number) object).intValue();
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for INTEGER from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Integer";
+ }
+ }
+
+ static class BigIntegerConverter extends NullableConverter<BigInteger>
+ {
+ @Override
+ public BigInteger convertInternal(Object object)
+ {
+ if (object instanceof BigInteger)
+ {
+ return (BigInteger) object;
+ }
+ else if (object instanceof scala.math.BigInt)
+ {
+ return ((scala.math.BigInt) object).bigInteger();
+ }
+ else if (object instanceof String)
+ {
+ return new BigInteger((String) object);
+ }
+ else if (object instanceof Long)
+ {
+ return BigInteger.valueOf((Long) object);
+ }
+ else if (object instanceof Integer)
+ {
+ return BigInteger.valueOf((Integer) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for VARCHAR from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "VarChar";
+ }
+ }
+
+ static class TimestampConverter extends NullableConverter<Date>
+ {
+ @Override
+ public Date convertInternal(Object object) throws RuntimeException
+ {
+ if (object instanceof Date)
+ {
+ return (Date) object;
+ }
+ else if (object instanceof Long)
+ {
+ return new Date((Long) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for DATE from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Timestamp";
+ }
+ }
+
+ static class DateConverter extends NullableConverter<Integer>
+ {
+ @Override
+ public String toString()
+ {
+ return "Date";
+ }
+
+ protected int fromDate(Date value)
+ {
+ long millisSinceEpoch = value.getTime();
+ return fromMillisSinceEpoch(millisSinceEpoch);
+ }
+
+ protected int fromMillisSinceEpoch(long millisSinceEpoch)
+ {
+ // NOTE: This code is lifted from org.apache.cassandra.serializers.SimpleDateSerializer#timeInMillisToDay.
+ // Reproduced here due to the difficulties of referencing classes from specific versions of Cassandra
+ // in the SBW.
+ int result = (int) TimeUnit.MILLISECONDS.toDays(millisSinceEpoch);
+ result -= Integer.MIN_VALUE;
+ return result;
+ }
+
+ @Override
+ public Integer convertInternal(Object object)
+ {
+ if (object instanceof Date)
+ {
+ return fromDate((Date) object);
+ }
+ else if (object instanceof Long)
+ {
+ return fromMillisSinceEpoch((Long) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for DATE from " + object.getClass().getTypeName());
+ }
+ }
+ }
+
+ static class TimeConverter extends NullableConverter<Long>
+ {
+ @Override
+ @SuppressWarnings("deprecation")
+ public Long convertInternal(Object object)
+ {
+ if (object instanceof Long)
+ {
+ long result = (Long) object;
+ if (result < 0 || result >= TimeUnit.DAYS.toNanos(1))
+ {
+ throw new IllegalArgumentException("Input value out of bounds for Cassandra Time field: " + result);
+ }
+ return (long) object;
+ }
+ else if (object instanceof Timestamp)
+ {
+ // Here, we truncate the date information and just use the time
+ Timestamp timestamp = (Timestamp) object;
+ long rawTime = 0;
+ rawTime += TimeUnit.HOURS.toNanos(timestamp.getHours());
+ rawTime += TimeUnit.MINUTES.toNanos(timestamp.getMinutes());
+ rawTime += TimeUnit.SECONDS.toNanos(timestamp.getSeconds());
+ rawTime += timestamp.getNanos();
+ return rawTime;
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for TIME from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Time";
+ }
+ }
+
+ static class UUIDConverter extends NullableConverter<UUID>
+ {
+ @Override
+ public UUID convertInternal(Object object)
+ {
+ if (object instanceof UUID)
+ {
+ return (UUID) object;
+ }
+ else if (object instanceof String)
+ {
+ return java.util.UUID.fromString((String) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for UUID from " + object.getClass().getTypeName());
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "UUID";
+ }
+ }
+
+ static class TimeUUIDConverter extends NullableConverter<UUID>
+ {
+ @Override
+ public UUID convertInternal(Object object)
+ {
+ UUID result;
+ if (object instanceof UUID)
+ {
+ result = (UUID) object;
+ }
+ else if (object instanceof String)
+ {
+ result = java.util.UUID.fromString((String) object);
+ }
+ else if (object instanceof Long)
+ {
+ result = UUIDs.startOf((Long) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for TIMEUUID from " + object.getClass().getTypeName());
+ }
+ if (result.version() == 1)
+ {
+ return result;
+ }
+ else
+ {
+ throw new RuntimeException("Attempted to convert a non-Timestamp UUID to a TimeUUID - UUID was " + object);
+ }
+ }
+
+ @Override
+ public String toString()
+ {
+ return "TimeUUID";
+ }
+ }
+
+ static class ListConverter<E> extends NullableConverter<List<E>>
+ {
+ private final Converter<?> innerConverter;
+
+ ListConverter(CqlField.CqlCollection cqlType)
+ {
+ innerConverter = getConverter(cqlType.type());
+ }
+
+ @Override
+ public List<E> convertInternal(Object object)
+ {
+ if (object instanceof scala.collection.Iterable)
+ {
+ return makeList(asJavaIterable((scala.collection.Iterable<?>) object));
+ }
+ else if (object instanceof Iterable)
+ {
+ return makeList((Iterable<?>) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for LIST from " + object.getClass().getTypeName());
+ }
+ }
+
+ private List<E> makeList(Iterable<?> iterable)
+ {
+ List<E> list = new ArrayList<>();
+ for (Object object : iterable)
+ {
+ list.add((E) innerConverter.convert(object));
+ }
+ return list;
+ }
+
+ @Override
+ public String toString()
+ {
+ return "List";
+ }
+ }
+
+ static class SetConverter<E> extends NullableConverter<Set<E>>
+ {
+ private final Converter<?> innerConverter;
+
+ SetConverter(CqlField.CqlCollection cqlType)
+ {
+ innerConverter = getConverter(cqlType.type());
+ }
+
+ @Override
+ public Set<E> convertInternal(Object object)
+ {
+ if (object instanceof scala.collection.Iterable)
+ {
+ return makeSet(asJavaIterable((scala.collection.Iterable<?>) object));
+ }
+ else if (object instanceof Iterable)
+ {
+ return makeSet((Iterable<?>) object);
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for SET from " + object.getClass().getTypeName());
+ }
+ }
+
+ private Set<E> makeSet(Iterable<?> iterable)
+ {
+ Set<E> set = new HashSet<>();
+ for (Object object : iterable)
+ {
+ set.add((E) innerConverter.convert(object));
+ }
+ return set;
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Set<" + innerConverter.toString() + ">";
+ }
+ }
+
+ static class MapConverter<K, V> extends NullableConverter<Map<K, V>>
+ {
+ private final Converter<?> keyConverter;
+ private final Converter<?> valConverter;
+
+ MapConverter(CqlField.CqlMap cqlType)
+ {
+ keyConverter = getConverter(cqlType.keyType());
+ valConverter = getConverter(cqlType.valueType());
+ }
+
+ @Override
+ public Map<K, V> convertInternal(Object object)
+ {
+ if (object instanceof scala.collection.Iterable)
+ {
+ return makeMap(asJavaIterable((scala.collection.Iterable<?>) object));
+ }
+ else if (object instanceof Iterable)
+ {
+ return makeMap((Iterable<?>) object);
+ }
+ else if (object instanceof Map)
+ {
+ return makeMap(((Map<K, V>) object).entrySet());
+ }
+ throw new RuntimeException("Unsupported conversion for MAP from " + object.getClass().getTypeName());
+ }
+
+ private Map<K, V> makeMap(Iterable<?> iterable)
+ {
+ Object key;
+ Object value;
+ Map<K, V> map = new HashMap<>();
+ for (Object object : iterable)
+ {
+ if (object instanceof Map.Entry)
+ {
+ key = ((Map.Entry<K, V>) object).getKey();
+ value = ((Map.Entry<K, V>) object).getValue();
+ }
+ else if (object instanceof scala.Tuple2)
+ {
+ key = ((Tuple2<K, V>) object)._1();
+ value = ((Tuple2<K, V>) object)._2();
+ }
+ else
+ {
+ throw new RuntimeException("Unsupported conversion for key/value pair in MAP from " + object.getClass().getTypeName());
+ }
+ map.put((K) keyConverter.convert(key), (V) valConverter.convert(value));
+ }
+ return map;
+ }
+
+ @Override
+ public String toString()
+ {
+ return "Map";
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamError.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamError.java
new file mode 100644
index 0000000..da5383d
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamError.java
@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+
+public class StreamError implements Serializable
+{
+ public final RingInstance instance;
+ public final String errMsg;
+
+ public StreamError(RingInstance instance, String errMsg)
+ {
+ this.instance = instance;
+ this.errMsg = errMsg;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamResult.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamResult.java
new file mode 100644
index 0000000..ee8d02b
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamResult.java
@@ -0,0 +1,59 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.List;
+
+import com.google.common.collect.Range;
+
+public class StreamResult implements Serializable
+{
+ public final String sessionID;
+ public final Range<BigInteger> tokenRange;
+ public final ArrayList<StreamError> failures;
+ public List<CommitResult> commitResults; // CHECKSTYLE IGNORE: Public mutable field
+ public final ArrayList<RingInstance> passed;
+
+ public StreamResult(String sessionID,
+ Range<BigInteger> tokenRange,
+ ArrayList<StreamError> failures,
+ ArrayList<RingInstance> passed)
+ {
+ this.sessionID = sessionID;
+ this.tokenRange = tokenRange;
+ this.failures = failures;
+ this.passed = passed;
+ }
+
+ public void setCommitResults(List<CommitResult> commitResult)
+ {
+ this.commitResults = commitResult;
+ }
+
+ @Override
+ public String toString()
+ {
+ return String.format("StreamResult{sessionID='%s', tokenRange=%s, failures=%s, commitResults=%s, passed=%s}",
+ sessionID, tokenRange, failures, commitResults, passed);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamSession.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamSession.java
new file mode 100644
index 0000000..cd6a51d
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/StreamSession.java
@@ -0,0 +1,300 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.File;
+import java.io.IOException;
+import java.math.BigInteger;
+import java.nio.file.DirectoryStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Range;
+import org.apache.commons.io.FileUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.bulkwriter.token.CassandraRing;
+import org.apache.cassandra.spark.bulkwriter.token.ReplicaAwareFailureHandler;
+import org.apache.cassandra.spark.common.MD5Hash;
+import org.apache.cassandra.spark.common.SSTables;
+
+public class StreamSession
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(StreamSession.class);
+ private final BulkWriterContext writerContext;
+ private final String sessionID;
+ private final Range<BigInteger> tokenRange;
+ final List<RingInstance> replicas;
+ private final ArrayList<StreamError> errors = new ArrayList<>();
+ private final ReplicaAwareFailureHandler<RingInstance> failureHandler;
+ private final AtomicInteger nextSSTableIdx = new AtomicInteger(1);
+ private final ExecutorService executor;
+ private final List<Future<?>> futures = new ArrayList<>();
+ private final CassandraRing<RingInstance> ring;
+ private static final String WRITE_PHASE = "UploadAndCommit";
+
+ public StreamSession(BulkWriterContext writerContext, String sessionID, Range<BigInteger> tokenRange)
+ {
+ this(writerContext, sessionID, tokenRange, Executors.newSingleThreadExecutor());
+ }
+
+ @VisibleForTesting
+ public StreamSession(BulkWriterContext writerContext,
+ String sessionID,
+ Range<BigInteger> tokenRange,
+ ExecutorService executor)
+ {
+ this.writerContext = writerContext;
+ this.ring = writerContext.cluster().getRing(true);
+ this.failureHandler = new ReplicaAwareFailureHandler<>(ring);
+ this.sessionID = sessionID;
+ this.tokenRange = tokenRange;
+ this.replicas = getReplicas();
+ this.executor = executor;
+ }
+
+ public void scheduleStream(SSTableWriter ssTableWriter)
+ {
+ Preconditions.checkState(!ssTableWriter.getTokenRange().isEmpty(), "Trying to stream empty SSTable");
+
+ Preconditions.checkState(tokenRange.encloses(ssTableWriter.getTokenRange()),
+ String.format("SSTable range %s should be enclosed in the partition range %s",
+ ssTableWriter.getTokenRange(), tokenRange));
+
+ futures.add(executor.submit(() -> sendSSTables(writerContext, ssTableWriter)));
+ }
+
+ public StreamResult close() throws ExecutionException, InterruptedException
+ {
+ for (Future future : futures)
+ {
+ try
+ {
+ future.get();
+ }
+ catch (Exception exception)
+ {
+ LOGGER.error("Unexpected stream errMsg. "
+ + "Stream errors should have converted to StreamError and sent to driver", exception);
+ throw new RuntimeException(exception);
+ }
+ }
+
+ executor.shutdown();
+ LOGGER.info("[{}]: Closing stream session. Sent {} SSTables", sessionID, futures.size());
+
+ if (futures.isEmpty())
+ {
+ return new StreamResult(sessionID, tokenRange, new ArrayList<>(), new ArrayList<>());
+ }
+ else
+ {
+ StreamResult streamResult = new StreamResult(sessionID, tokenRange, errors, new ArrayList<>(replicas));
+ List<CommitResult> cr = commit(streamResult);
+ streamResult.setCommitResults(cr);
+ LOGGER.debug("StreamResult: {}", streamResult);
+ BulkWriteValidator.validateClOrFail(failureHandler, LOGGER, WRITE_PHASE, writerContext.job());
+ return streamResult;
+ }
+ }
+
+ private List<CommitResult> commit(StreamResult streamResult) throws ExecutionException, InterruptedException
+ {
+ try (CommitCoordinator cc = CommitCoordinator.commit(writerContext, new StreamResult[]{streamResult}))
+ {
+ List<CommitResult> commitResults = cc.get();
+ LOGGER.debug("All CommitResults: {}", commitResults);
+ commitResults.forEach(cr -> BulkWriteValidator.updateFailureHandler(cr, WRITE_PHASE, failureHandler));
+ return commitResults;
+ }
+ }
+
+ @VisibleForTesting
+ List<RingInstance> getReplicas()
+ {
+ Map<Range<BigInteger>, List<RingInstance>> overlappingRanges = ring.getSubRanges(tokenRange).asMapOfRanges();
+
+ Preconditions.checkState(overlappingRanges.keySet().size() == 1,
+ String.format("Partition range %s is mapping more than one range %s",
+ tokenRange, overlappingRanges));
+
+ List<RingInstance> replicaList = overlappingRanges.values().stream()
+ .flatMap(Collection::stream)
+ .distinct()
+ .collect(Collectors.toList());
+ List<RingInstance> availableReplicas = validateReplicas(replicaList);
+ // In order to better utilize replicas, shuffle the replicaList so each session starts writing to a different replica first
+ Collections.shuffle(availableReplicas);
+ return availableReplicas;
+ }
+
+ private List<RingInstance> validateReplicas(List<RingInstance> replicaList)
+ {
+ Map<Boolean, List<RingInstance>> groups = replicaList.stream()
+ .collect(Collectors.partitioningBy(writerContext.cluster()::instanceIsAvailable));
+ groups.get(false).forEach(instance -> {
+ String errorMessage = String.format("Instance %s is not available.", instance.getNodeName());
+ failureHandler.addFailure(tokenRange, instance, errorMessage);
+ errors.add(new StreamError(instance, errorMessage));
+ });
+ return groups.get(true);
+ }
+
+ private void sendSSTables(BulkWriterContext writerContext, SSTableWriter ssTableWriter)
+ {
+ try (DirectoryStream<Path> dataFileStream = Files.newDirectoryStream(ssTableWriter.getOutDir(), "*Data.db"))
+ {
+ for (Path dataFile : dataFileStream)
+ {
+ int ssTableIdx = nextSSTableIdx.getAndIncrement();
+
+ LOGGER.info("[{}]: Pushing SSTable {} to replicas {}",
+ sessionID, dataFile, replicas.stream()
+ .map(RingInstance::getNodeName)
+ .collect(Collectors.joining(",")));
+ replicas.removeIf(replica -> !trySendSSTableToReplica(writerContext, ssTableWriter, dataFile, ssTableIdx, replica));
+ }
+ }
+ catch (IOException exception)
+ {
+ LOGGER.error("[{}]: Unexpected exception while streaming SSTables {}",
+ sessionID, ssTableWriter.getOutDir());
+ cleanAllReplicas();
+ throw new RuntimeException(exception);
+ }
+ finally
+ {
+ // Clean up SSTable files once the task is complete
+ File tempDir = ssTableWriter.getOutDir().toFile();
+ LOGGER.info("[{}]:Removing temporary files after stream session from {}", sessionID, tempDir);
+ try
+ {
+ FileUtils.deleteDirectory(tempDir);
+ }
+ catch (IOException exception)
+ {
+ LOGGER.warn("[{}]:Failed to delete temporary directory {}", sessionID, tempDir, exception);
+ }
+ }
+ }
+
+ private boolean trySendSSTableToReplica(BulkWriterContext writerContext,
+ SSTableWriter ssTableWriter,
+ Path dataFile,
+ int ssTableIdx,
+ RingInstance replica)
+ {
+ try
+ {
+ sendSSTableToReplica(writerContext, dataFile, ssTableIdx, replica, ssTableWriter.getFileHashes());
+ return true;
+ }
+ catch (Exception exception)
+ {
+ LOGGER.error("[{}]: Failed to stream range {} to instance {}",
+ sessionID, tokenRange, replica.getNodeName(), exception);
+ writerContext.cluster().refreshClusterInfo();
+ failureHandler.addFailure(tokenRange, replica, exception.getMessage());
+ errors.add(new StreamError(replica, exception.getMessage()));
+ clean(writerContext, replica, sessionID);
+ return false;
+ }
+ }
+
+ /**
+ * Get all replicas and clean temporary state on them
+ */
+ private void cleanAllReplicas()
+ {
+ Set<RingInstance> instances = new HashSet<>(replicas);
+ errors.forEach(streamError -> instances.add(streamError.instance));
+ instances.forEach(instance -> clean(writerContext, instance, sessionID));
+ }
+
+ private void sendSSTableToReplica(BulkWriterContext writerContext,
+ Path dataFile,
+ int ssTableIdx,
+ RingInstance instance,
+ Map<Path, MD5Hash> fileHashes) throws Exception
+ {
+ try (DirectoryStream<Path> componentFileStream =
+ Files.newDirectoryStream(dataFile.getParent(), SSTables.getSSTableBaseName(dataFile) + "*"))
+ {
+ for (Path componentFile : componentFileStream)
+ {
+ if (componentFile.getFileName().toString().endsWith("Data.db"))
+ {
+ continue;
+ }
+ sendSSTableComponent(writerContext, componentFile, ssTableIdx, instance, fileHashes.get(componentFile));
+ }
+ sendSSTableComponent(writerContext, dataFile, ssTableIdx, instance, fileHashes.get(dataFile));
+ }
+ }
+
+ private void sendSSTableComponent(BulkWriterContext writerContext,
+ Path componentFile,
+ int ssTableIdx,
+ RingInstance instance,
+ MD5Hash fileHash) throws Exception
+ {
+ Preconditions.checkNotNull(fileHash, "All files must have a hash. SSTableWriter should have calculated these. This is a bug.");
+ long fileSize = Files.size(componentFile);
+ LOGGER.info("[{}]: Uploading {} to {}: Size is {}", sessionID, componentFile, instance.getNodeName(), fileSize);
+ writerContext.transfer().uploadSSTableComponent(componentFile, ssTableIdx, instance, sessionID, fileHash);
+ }
+
+ public static void clean(BulkWriterContext writerContext, RingInstance instance, String sessionID)
+ {
+ if (writerContext.job().getSkipClean())
+ {
+ LOGGER.info("Skip clean requested - not cleaning SSTable session {} on instance {}",
+ sessionID, instance.getNodeName());
+ return;
+ }
+ String jobID = writerContext.job().getId().toString();
+ LOGGER.info("Cleaning SSTable session {} on instance {}", sessionID, instance.getNodeName());
+ try
+ {
+ writerContext.transfer().cleanUploadSession(instance, sessionID, jobID);
+ }
+ catch (Exception exception)
+ {
+ LOGGER.warn("Failed to clean SSTables on {} for session {} and ignoring errMsg",
+ instance.getNodeName(), sessionID, exception);
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TableInfoProvider.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TableInfoProvider.java
new file mode 100644
index 0000000..aaf1ff5
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TableInfoProvider.java
@@ -0,0 +1,48 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.util.List;
+
+import org.apache.cassandra.spark.common.schema.ColumnType;
+import org.apache.cassandra.spark.data.CqlField;
+
+public interface TableInfoProvider
+{
+ CqlField.CqlType getColumnType(String columnName);
+
+ List<ColumnType<?>> getPartitionKeyTypes();
+
+ boolean columnExists(String columnName);
+
+ List<String> getPartitionKeyColumnNames();
+
+ String getCreateStatement();
+
+ List<String> getPrimaryKeyColumnNames();
+
+ String getName();
+
+ String getKeyspaceName();
+
+ boolean hasSecondaryIndex();
+
+ List<String> getColumnNames();
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TableSchema.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TableSchema.java
new file mode 100644
index 0000000..6e5b0ae
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TableSchema.java
@@ -0,0 +1,234 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.LinkedHashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import com.google.common.base.Preconditions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.common.schema.ColumnType;
+import org.apache.spark.sql.types.StructType;
+
+public class TableSchema implements Serializable
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(TableSchema.class);
+
+ final String createStatement;
+ final String modificationStatement;
+ final List<String> partitionKeyColumns;
+ final List<ColumnType<?>> partitionKeyColumnTypes;
+ final List<SqlToCqlTypeConverter.Converter<?>> converters;
+ private final List<Integer> keyFieldPositions;
+ private final WriteMode writeMode;
+
+ public TableSchema(StructType dfSchema, TableInfoProvider tableInfo, WriteMode writeMode)
+ {
+ this.writeMode = writeMode;
+
+ validateDataFrameCompatibility(dfSchema, tableInfo);
+ validateNoSecondaryIndexes(tableInfo);
+
+ this.createStatement = getCreateStatement(tableInfo);
+ this.modificationStatement = getModificationStatement(dfSchema, tableInfo);
+ this.partitionKeyColumns = getPartitionKeyColumnNames(tableInfo);
+ this.partitionKeyColumnTypes = getPartitionKeyColumnTypes(tableInfo);
+ this.converters = getConverters(dfSchema, tableInfo);
+ LOGGER.info("Converters: {}", converters);
+ this.keyFieldPositions = getKeyFieldPositions(dfSchema, tableInfo.getColumnNames(), getRequiredKeyColumns(tableInfo));
+ }
+
+ private List<String> getRequiredKeyColumns(TableInfoProvider tableInfo)
+ {
+ switch (writeMode)
+ {
+ case INSERT:
+ // Inserts require all primary key columns
+ return tableInfo.getPrimaryKeyColumnNames();
+ case DELETE_PARTITION:
+ // To delete a partition, we only need the partition key columns, not all primary key columns
+ return tableInfo.getPartitionKeyColumnNames();
+ default:
+ throw new UnsupportedOperationException("Unknown WriteMode provided");
+ }
+ }
+
+ public Object[] normalize(Object[] row)
+ {
+ for (int index = 0; index < row.length; index++)
+ {
+ row[index] = converters.get(index).convert(row[index]);
+ }
+ return row;
+ }
+
+ public Object[] getKeyColumns(Object[] allColumns)
+ {
+ Object[] result = new Object[keyFieldPositions.size()];
+ for (int keyFieldPosition = 0; keyFieldPosition < keyFieldPositions.size(); keyFieldPosition++)
+ {
+ Object colVal = allColumns[keyFieldPositions.get(keyFieldPosition)];
+ Preconditions.checkNotNull(colVal, "Found a null primary or composite key column in source data. All key columns must be non-null.");
+ result[keyFieldPosition] = colVal;
+ }
+ return result;
+ }
+
+ private static List<SqlToCqlTypeConverter.Converter<?>> getConverters(StructType dfSchema,
+ TableInfoProvider tableInfo)
+ {
+ return Arrays.stream(dfSchema.fieldNames())
+ .map(tableInfo::getColumnType)
+ .map(SqlToCqlTypeConverter::getConverter)
+ .collect(Collectors.toList());
+ }
+
+ private static List<ColumnType<?>> getPartitionKeyColumnTypes(TableInfoProvider tableInfo)
+ {
+ return tableInfo.getPartitionKeyTypes();
+ }
+
+ private static List<String> getPartitionKeyColumnNames(TableInfoProvider tableInfo)
+ {
+ return tableInfo.getPartitionKeyColumnNames();
+ }
+
+ private static String getCreateStatement(TableInfoProvider tableInfo)
+ {
+ String createStatement = tableInfo.getCreateStatement();
+ LOGGER.info("CQL create statement for the table {}", createStatement);
+ return createStatement;
+ }
+
+ private String getModificationStatement(StructType dfSchema, TableInfoProvider tableInfo)
+ {
+ switch (writeMode)
+ {
+ case INSERT:
+ return getInsertStatement(dfSchema, tableInfo);
+ case DELETE_PARTITION:
+ return getDeleteStatement(dfSchema, tableInfo);
+ default:
+ throw new UnsupportedOperationException("Unknown WriteMode provided");
+ }
+ }
+
+ private static String getInsertStatement(StructType dfSchema, TableInfoProvider tableInfo)
+ {
+ String insertStatement = String.format("INSERT INTO %s.%s (%s) VALUES (%s);",
+ tableInfo.getKeyspaceName(),
+ tableInfo.getName(),
+ String.join(",", dfSchema.fieldNames()),
+ Arrays.stream(dfSchema.fieldNames())
+ .map(field -> "?")
+ .collect(Collectors.joining(",")));
+
+ LOGGER.info("CQL insert statement for the RDD {}", insertStatement);
+ return insertStatement;
+ }
+
+ private String getDeleteStatement(StructType dfSchema, TableInfoProvider tableInfo)
+ {
+ Stream<String> fieldEqualityStatements = Arrays.stream(dfSchema.fieldNames()).map(key -> key + "=?");
+ String deleteStatement = String.format("DELETE FROM %s.%s where %s;",
+ tableInfo.getKeyspaceName(),
+ tableInfo.getName(),
+ fieldEqualityStatements.collect(Collectors.joining(" AND ")));
+
+ LOGGER.info("CQL delete statement for the RDD {}", deleteStatement);
+ return deleteStatement;
+ }
+
+ private void validateDataFrameCompatibility(StructType dfSchema, TableInfoProvider tableInfo)
+ {
+ Set<String> dfFields = new LinkedHashSet<>();
+ Collections.addAll(dfFields, dfSchema.fieldNames());
+
+ validatePrimaryKeyColumnsProvided(tableInfo, dfFields);
+
+ switch (writeMode)
+ {
+ case INSERT:
+ validateDataframeFieldsInTable(tableInfo, dfFields);
+ return;
+ case DELETE_PARTITION:
+ validateOnlyPartitionKeyColumnsInDataframe(tableInfo, dfFields);
+ return;
+ default:
+ LOGGER.warn("Unrecognized write mode {}", writeMode);
+ }
+ }
+
+ private void validateOnlyPartitionKeyColumnsInDataframe(TableInfoProvider tableInfo, Set<String> dfFields)
+ {
+ Set<String> requiredKeyColumns = new LinkedHashSet<>(getRequiredKeyColumns(tableInfo));
+ Preconditions.checkArgument(requiredKeyColumns.equals(dfFields),
+ String.format("Only partition key columns (%s) are supported in the input Dataframe"
+ + " when WRITE_MODE=DELETE_PARTITION but (%s) columns were provided",
+ String.join(",", requiredKeyColumns), String.join(",", dfFields)));
+ }
+
+ private void validatePrimaryKeyColumnsProvided(TableInfoProvider tableInfo, Set<String> dfFields)
+ {
+ // Make sure all primary key columns are provided
+ List<String> requiredKeyColumns = getRequiredKeyColumns(tableInfo);
+ Preconditions.checkArgument(dfFields.containsAll(requiredKeyColumns),
+ "Missing some required key components in DataFrame => " + requiredKeyColumns
+ .stream()
+ .filter(column -> !dfFields.contains(column))
+ .collect(Collectors.joining(",")));
+ }
+
+ private static void validateDataframeFieldsInTable(TableInfoProvider tableInfo, Set<String> dfFields)
+ {
+ // Make sure all fields in DF schema are part of table
+ String unknownFields = dfFields.stream()
+ .filter(columnName -> !tableInfo.columnExists(columnName))
+ .collect(Collectors.joining(","));
+ Preconditions.checkArgument(unknownFields.isEmpty(), "Unknown fields in data frame => " + unknownFields);
+ }
+
+ private static void validateNoSecondaryIndexes(TableInfoProvider tableInfo)
+ {
+ if (tableInfo.hasSecondaryIndex())
+ {
+ throw new RuntimeException("Bulkwriter doesn't support secondary indexes");
+ }
+ }
+
+ private static List<Integer> getKeyFieldPositions(StructType dfSchema,
+ List<String> columnNames,
+ List<String> keyFieldNames)
+ {
+ List<String> dfFieldNames = Arrays.asList(dfSchema.fieldNames());
+ return columnNames.stream()
+ .filter(keyFieldNames::contains)
+ .map(dfFieldNames::indexOf)
+ .collect(Collectors.toList());
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TokenPartitioner.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TokenPartitioner.java
new file mode 100644
index 0000000..06d51f1
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/TokenPartitioner.java
@@ -0,0 +1,243 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.math.BigInteger;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Collectors;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeMap;
+import com.google.common.collect.RangeSet;
+import com.google.common.collect.TreeRangeMap;
+import com.google.common.collect.TreeRangeSet;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import org.apache.cassandra.spark.bulkwriter.token.CassandraRing;
+import org.apache.cassandra.spark.bulkwriter.token.RangeUtils;
+import org.apache.spark.Partitioner;
+
+public class TokenPartitioner extends Partitioner
+{
+ private static final Logger LOGGER = LoggerFactory.getLogger(TokenPartitioner.class);
+
+ private transient int nrPartitions;
+ private transient RangeMap<BigInteger, Integer> partitionMap;
+ private transient Map<Integer, Range<BigInteger>> reversePartitionMap;
+ private final CassandraRing<RingInstance> ring;
+ private final Integer numberSplits;
+
+ public TokenPartitioner(CassandraRing<RingInstance> ring,
+ Integer numberSplits,
+ int defaultParallelism,
+ Integer cores)
+ {
+ this(ring, numberSplits, defaultParallelism, cores, true);
+ }
+
+ @VisibleForTesting
+ public TokenPartitioner(CassandraRing<RingInstance> ring,
+ Integer numberSplits,
+ int defaultParallelism,
+ Integer cores,
+ boolean randomize)
+ {
+ this.ring = ring;
+ this.numberSplits = calculateSplits(ring, numberSplits, defaultParallelism, cores);
+ setupTokenRangeMap(randomize);
+ validate(); // Intentionally not keeping this in readObject(), it is enough to validate in constructor alone
+ LOGGER.info("Partition map " + partitionMap);
+ LOGGER.info("Reverse partition map " + reversePartitionMap);
+ LOGGER.info("Number of partitions {}", nrPartitions);
+ }
+
+ @Override
+ public int numPartitions()
+ {
+ return nrPartitions;
+ }
+
+ @SuppressWarnings("ConstantConditions")
+ @Override
+ public int getPartition(Object key)
+ {
+ DecoratedKey decoratedKey = (DecoratedKey) key;
+
+ return partitionMap.get(decoratedKey.getToken());
+ }
+
+ public int numSplits()
+ {
+ return numberSplits;
+ }
+
+ public Range<BigInteger> getTokenRange(int partitionId)
+ {
+ return reversePartitionMap.get(partitionId);
+ }
+
+ private void setupTokenRangeMap(boolean randomize)
+ {
+ partitionMap = TreeRangeMap.create();
+ reversePartitionMap = new HashMap<>();
+
+ AtomicInteger nextPartitionId = new AtomicInteger(0);
+ List<Range<BigInteger>> subRanges = ring.getRangeMap().asMapOfRanges().keySet().stream()
+ .flatMap(tr -> RangeUtils.split(tr, numberSplits).stream())
+ .collect(Collectors.toList());
+ if (randomize)
+ {
+ // In order to help distribute the upload load more evenly, shuffle the subranges before assigning a partition
+ Collections.shuffle(subRanges);
+ }
+ subRanges.forEach(tr -> {
+ int partitionId = nextPartitionId.getAndIncrement();
+
+ partitionMap.put(tr, partitionId);
+ reversePartitionMap.put(partitionId, tr);
+ });
+
+ this.nrPartitions = nextPartitionId.get();
+ }
+
+ private void validate()
+ {
+ validateMapSizes();
+ validateCompleteRangeCoverage();
+ validateRangesDoNotOverlap();
+ }
+
+ private void validateRangesDoNotOverlap()
+ {
+ List<Range<BigInteger>> sortedRanges = partitionMap.asMapOfRanges().keySet().stream()
+ .sorted(Comparator.comparing(Range::lowerEndpoint))
+ .collect(Collectors.toList());
+ Range<BigInteger> previous = null;
+ for (Range<BigInteger> current : sortedRanges)
+ {
+ if (previous != null)
+ {
+ Preconditions.checkState(!current.isConnected(previous) || current.intersection(previous).isEmpty(),
+ String.format("Two ranges in partition map are overlapping %s %s", previous, current));
+ }
+
+ previous = current;
+ }
+ }
+
+ private void validateCompleteRangeCoverage()
+ {
+ RangeSet<BigInteger> missingRangeSet = TreeRangeSet.create();
+ missingRangeSet.add(Range.closed(ring.getPartitioner().minToken(),
+ ring.getPartitioner().maxToken()));
+
+ partitionMap.asMapOfRanges().keySet().forEach(missingRangeSet::remove);
+
+ List<Range<BigInteger>> missingRanges = missingRangeSet.asRanges().stream()
+ .filter(Range::isEmpty)
+ .collect(Collectors.toList());
+ // noinspection unchecked
+ Preconditions.checkState(missingRanges.isEmpty(),
+ "There should be no missing ranges, but found " + missingRanges.toString());
+ }
+
+ private void validateMapSizes()
+ {
+ Preconditions.checkState(nrPartitions == partitionMap.asMapOfRanges().keySet().size(),
+ String.format("Number of partitions %d not matching with partition map size %d",
+ nrPartitions, partitionMap.asMapOfRanges().keySet().size()));
+ Preconditions.checkState(nrPartitions == reversePartitionMap.keySet().size(),
+ String.format("Number of partitions %d not matching with reverse partition map size %d",
+ nrPartitions, reversePartitionMap.keySet().size()));
+ Preconditions.checkState(nrPartitions >= ring.getRangeMap().asMapOfRanges().keySet().size(),
+ String.format("Number of partitions %d supposed to be more than number of token ranges %d",
+ nrPartitions, ring.getRangeMap().asMapOfRanges().keySet().size()));
+ Preconditions.checkState(nrPartitions >= ring.getTokenRanges().keySet().size(),
+ String.format("Number of partitions %d supposed to be more than number of instances %d",
+ nrPartitions, ring.getTokenRanges().keySet().size()));
+ Preconditions.checkState(partitionMap.asMapOfRanges().keySet().size() == reversePartitionMap.keySet().size(),
+ String.format("You must be kidding me! Partition map %d and reverse map %d are not of same size",
+ partitionMap.asMapOfRanges().keySet().size(), reversePartitionMap.keySet().size()));
+ }
+
+ private void writeObject(ObjectOutputStream out) throws IOException
+ {
+ out.defaultWriteObject();
+ HashMap<Range<BigInteger>, Integer> partitionEntires = new HashMap<>();
+ partitionMap.asMapOfRanges().forEach(partitionEntires::put);
+ out.writeObject(partitionEntires);
+ }
+
+ @SuppressWarnings("unchecked")
+ private void readObject(ObjectInputStream in) throws ClassNotFoundException, IOException
+ {
+ in.defaultReadObject();
+ HashMap<Range<BigInteger>, Integer> partitionEntires = (HashMap<Range<BigInteger>, Integer>) in.readObject();
+ partitionMap = TreeRangeMap.create();
+ reversePartitionMap = new HashMap<>();
+ partitionEntires.forEach((r, i) -> {
+ partitionMap.put(r, i);
+ reversePartitionMap.put(i, r);
+ nrPartitions++;
+ });
+ LOGGER.info("Partition map " + partitionMap);
+ LOGGER.info("Reverse partition map " + reversePartitionMap);
+ LOGGER.info("Number of partitions {}", nrPartitions);
+ }
+
+ // In order to best utilize the number of Spark cores while minimizing the number of commit calls,
+ // we calculate the number of splits that will just match or exceed the total number of available Spark cores.
+ // NOTE: The actual number of partitions that result from this should always be at least
+ // the number of token ranges times the number of splits, but can be slightly more.
+ public int calculateSplits(CassandraRing<RingInstance> ring,
+ Integer numberSplits,
+ int defaultParallelism,
+ Integer cores)
+ {
+ if (numberSplits >= 0)
+ {
+ return numberSplits;
+ }
+ int tasksToRun = Math.max(cores, defaultParallelism);
+ Map<Range<BigInteger>, List<RingInstance>> rangeListMap = ring.getRangeMap().asMapOfRanges();
+ LOGGER.info("Initial ranges: {}", rangeListMap);
+ int ranges = rangeListMap.size();
+ LOGGER.info("Number of ranges: {}", ranges);
+ int calculatedSplits = divCeil(tasksToRun, ranges);
+ LOGGER.info("Calculated number of splits as {}", calculatedSplits);
+ return calculatedSplits;
+ }
+
+ int divCeil(int a, int b)
+ {
+ return (a + b - 1) / b;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/Tokenizer.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/Tokenizer.java
new file mode 100644
index 0000000..cdd20c2
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/Tokenizer.java
@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+
+import org.apache.cassandra.spark.bulkwriter.token.TokenUtils;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+
+public class Tokenizer implements Serializable
+{
+ private final TokenUtils tokenUtils;
+ private final TableSchema tableSchema;
+
+ public Tokenizer(BulkWriterContext writerContext)
+ {
+ this.tableSchema = writerContext.schema().getTableSchema();
+ this.tokenUtils = new TokenUtils(tableSchema.partitionKeyColumns,
+ tableSchema.partitionKeyColumnTypes,
+ writerContext.cluster().getPartitioner() == Partitioner.Murmur3Partitioner);
+ }
+
+ public DecoratedKey getDecoratedKey(Object[] columns)
+ {
+ Object[] keyColumns = tableSchema.getKeyColumns(columns);
+ ByteBuffer key = tokenUtils.getCompositeKey(keyColumns);
+ return new DecoratedKey(tokenUtils.getToken(key), key);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriteMode.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriteMode.java
new file mode 100644
index 0000000..7d488fa
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriteMode.java
@@ -0,0 +1,26 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public enum WriteMode
+{
+ INSERT,
+ DELETE_PARTITION
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOption.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOption.java
new file mode 100644
index 0000000..65bf8e3
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOption.java
@@ -0,0 +1,25 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public interface WriterOption
+{
+ String name();
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOptions.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOptions.java
new file mode 100644
index 0000000..7b48df8
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/WriterOptions.java
@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter;
+
+public enum WriterOptions implements WriterOption
+{
+ SIDECAR_INSTANCES,
+ KEYSPACE,
+ TABLE,
+ BULK_WRITER_CL,
+ LOCAL_DC,
+ NUMBER_SPLITS,
+ BATCH_SIZE,
+ COMMIT_THREADS_PER_INSTANCE,
+ COMMIT_BATCH_SIZE,
+ VALIDATE_SSTABLES,
+ SKIP_EXTENDED_VERIFY,
+ WRITE_MODE,
+ KEYSTORE_PASSWORD,
+ KEYSTORE_PATH,
+ KEYSTORE_BASE64_ENCODED,
+ KEYSTORE_TYPE,
+ TRUSTSTORE_PASSWORD,
+ TRUSTSTORE_TYPE,
+ TRUSTSTORE_PATH,
+ TRUSTSTORE_BASE64_ENCODED,
+ SIDECAR_PORT,
+ ROW_BUFFER_MODE,
+ SSTABLE_DATA_SIZE_IN_MB
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/CassandraRing.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/CassandraRing.java
new file mode 100644
index 0000000..22c0519
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/CassandraRing.java
@@ -0,0 +1,291 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter.token;
+
+import java.io.IOException;
+import java.io.NotSerializableException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.io.Serializable;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeMap;
+import com.google.common.collect.TreeRangeMap;
+import org.apache.commons.lang.builder.HashCodeBuilder;
+
+import org.apache.cassandra.spark.bulkwriter.RingInstance;
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+import org.apache.cassandra.spark.data.ReplicationFactor;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+
+import static org.apache.cassandra.spark.data.ReplicationFactor.ReplicationStrategy;
+
+/**
+ * CassandraRing is designed to have one unique way of handling Cassandra token/topology information across all Cassandra
+ * tooling. This class is made Serializable so it's easy to use it from Hadoop/Spark. As Cassandra token ranges are
+ * dependent on Replication strategy, ring makes sense for a specific keyspace only. It is made to be immutable for the
+ * sake of simplicity.
+ *
+ * Token ranges are calculated assuming Cassandra racks are not being used, but controlled by assigning tokens properly.
+ * This is the case for now with managed clusters. We should re-think about this if it changes in future.
+ *
+ * {@link #equals(Object)} and {@link #hashCode()} doesn't take {@link #replicas} and {@link #tokenRangeMap} into
+ * consideration as they are just derived fields.
+ *
+ * {@link CassandraInstance} doesn't extend {@link Serializable}. So, it is possible to use {@link CassandraRing} with
+ * non serializable {@link CassandraInstance}. If we try to serialize ring with those instances, you will see
+ * {@link NotSerializableException}. One of the serializable implementation of {@link CassandraInstance} is at
+ * {@link RingInstance}.
+ *
+ * @param <Instance> ring instance type
+ */
+public class CassandraRing<Instance extends CassandraInstance> implements Serializable
+{
+ private final Partitioner partitioner;
+ private final String keyspace;
+ private transient ReplicationFactor replicationFactor;
+ private transient ArrayList<Instance> instances;
+
+ private transient RangeMap<BigInteger, List<Instance>> replicas;
+ private transient Multimap<Instance, Range<BigInteger>> tokenRangeMap;
+
+ /**
+ * Add a replica with given range to replicaMap (RangeMap pointing to replicas).
+ *
+ * replicaMap starts with full range (representing complete ring) with empty list of replicas. So, it is
+ * guaranteed that range will match one or many ranges in replicaMap.
+ *
+ * Scheme to add a new replica for a range:
+ * * Find overlapping rangeMap entries from replicaMap
+ * * For each overlapping range, create new replica list by adding new replica to the existing list and add it
+ * back to replicaMap.
+ */
+ private static <Instance extends CassandraInstance> void addReplica(Instance replica,
+ Range<BigInteger> range,
+ RangeMap<BigInteger, List<Instance>> replicaMap)
+ {
+ Preconditions.checkArgument(range.lowerEndpoint().compareTo(range.upperEndpoint()) <= 0,
+ "Range calculations assume range is not wrapped");
+
+ RangeMap<BigInteger, List<Instance>> replicaRanges = replicaMap.subRangeMap(range);
+ RangeMap<BigInteger, List<Instance>> mappingsToAdd = TreeRangeMap.create();
+
+ replicaRanges.asMapOfRanges().forEach((key, value) -> {
+ List<Instance> replicas = new ArrayList<>(value);
+ replicas.add(replica);
+ mappingsToAdd.put(key, replicas);
+ });
+ replicaMap.putAll(mappingsToAdd);
+ }
+
+ /**
+ * This is the only way of calculating {@link #replicas} and {@link #tokenRangeMap} either it is from Constructor or
+ * from {@link #readObject(ObjectInputStream)}
+ */
+ private void setupTokenRangeMap()
+ {
+ replicas = TreeRangeMap.create();
+ tokenRangeMap = ArrayListMultimap.create();
+
+ // Calculate instance to token ranges mapping
+ if (replicationFactor.getReplicationStrategy() == ReplicationStrategy.SimpleStrategy)
+ {
+ tokenRangeMap.putAll(RangeUtils.calculateTokenRanges(instances,
+ replicationFactor.getTotalReplicationFactor(),
+ partitioner));
+ }
+ else if (replicationFactor.getReplicationStrategy() == ReplicationStrategy.NetworkTopologyStrategy)
+ {
+ for (String dataCenter : getDataCenters())
+ {
+ int rf = replicationFactor.getOptions().get(dataCenter);
+ if (rf == 0)
+ {
+ // Apparently, it is valid to have zero replication factor in Cassandra
+ continue;
+ }
+ List<Instance> dcInstances = instances.stream()
+ .filter(instance -> instance.getDataCenter().matches(dataCenter))
+ .collect(Collectors.toList());
+ tokenRangeMap.putAll(RangeUtils.calculateTokenRanges(dcInstances,
+ replicationFactor.getOptions().get(dataCenter),
+ partitioner));
+ }
+ }
+ else
+ {
+ throw new UnsupportedOperationException("Unsupported replication strategy");
+ }
+
+ // Calculate token range to replica mapping
+ replicas.put(Range.closed(partitioner.minToken(), partitioner.maxToken()), Collections.emptyList());
+ tokenRangeMap.asMap().forEach((instance, ranges) -> ranges.forEach(range -> addReplica(instance, range, replicas)));
+ }
+
+ public CassandraRing(Partitioner partitioner,
+ String keyspace,
+ ReplicationFactor replicationFactor,
+ Collection<Instance> instances)
+ {
+ this.partitioner = partitioner;
+ this.keyspace = keyspace;
+ this.replicationFactor = replicationFactor;
+ this.instances = instances.stream()
+ .sorted(Comparator.comparing(instance -> new BigInteger(instance.getToken())))
+ .collect(Collectors.toCollection(ArrayList::new));
+
+ this.setupTokenRangeMap();
+ }
+
+ public Partitioner getPartitioner()
+ {
+ return partitioner;
+ }
+
+ public String getKeyspace()
+ {
+ return keyspace;
+ }
+
+ public ReplicationFactor getReplicationFactor()
+ {
+ return replicationFactor;
+ }
+
+ public Collection<Instance> getInstances()
+ {
+ return instances;
+ }
+
+ public RangeMap<BigInteger, List<Instance>> getRangeMap()
+ {
+ return replicas;
+ }
+
+ public RangeMap<BigInteger, List<Instance>> getSubRanges(Range<BigInteger> tokenRange)
+ {
+ return replicas.subRangeMap(tokenRange);
+ }
+
+ public Collection<BigInteger> getTokens()
+ {
+ return instances.stream()
+ .map(CassandraInstance::getToken)
+ .map(BigInteger::new)
+ .sorted()
+ .collect(Collectors.toList());
+ }
+
+ public Collection<Range<BigInteger>> getTokenRanges(Instance instance)
+ {
+ return tokenRangeMap.get(instance);
+ }
+
+ public Multimap<Instance, Range<BigInteger>> getTokenRanges()
+ {
+ return tokenRangeMap;
+ }
+
+ public Collection<String> getDataCenters()
+ {
+ return replicationFactor.getReplicationStrategy() == ReplicationStrategy.SimpleStrategy
+ ? Collections.emptySet()
+ : replicationFactor.getOptions().keySet();
+ }
+
+ @Override
+ public boolean equals(Object other)
+ {
+ if (this == other)
+ {
+ return true;
+ }
+ if (other == null || this.getClass() != other.getClass())
+ {
+ return false;
+ }
+
+ CassandraRing<?> that = (CassandraRing<?>) other;
+ return this.partitioner == that.partitioner
+ && this.keyspace.equals(that.keyspace)
+ && this.replicationFactor.getReplicationStrategy() == that.replicationFactor.getReplicationStrategy()
+ && this.replicationFactor.getOptions().equals(that.replicationFactor.getOptions())
+ && this.instances.equals(that.instances);
+ }
+
+ @Override
+ public int hashCode()
+ {
+ return new HashCodeBuilder()
+ .append(partitioner)
+ .append(keyspace)
+ .append(replicationFactor)
+ .append(instances)
+ .hashCode();
+ }
+
+ @SuppressWarnings("unchecked")
+ private void readObject(ObjectInputStream in) throws ClassNotFoundException, IOException
+ {
+ in.defaultReadObject();
+
+ // Get replication factor from options - HashMap
+ replicationFactor = new ReplicationFactor((HashMap<String, String>) in.readObject());
+
+ // Read list of SerializableCasInstance
+ instances = (ArrayList<Instance>) in.readObject();
+
+ setupTokenRangeMap();
+ }
+
+ @SuppressWarnings("unchecked")
+ private void writeObject(ObjectOutputStream out) throws IOException
+ {
+ out.defaultWriteObject();
+
+ // ReplicationFactor can be built from options. Build and serialize options
+ HashMap<String, String> replicationOptions = new HashMap<>();
+ replicationOptions.put("class", replicationFactor.getReplicationStrategy() == ReplicationStrategy.SimpleStrategy
+ ? "org.apache.cassandra.locator.SimpleStrategy"
+ : "org.apache.cassandra.locator.NetworkTopologyStrategy");
+ replicationFactor.getOptions().keySet().forEach(option -> replicationOptions.put(option, replicationFactor.getOptions().get(option).toString()));
+ out.writeObject(replicationOptions);
+
+ if (!instances.isEmpty() && !(instances.get(0) instanceof Serializable))
+ {
+ throw new NotSerializableException(String.format("Cassandra instances(%s) are not Serializable",
+ instances.get(0).getClass()));
+ }
+
+ // Write instances
+ out.writeObject(instances);
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/ConsistencyLevel.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/ConsistencyLevel.java
new file mode 100644
index 0000000..9153cbd
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/ConsistencyLevel.java
@@ -0,0 +1,175 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter.token;
+
+import java.util.Collection;
+
+import com.google.common.base.Preconditions;
+
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+import org.apache.cassandra.spark.data.ReplicationFactor;
+
+public interface ConsistencyLevel
+{
+ boolean isLocal();
+
+ boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts, ReplicationFactor replicationFactor, String localDC);
+
+ enum CL implements ConsistencyLevel
+ {
+ ALL
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return false;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ return failedInsts.isEmpty();
+ }
+ },
+ EACH_QUORUM
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return false;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ Preconditions.checkArgument(replicationFactor.getReplicationStrategy() != ReplicationFactor.ReplicationStrategy.SimpleStrategy,
+ "EACH_QUORUM doesn't make sense for SimpleStrategy keyspaces");
+
+ for (String datacenter : replicationFactor.getOptions().keySet())
+ {
+ int rf = replicationFactor.getOptions().get(datacenter);
+ if (failedInsts.stream()
+ .filter(instance -> instance.getDataCenter().matches(datacenter))
+ .count() > (rf - (rf / 2 + 1)))
+ {
+ return false;
+ }
+ }
+
+ return true;
+ }
+ },
+ QUORUM
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return false;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ int rf = replicationFactor.getTotalReplicationFactor();
+ return failedInsts.size() <= (rf - (rf / 2 + 1));
+ }
+ },
+ LOCAL_QUORUM
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return true;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ Preconditions.checkArgument(replicationFactor.getReplicationStrategy() != ReplicationFactor.ReplicationStrategy.SimpleStrategy,
+ "LOCAL_QUORUM doesn't make sense for SimpleStrategy keyspaces");
+
+ int rf = replicationFactor.getOptions().get(localDC);
+ return failedInsts.stream().filter(instance -> instance.getDataCenter().matches(localDC)).count() <= (rf - (rf / 2 + 1));
+ }
+ },
+ ONE
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return false;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ int rf = replicationFactor.getTotalReplicationFactor();
+ return failedInsts.size() <= rf - 1;
+ }
+ },
+ TWO
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return false;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ int rf = replicationFactor.getTotalReplicationFactor();
+ return failedInsts.size() <= rf - 2;
+ }
+ },
+ LOCAL_ONE
+ {
+ @Override
+ public boolean isLocal()
+ {
+ return true;
+ }
+
+ @Override
+ public boolean checkConsistency(Collection<? extends CassandraInstance> failedInsts,
+ ReplicationFactor replicationFactor,
+ String localDC)
+ {
+ Preconditions.checkArgument(replicationFactor.getReplicationStrategy() != ReplicationFactor.ReplicationStrategy.SimpleStrategy,
+ "LOCAL_QUORUM doesn't make sense for SimpleStrategy keyspaces");
+
+ int rf = replicationFactor.getOptions().get(localDC);
+ return failedInsts.stream().filter(instance -> instance.getDataCenter().matches(localDC)).count() <= (rf - 1);
+ }
+ }
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/RangeUtils.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/RangeUtils.java
new file mode 100644
index 0000000..3165aa2
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/RangeUtils.java
@@ -0,0 +1,148 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter.token;
+
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+import com.google.common.base.Preconditions;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.BoundType;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Range;
+
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+
+/**
+ * Common Cassandra range operations on Guava ranges. Assumes ranges are not wrapped around.
+ * It's the responsibility of caller to unwrap ranges. For example, [100, 1] should become
+ * [100, MAX] and [MIN, 1]. MIN and MAX depend on {@link Partitioner}.
+ */
+public final class RangeUtils
+{
+ private RangeUtils()
+ {
+ }
+
+ public static BigInteger sizeOf(Range<BigInteger> range)
+ {
+ Preconditions.checkArgument(range.lowerEndpoint().compareTo(range.upperEndpoint()) <= 0,
+ "RangeUtils assume ranges are not wrap-around");
+
+ if (range.isEmpty())
+ {
+ return BigInteger.ZERO;
+ }
+
+ BigInteger size = range.upperEndpoint().subtract(range.lowerEndpoint()).add(BigInteger.ONE);
+
+ if (range.lowerBoundType() == BoundType.OPEN)
+ {
+ size = size.subtract(BigInteger.ONE);
+ }
+
+ if (range.upperBoundType() == BoundType.OPEN)
+ {
+ size = size.subtract(BigInteger.ONE);
+ }
+
+ return size;
+ }
+
+ /**
+ * Splits the given range into equal sized small ranges. Number of splits can be controlled by
+ * nrSplits. If nrSplits are smaller than size of the range, split size would be set to 1.
+ *
+ * This is best effort scheme, nrSplits not necessarily as promised and not all splits may not be
+ * exact same size.
+ *
+ * @param range the range to split
+ * @param nrSplits the number of sub-ranges into which the range should be divided
+ * @return a list of sub-ranges
+ */
+ public static List<Range<BigInteger>> split(Range<BigInteger> range, int nrSplits)
+ {
+ Preconditions.checkArgument(range.lowerEndpoint().compareTo(range.upperEndpoint()) <= 0,
+ "RangeUtils assume ranges are not wrap-around");
+
+ if (range.isEmpty())
+ {
+ return Collections.emptyList();
+ }
+
+ Preconditions.checkArgument(nrSplits >= 1, "nrSplits must be greater than or equal to 1");
+
+ // Make sure split size is not 0
+ BigInteger splitSize = sizeOf(range).divide(BigInteger.valueOf(nrSplits));
+ if (splitSize.compareTo(BigInteger.ZERO) == 0)
+ {
+ splitSize = BigInteger.ONE;
+ }
+
+ // Start from range lower endpoint and spit ranges of size splitSize, until we cross the range
+ BigInteger nextLowerEndpoint = range.lowerBoundType() == BoundType.CLOSED
+ ? range.lowerEndpoint()
+ : range.lowerEndpoint().add(BigInteger.ONE);
+ List<Range<BigInteger>> splits = new ArrayList<>();
+ while (range.contains(nextLowerEndpoint))
+ {
+ BigInteger upperEndpoint = nextLowerEndpoint.add(splitSize);
+ splits.add(range.intersection(Range.closedOpen(nextLowerEndpoint, upperEndpoint)));
+ nextLowerEndpoint = upperEndpoint;
+ }
+
+ return splits;
+ }
+
+ public static <Instance extends CassandraInstance> Multimap<Instance, Range<BigInteger>> calculateTokenRanges(
+ List<Instance> instances,
+ int replicationFactor,
+ Partitioner partitioner)
+ {
+ Preconditions.checkArgument(replicationFactor != 0, "Calculation token ranges wouldn't work with RF 0");
+ Preconditions.checkArgument(instances.size() == 0 || replicationFactor <= instances.size(),
+ "Calculation token ranges wouldn't work when RF (" + replicationFactor
+ + ") is greater than number of Cassandra instances " + instances.size());
+ Multimap<Instance, Range<BigInteger>> tokenRanges = ArrayListMultimap.create();
+ for (int index = 0; index < instances.size(); index++)
+ {
+ Instance instance = instances.get(index);
+ int disjointReplica = ((instances.size() + index) - replicationFactor) % instances.size();
+ BigInteger rangeStart = new BigInteger(instances.get(disjointReplica).getToken());
+ BigInteger rangeEnd = new BigInteger(instance.getToken());
+
+ // If start token is not strictly smaller than end token we are looking at a wrap around range, split it
+ if (rangeStart.compareTo(rangeEnd) >= 0)
+ {
+ tokenRanges.put(instance, Range.range(rangeStart, BoundType.OPEN, partitioner.maxToken(), BoundType.CLOSED));
+ tokenRanges.put(instance, Range.range(partitioner.minToken(), BoundType.CLOSED, rangeEnd, BoundType.CLOSED));
+ }
+ else
+ {
+ tokenRanges.put(instance, Range.openClosed(rangeStart, rangeEnd));
+ }
+ }
+
+ return tokenRanges;
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/ReplicaAwareFailureHandler.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/ReplicaAwareFailureHandler.java
new file mode 100644
index 0000000..e57c571
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/ReplicaAwareFailureHandler.java
@@ -0,0 +1,134 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter.token;
+
+import java.math.BigInteger;
+import java.util.AbstractMap;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeMap;
+import com.google.common.collect.TreeRangeMap;
+
+import org.apache.cassandra.spark.common.model.CassandraInstance;
+
+public class ReplicaAwareFailureHandler<Instance extends CassandraInstance>
+{
+ private final CassandraRing<Instance> ring;
+ private final RangeMap<BigInteger, Multimap<Instance, String>> failedRangesMap = TreeRangeMap.create();
+
+ public ReplicaAwareFailureHandler(CassandraRing<Instance> ring)
+ {
+ this.ring = ring;
+ failedRangesMap.put(Range.closed(ring.getPartitioner().minToken(),
+ ring.getPartitioner().maxToken()),
+ ArrayListMultimap.create());
+ }
+
+ /**
+ * Adds a new token range as a failed token range, with errors on given instance.
+ *
+ * It's guaranteed that failedRangesMap has overlapping ranges for the range we are trying to insert (Check
+ * constructor, we are adding complete ring first).
+ *
+ * So the scheme is to get list of overlapping ranges first. For each overlapping range get the failure map.
+ * Make a copy of the map and add new failure to this map. It's important we make the copy and not use the
+ * one returned from failedRangesMap map. As our range could be overlapping partially and the map could be used
+ * by other range.
+ *
+ * @param tokenRange the range which failed
+ * @param casInstance the instance on which the range failed
+ * @param errMessage the error that occurred for this particular range/instance pair
+ */
+ public void addFailure(Range<BigInteger> tokenRange, Instance casInstance, String errMessage)
+ {
+ RangeMap<BigInteger, Multimap<Instance, String>> overlappingFailures = failedRangesMap.subRangeMap(tokenRange);
+ RangeMap<BigInteger, Multimap<Instance, String>> mappingsToAdd = TreeRangeMap.create();
+
+ for (Map.Entry<Range<BigInteger>, Multimap<Instance, String>> entry : overlappingFailures.asMapOfRanges().entrySet())
+ {
+ Multimap<Instance, String> newErrorMap = ArrayListMultimap.create(entry.getValue());
+
+ newErrorMap.put(casInstance, errMessage);
+ mappingsToAdd.put(entry.getKey(), newErrorMap);
+ }
+ failedRangesMap.putAll(mappingsToAdd);
+ }
+
+ public boolean hasFailed(ConsistencyLevel consistencyLevel, String localDC)
+ {
+ return !getFailedEntries(consistencyLevel, localDC).isEmpty();
+ }
+
+ @SuppressWarnings("unused") // Convenience method can become useful in the future
+ public Collection<Range<BigInteger>> getFailedRanges(ConsistencyLevel consistencyLevel, String localDC)
+ {
+ return getFailedEntries(consistencyLevel, localDC).stream()
+ .map(AbstractMap.SimpleEntry::getKey)
+ .collect(Collectors.toList());
+ }
+
+ @SuppressWarnings("unused") // Convenience method can become useful in the future
+ public Multimap<Instance, String> getFailedInstances(ConsistencyLevel consistencyLevel, String localDC)
+ {
+ Multimap<Instance, String> failedInstances = ArrayListMultimap.create();
+ getFailedEntries(consistencyLevel, localDC).stream()
+ .flatMap(failedMultiEntry -> failedMultiEntry.getValue().entries().stream())
+ .forEach(failedEntry -> failedInstances.put(failedEntry.getKey(), failedEntry.getValue()));
+
+ return failedInstances;
+ }
+
+ public Collection<AbstractMap.SimpleEntry<Range<BigInteger>, Multimap<Instance, String>>> getFailedEntries(
+ ConsistencyLevel consistencyLevel,
+ String localDC)
+ {
+ List<AbstractMap.SimpleEntry<Range<BigInteger>, Multimap<Instance, String>>> failedEntries = new ArrayList<>();
+
+ for (Map.Entry<Range<BigInteger>, Multimap<Instance, String>> failedRangeEntry : failedRangesMap.asMapOfRanges().entrySet())
+ {
+ Multimap<Instance, String> errorMap = failedRangeEntry.getValue();
+ Collection<Instance> failedInstances = errorMap.keySet().stream()
+ .filter(instance -> !errorMap.get(instance).isEmpty())
+ .collect(Collectors.toList());
+ if (!consistencyLevel.checkConsistency(failedInstances, ring.getReplicationFactor(), localDC))
+ {
+ failedEntries.add(new AbstractMap.SimpleEntry<>(failedRangeEntry.getKey(), failedRangeEntry.getValue()));
+ }
+ }
+
+ return failedEntries;
+ }
+
+ public String toString()
+ {
+ return "CassandraRing: Tokens: " + ring.getTokens().stream()
+ .map(BigInteger::toString)
+ .collect(Collectors.joining(",", "[", "]")) + ", "
+ + "ReplicationFactor: " + ring.getReplicationFactor().getReplicationStrategy().name() + " "
+ + ring.getReplicationFactor().getOptions();
+ }
+}
diff --git a/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/TokenUtils.java b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/TokenUtils.java
new file mode 100644
index 0000000..b0169ae
--- /dev/null
+++ b/cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/bulkwriter/token/TokenUtils.java
@@ -0,0 +1,164 @@
+/*
+ * 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.
+ */
+
+package org.apache.cassandra.spark.bulkwriter.token;
+
+import java.io.Serializable;
+import java.math.BigInteger;
+import java.nio.Buffer;
+import java.nio.ByteBuffer;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+
+import org.apache.cassandra.spark.common.schema.ColumnType;
+import org.apache.cassandra.spark.common.schema.ColumnTypes;
+import org.apache.cassandra.spark.data.partitioner.Partitioner;
+
+/**
+ * This is utility class for computing Cassandra token for a CQL row.
+ * This is mainly written to partition Cassandra ring and split it between Spark executors.
+ * And to make sure each executor is only talking to one Cassandra replica set.
+ * This reduces number of SSTables that get created in Cassandra by the bulk writing job.
+ * Fewer SSTables will result in lower read latencies and lower compaction overhead.
+ */
+@SuppressWarnings("WeakerAccess")
+public class TokenUtils implements Serializable
+{
+ private final String[] partitionKeyColumns;
+ private final boolean isMurmur3Partitioner;
+ private final ColumnType[] partitionKeyColumnTypes;
+
+ public TokenUtils(List<String> partitionKeyColumns,
+ List<ColumnType<?>> partitionKeyColumnTypes,
+ boolean isMurmur3Partitioner)
+ {
+ this.partitionKeyColumns = partitionKeyColumns.toArray(new String[0]);
+ this.partitionKeyColumnTypes = partitionKeyColumnTypes.toArray(new ColumnType[partitionKeyColumns.size()]);
+ this.isMurmur3Partitioner = isMurmur3Partitioner;
+ }
+
+ // noinspection unchecked
+ private ByteBuffer getByteBuffer(Object columnValue, int partitionKeyColumnIdx)
+ {
+ ColumnType columnType = partitionKeyColumnTypes[partitionKeyColumnIdx];
+ if (columnType == ColumnTypes.UUID)
+ {
+ return columnType.serialize(UUID.fromString((String) columnValue));
+ }
+ else
+ {
+ return columnType.serialize(columnValue);
+ }
+ }
+
+ private BigInteger getToken(ByteBuffer[] byteBuffers)
+ {
+ return getToken(getCompositeKey(byteBuffers));
+ }
+
+ public BigInteger getToken(ByteBuffer key)
+ {
+ return isMurmur3Partitioner ? Partitioner.Murmur3Partitioner.hash(key) : Partitioner.RandomPartitioner.hash(key);
+ }
+
+ public ByteBuffer getCompositeKey(Object[] columnValues)
+ {
+ ByteBuffer[] byteBuffers = new ByteBuffer[partitionKeyColumns.length];
+
+ for (int column = 0; column < partitionKeyColumns.length; column++)
+ {
+ byteBuffers[column] = getByteBuffer(columnValues[column], column);
+ }
+
+ return getCompositeKey(byteBuffers);
+ }
+
+ public ByteBuffer getCompositeKey(ByteBuffer[] byteBuffers)
+ {
+ ByteBuffer key;
+ if (byteBuffers.length == 1)
+ {
+ key = byteBuffers[0];
+ }
+ else
+ {
+ // Calculate length of the key
+ int length = 0;
+ for (ByteBuffer buffer : byteBuffers)
+ {
+ length += 2;
+ length += buffer.remaining();
+ length += 1;
+ }
+
+ // Add buffers one after another
+ key = ByteBuffer.allocate(length);
+ for (ByteBuffer buffer : byteBuffers)
+ {
+ key.putShort((short) buffer.remaining());
+ key.put(buffer);
+ key.put((byte) 0x00);
+ }
+ ((Buffer) key).flip();
+ }
+
+ return key;
+ }
+
+ /**
+ * Calculate Cassandra token for CQL row
+ *
+ * @param columnValues Map of partitioner key columns to column value as Java type
+ * @return Cassandra token
+ */
+ public BigInteger getToken(Map<String, Object> columnValues)
+ {
+ ByteBuffer[] byteBuffers = new ByteBuffer[partitionKeyColumns.length];
+
+ for (int column = 0; column < partitionKeyColumns.length; column++)
+ {
+ byteBuffers[column] = getByteBuffer(columnValues.get(partitionKeyColumns[column]), column);
+ }
+
+ return getToken(byteBuffers);
+ }
+
+ /**
... 56393 lines suppressed ...
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@cassandra.apache.org
For additional commands, e-mail: commits-help@cassandra.apache.org