Optimize distribute pipeline

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,77 +117,109 @@ 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!")
         remaining_time = args.exit_after
         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")

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")

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,16 +94,14 @@ 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
 
     Parameters
     ----------
     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()
     )