com.nvidia.spark.rapids.GpuPackedTableColumn Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of rapids-4-spark_2.13 Show documentation
Creates the distribution package of the RAPIDS plugin for Apache Spark
The newest version!
/*
 * Copyright (c) 2021, NVIDIA CORPORATION.
 *
 * 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.
 */

package com.nvidia.spark.rapids;

import ai.rapids.cudf.ContiguousTable;
import ai.rapids.cudf.DeviceMemoryBuffer;
import org.apache.spark.sql.vectorized.ColumnVector;
import org.apache.spark.sql.vectorized.ColumnarBatch;

import static org.apache.spark.sql.types.DataTypes.NullType;

/**
 * A GPU column tracking a packed table such as one generated by contiguous split.
 * Unlike {@link GpuColumnVectorFromBuffer}, the columnar data cannot be accessed directly.
 *
 * This class primarily serves the role of tracking the packed table data in a ColumnarBatch
 * without requiring the underlying table to be manifested along with all of the child columns.
 * The typical use-case generates one of these columns per task output partition, and then the
 * RAPIDS shuffle transmits the opaque host metadata and GPU data buffer to another host.
 *
 * NOTE: There should only be one instance of this column per ColumnarBatch as the
 */
public final class GpuPackedTableColumn extends GpuColumnVectorBase implements WithTableBuffer {
  private static final String BAD_ACCESS_MSG = "Column is packed";

  private final ContiguousTable contigTable;

  public static ColumnarBatch from(ContiguousTable contigTable) {
    ColumnVector column = new GpuPackedTableColumn(contigTable);
    return new ColumnarBatch(new ColumnVector[] { column }, (int) contigTable.getRowCount());
  }

  /** Returns true if this columnar batch uses a packed table */
  public static boolean isBatchPacked(ColumnarBatch batch) {
    return batch.numCols() == 1 && batch.column(0) instanceof GpuPackedTableColumn;
  }

  GpuPackedTableColumn(ContiguousTable contigTable) {
    super(NullType);
    this.contigTable = contigTable;
  }

  /**
   * Returns the contiguous table underneath this column.
   * NOTE: The contiguous table is still owned by this column instance, so the
   *       resulting contiguous table should NOT be closed by the caller.
   */
  public ContiguousTable getContiguousTable() {
    return contigTable;
  }

  @Override
  public DeviceMemoryBuffer getTableBuffer() {
    return contigTable.getBuffer();
  }

  @Override
  public void close() {
    contigTable.close();
  }

  @Override
  public boolean hasNull() {
    throw new IllegalStateException(BAD_ACCESS_MSG);
  }

  @Override
  public int numNulls() {
    throw new IllegalStateException(BAD_ACCESS_MSG);
  }
}