diff --git a/bin/rubin/distribute.py b/bin/rubin/distribute.py index b6760561..0f47f1d5 100644 --- a/bin/rubin/distribute.py +++ b/bin/rubin/distribute.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # Copyright 2019-2026 AstroLab Software -# Author: Julien Peloton +# Author: Julien Peloton, Massinissa MACHTER # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,6 +23,7 @@ """ from pyspark.sql.types import BooleanType +import pyspark.sql.functions as F import pkgutil import argparse @@ -44,6 +45,7 @@ expand_function_from_string, FinkUDF, ) + import fink_filters.rubin.livestream as ffrl from fink_broker.rubin.hbase_utils import ingest_section @@ -115,69 +117,100 @@ def main(): logger.warn(msg) spark.stop() - for userfilter in userfilters: - filter_func, colnames = expand_function_from_string(df, userfilter) - - tag = userfilter.split(".")[-1] - fink_filter = FinkUDF( - filter_func, - BooleanType(), - tag, - ) + # Kafka ingestion + if args.noscience: + topics = [] - # Apply or not the filtering - if args.noscience: + for userfilter in userfilters: logger.debug( "Do not apply user-defined filter %s in no-science mode", userfilter ) - df_filtered = df - else: + topicname = args.substream_prefix + userfilter.split(".")[-1] + "_lsst" + topics.append(F.lit(topicname)) + + df_with_topics = df.withColumn("topics", F.array(*topics)) + + df_filtered = df_with_topics.withColumn("topic", F.explode("topics")) + elif args.no_kafka_ingest: + logger.warning("Skipping all Kafka ingestion") + kafka_query = FakeQuery() + else: + topic_exprs = [] + for userfilter in userfilters: logger.debug("Apply user-defined filter %s", userfilter) - df_filtered = df.filter(fink_filter.for_spark(*colnames)) - - # HBase support requires fink-filters>=7.34 - module = userfilter.rsplit(".", maxsplit=1)[0] - hbase_support = importlib.import_module(module).HBASE_SUPPORT - if not args.no_hbase_ingest and hbase_support: - # HBase ingestion - major_version, minor_version = get_schema_from_parquet(scitmpdatapath) - - # Key is time_oid to perform date range search - cols_row_key_name = ["midpointMjdTai", "diaObjectId"] - row_key_name = "_".join(cols_row_key_name) - table_name = "{}.tag_{}".format(args.science_db_name, tag) - - hbase_query = ingest_section( - df_filtered, - major_version, - minor_version, - row_key_name, - table_name=table_name, - catfolder=args.science_db_catalogs, - cols_row_key_name=cols_row_key_name, - streaming=True, - checkpoint_path=checkpointpath_hbase + "/" + tag, + + # build filter function expr dynamically + filter_func, colnames = expand_function_from_string(df, userfilter) + tag = userfilter.split(".")[-1] + fink_filter = FinkUDF( + filter_func, + BooleanType(), + tag, ) - else: - logger.warning("Skipping HBase ingestion for filter {}".format(userfilter)) - hbase_query = FakeQuery() + expr = fink_filter.for_spark(*colnames) - if not args.no_kafka_ingest: - # Kafka distribution topicname = args.substream_prefix + tag + "_lsst" - kafka_query = push_to_kafka( - df_filtered, - topicname, - cnames, - checkpointpath_kafka, - args.tinterval, - kafka_cfg, - npart=10, - ) - else: - logger.warning("Skipping Kafka ingestion for filter {}".format(userfilter)) - kafka_query = FakeQuery() + topic_exprs.append(F.when(expr, F.lit(topicname))) + + # array_compact for delete NULL values in array + df_with_topics = df.withColumn("topics", F.array_compact(F.array(*topic_exprs))) + + df_filtered = df_with_topics.withColumn("topic", F.explode("topics")) + + # All filters distributed to multiple kafka topics with 1 writeStream + kafka_query = push_to_kafka( + df_filtered, + cnames, + checkpointpath_kafka + "/{}filters_lsst".format(args.substream_prefix), + args.tinterval, + kafka_cfg, + npart=10, + ) + + # Hbase ingestion + if not args.no_hbase_ingest: + major_version, minor_version = get_schema_from_parquet(scitmpdatapath) + + # Key is time_oid to perform date range search + cols_row_key_name = ["midpointMjdTai", "diaObjectId"] + row_key_name = "_".join(cols_row_key_name) + + hbase_queries = [] + + # Loop over filters as the Spark-HBase connector does not + # support dynamic routing via column "table" + OneWriteStream + for userfilter in userfilters: + module = userfilter.rsplit(".", maxsplit=1)[0] + hbase_support = importlib.import_module(module).HBASE_SUPPORT + + # Push only to tables with HBase support + if hbase_support: + tag = userfilter.split(".")[-1] + table_name = "{}.tag_{}".format(args.science_db_name, tag) + topicname = args.substream_prefix + tag + "_lsst" + + df_filtered_tag = df_filtered.filter(F.col("topic") == topicname) + + hbase_query = ingest_section( + df_filtered_tag, + major_version, + minor_version, + row_key_name, + table_name=table_name, + catfolder=args.science_db_catalogs, + cols_row_key_name=cols_row_key_name, + streaming=True, + checkpoint_path=checkpointpath_hbase + "/" + tag, + ) + hbase_queries.append(hbase_query) + else: + logger.warning( + "Skipping HBase ingestion for filter {}".format(userfilter) + ) + else: + logger.warning("Skipping all HBase ingestion") + hbase_queries = [FakeQuery()] if args.exit_after is not None: logger.debug("Keep the Streaming running until something or someone ends it!") @@ -185,7 +218,8 @@ def main(): remaining_time = remaining_time if remaining_time > 0 else 0 time.sleep(remaining_time) kafka_query.stop() - hbase_query.stop() + for hbase_query in hbase_queries: + hbase_query.stop() logger.info("Exiting the distribute service normally...") else: logger.debug("Wait for the end of queries") diff --git a/bin/ztf/distribute.py b/bin/ztf/distribute.py index 3888875b..1898c9a7 100644 --- a/bin/ztf/distribute.py +++ b/bin/ztf/distribute.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # Copyright 2019-2025 AstroLab Software -# Author: Abhishek Chauhan, Julien Peloton +# Author: Abhishek Chauhan, Julien Peloton, Massinissa MACHTER # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,7 +37,7 @@ from fink_broker.common.distribution_utils import push_to_kafka from fink_broker.common.logging_utils import init_logger from fink_utils.spark.utils import concat_col -from fink_utils.spark.utils import apply_user_defined_filter +from fink_utils.spark.utils import expand_function_from_string import fink_filters.ztf.livestream as ffzl @@ -164,37 +164,52 @@ def main(): logger.warn(msg) spark.stop() - for userfilter in userfilters: - if args.noscience: + if args.noscience: + topics = [] + for userfilter in userfilters: logger.debug( "Do not apply user-defined filter %s in no-science mode", userfilter ) - df_tmp = df - else: + + topicname = args.substream_prefix + userfilter.split(".")[-1] + "_ztf" + topics.append(F.lit(topicname)) + df_with_topics = df.withColumn("topics", F.array(*topics)) + + df_filtered = df_with_topics.withColumn("topic", F.explode("topics")) + + else: + topic_exprs = [] + for userfilter in userfilters: logger.debug("Apply user-defined filter %s", userfilter) - df_tmp = apply_user_defined_filter(df, userfilter, _LOG) - - # The topic name is the filter name - topicname = args.substream_prefix + userfilter.split(".")[-1] + "_ztf" - - # FIXME: shouldn't we collect in a list the disquery? - disquery = push_to_kafka( - df_tmp, - topicname, - cnames, - checkpointpath_kafka, - args.tinterval, - kafka_cfg, - npart=None, - ) + + # build filter function dynamically + filter_func, colnames = expand_function_from_string(df, userfilter) + + topicname = args.substream_prefix + userfilter.split(".")[-1] + "_ztf" + + topic_exprs.append(F.when(filter_func(*colnames), F.lit(topicname))) + # array_compact for delete NULL in array + df_with_topics = df.withColumn("topics", F.array_compact(F.array(*topic_exprs))) + + df_filtered = df_with_topics.withColumn("topic", F.explode("topics")) + + # push to kafka (df_filtred) with One writeStream, using the column topic (not .option("topic",...) + disquery1 = push_to_kafka( + df_filtered, + cnames, + checkpointpath_kafka + "/{}filters_ztf".format(args.substream_prefix), + args.tinterval, + kafka_cfg, + npart=10, + ) # Special filter to count alerts topicname = "fink_ztf_{}".format(args.night) - disquery = push_to_kafka( + df = df.withColumn("topic", F.lit(topicname)) + disquery2 = push_to_kafka( df, - topicname, ["objectId"], - checkpointpath_kafka, + checkpointpath_kafka + "/" + topicname, args.tinterval, kafka_cfg, ) @@ -204,7 +219,8 @@ def main(): remaining_time = remaining_time if remaining_time > 0 else 0 logger.debug("Keep the Streaming for %s seconds", remaining_time) time.sleep(remaining_time) - disquery.stop() + disquery1.stop() + disquery2.stop() logger.info("Exiting the distribute service normally...") else: logger.debug("Wait for the end of queries") diff --git a/fink_broker/common/distribution_utils.py b/fink_broker/common/distribution_utils.py index 1b6964f5..764d53c8 100644 --- a/fink_broker/common/distribution_utils.py +++ b/fink_broker/common/distribution_utils.py @@ -1,5 +1,5 @@ # Copyright 2019-2026 AstroLab Software -# Author: Abhishek Chauhan, Julien Peloton +# Author: Abhishek Chauhan, Julien Peloton, Massinissa MACHTER # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -69,7 +69,9 @@ def get_kafka_df(df: DataFrame, key: str, elasticc: bool = False) -> DataFrame: # Create a StructType column in the df for distribution. # The contents and schema of the df can change over time - df_struct = df.select(struct(df.columns).alias("struct")) + df_struct = df.select( + struct(*[c for c in df.columns if c != "topic"]).alias("struct"), "topic" + ) # Convert into avro and save the schema if elasticc: @@ -79,9 +81,11 @@ def get_kafka_df(df: DataFrame, key: str, elasticc: bool = False) -> DataFrame: "/home/julien.peloton/elasticc/alert_schema/elasticc.v0_9.brokerClassification.avsc", "r", ).read() - df_kafka = df_struct.select(to_avro_native("struct", jsonschema).alias("value")) + df_kafka = df_struct.select( + to_avro_native("struct", jsonschema).alias("value"), "topic" + ) else: - df_kafka = df_struct.select(to_avro("struct").alias("value")) + df_kafka = df_struct.select(to_avro("struct").alias("value"), "topic") # Add a key based on schema versions df_kafka = df_kafka.withColumn("key", lit(key)) @@ -90,7 +94,7 @@ def get_kafka_df(df: DataFrame, key: str, elasticc: bool = False) -> DataFrame: def push_to_kafka( - df_in, topicname, cnames, checkpointpath_kafka, tinterval, kafka_cfg, npart=None + df_in, cnames, checkpointpath_kafka, tinterval, kafka_cfg, npart=None ): """Push data to a Kafka custer @@ -98,8 +102,6 @@ def push_to_kafka( ---------- df_in: Spark DataFrame Alert DataFrame - topicname: str - Name of the Kafka topic to create cnames: list of str List of columns to transfer in the stream checkpointpath_kafka: str @@ -115,7 +117,7 @@ def push_to_kafka( ------- out: Streaming DataFrame """ - df_in = df_in.selectExpr(cnames) + df_in = df_in.selectExpr(cnames + ["topic"]) # get schema from the streaming dataframe to # avoid non-nullable bug #852 @@ -131,8 +133,7 @@ def push_to_kafka( disquery = ( df_kafka.writeStream.format("kafka") .options(**kafka_cfg) - .option("topic", topicname) - .option("checkpointLocation", checkpointpath_kafka + "/" + topicname) + .option("checkpointLocation", checkpointpath_kafka) .trigger(processingTime="{} seconds".format(tinterval)) .start() )