Source code for airflow.providers.apache.druid.transfers.hive_to_druid

# 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
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""This module contains operator to move data from Hive to Druid."""
from __future__ import annotations

from typing import TYPE_CHECKING, Any, Sequence

from airflow.models import BaseOperator
from airflow.providers.apache.druid.hooks.druid import DruidHook
from airflow.providers.apache.hive.hooks.hive import HiveCliHook, HiveMetastoreHook

    from airflow.utils.context import Context

[docs]class HiveToDruidOperator(BaseOperator): """ Moves data from Hive to Druid, [del]note that for now the data is loaded into memory before being pushed to Druid, so this operator should be used for smallish amount of data.[/del] :param sql: SQL query to execute against the Druid database. (templated) :param druid_datasource: the datasource you want to ingest into in druid :param ts_dim: the timestamp dimension :param metric_spec: the metrics you want to define for your data :param hive_cli_conn_id: the hive connection id :param druid_ingest_conn_id: the druid ingest connection id :param metastore_conn_id: the metastore connection id :param hadoop_dependency_coordinates: list of coordinates to squeeze int the ingest json :param intervals: list of time intervals that defines segments, this is passed as is to the json object. (templated) :param num_shards: Directly specify the number of shards to create. :param target_partition_size: Target number of rows to include in a partition, :param query_granularity: The minimum granularity to be able to query results at and the granularity of the data inside the segment. E.g. a value of "minute" will mean that data is aggregated at minutely granularity. That is, if there are collisions in the tuple (minute(timestamp), dimensions), then it will aggregate values together using the aggregators instead of storing individual rows. A granularity of 'NONE' means millisecond granularity. :param segment_granularity: The granularity to create time chunks at. Multiple segments can be created per time chunk. For example, with 'DAY' segmentGranularity, the events of the same day fall into the same time chunk which can be optionally further partitioned into multiple segments based on other configurations and input size. :param hive_tblproperties: additional properties for tblproperties in hive for the staging table :param job_properties: additional properties for job """
[docs] template_fields: Sequence[str] = ("sql", "intervals")
[docs] template_ext: Sequence[str] = (".sql",)
[docs] template_fields_renderers = {"sql": "hql"}
def __init__( self, *, sql: str, druid_datasource: str, ts_dim: str, metric_spec: list[Any] | None = None, hive_cli_conn_id: str = "hive_cli_default", druid_ingest_conn_id: str = "druid_ingest_default", metastore_conn_id: str = "metastore_default", hadoop_dependency_coordinates: list[str] | None = None, intervals: list[Any] | None = None, num_shards: float = -1, target_partition_size: int = -1, query_granularity: str = "NONE", segment_granularity: str = "DAY", hive_tblproperties: dict[Any, Any] | None = None, job_properties: dict[Any, Any] | None = None, **kwargs: Any, ) -> None: super().__init__(**kwargs) self.sql = sql self.druid_datasource = druid_datasource self.ts_dim = ts_dim self.intervals = intervals or ["{{ ds }}/{{ tomorrow_ds }}"] self.num_shards = num_shards self.target_partition_size = target_partition_size self.query_granularity = query_granularity self.segment_granularity = segment_granularity self.metric_spec = metric_spec or [{"name": "count", "type": "count"}] self.hive_cli_conn_id = hive_cli_conn_id self.hadoop_dependency_coordinates = hadoop_dependency_coordinates self.druid_ingest_conn_id = druid_ingest_conn_id self.metastore_conn_id = metastore_conn_id self.hive_tblproperties = hive_tblproperties or {} self.job_properties = job_properties
[docs] def execute(self, context: Context) -> None: hive = HiveCliHook(hive_cli_conn_id=self.hive_cli_conn_id)"Extracting data from Hive") hive_table = "druid." + context["task_instance_key_str"].replace(".", "_") sql = self.sql.strip().strip(";") tblproperties = "".join(f", '{k}' = '{v}'" for k, v in self.hive_tblproperties.items()) hql = f"""\ SET mapred.output.compress=false; SET hive.exec.compress.output=false; DROP TABLE IF EXISTS {hive_table}; CREATE TABLE {hive_table} ROW FORMAT DELIMITED FIELDS TERMINATED BY '\t' STORED AS TEXTFILE TBLPROPERTIES ('serialization.null.format' = ''{tblproperties}) AS {sql} """"Running command:\n %s", hql) hive.run_cli(hql) meta_hook = HiveMetastoreHook(self.metastore_conn_id) # Get the Hive table and extract the columns table = meta_hook.get_table(hive_table) columns = [ for col in] # Get the path on hdfs static_path = meta_hook.get_table(hive_table).sd.location druid = DruidHook(druid_ingest_conn_id=self.druid_ingest_conn_id) try: index_spec = self.construct_ingest_query( static_path=static_path, columns=columns, )"Inserting rows into Druid, hdfs path: %s", static_path) druid.submit_indexing_job(index_spec)"Load seems to have succeeded!") finally:"Cleaning up by dropping the temp Hive table %s", hive_table) hql = f"DROP TABLE IF EXISTS {hive_table}" hive.run_cli(hql)
[docs] def construct_ingest_query(self, static_path: str, columns: list[str]) -> dict[str, Any]: """ Builds an ingest query for an HDFS TSV load. :param static_path: The path on hdfs where the data is :param columns: List of all the columns that are available """ # backward compatibility for num_shards, # but target_partition_size is the default setting # and overwrites the num_shards num_shards = self.num_shards target_partition_size = self.target_partition_size if self.target_partition_size == -1: if self.num_shards == -1: target_partition_size = DEFAULT_TARGET_PARTITION_SIZE else: num_shards = -1 metric_names = [m["fieldName"] for m in self.metric_spec if m["type"] != "count"] # Take all the columns, which are not the time dimension # or a metric, as the dimension columns dimensions = [c for c in columns if c not in metric_names and c != self.ts_dim] ingest_query_dict: dict[str, Any] = { "type": "index_hadoop", "spec": { "dataSchema": { "metricsSpec": self.metric_spec, "granularitySpec": { "queryGranularity": self.query_granularity, "intervals": self.intervals, "type": "uniform", "segmentGranularity": self.segment_granularity, }, "parser": { "type": "string", "parseSpec": { "columns": columns, "dimensionsSpec": { "dimensionExclusions": [], "dimensions": dimensions, # list of names "spatialDimensions": [], }, "timestampSpec": {"column": self.ts_dim, "format": "auto"}, "format": "tsv", }, }, "dataSource": self.druid_datasource, }, "tuningConfig": { "type": "hadoop", "jobProperties": { "mapreduce.job.user.classpath.first": "false", "": "false", "mapreduce.output.fileoutputformat.compress": "false", }, "partitionsSpec": { "type": "hashed", "targetPartitionSize": target_partition_size, "numShards": num_shards, }, }, "ioConfig": {"inputSpec": {"paths": static_path, "type": "static"}, "type": "hadoop"}, }, } if self.job_properties: ingest_query_dict["spec"]["tuningConfig"]["jobProperties"].update(self.job_properties) if self.hadoop_dependency_coordinates: ingest_query_dict["hadoopDependencyCoordinates"] = self.hadoop_dependency_coordinates return ingest_query_dict

Was this entry helpful?