DataStream.transform_sql

This is a SQL version of the transform method. It allows you to write SQL expressions that can be applied to each batch in the DataStream. This is the X in select X from Y. The Y is the DataStream. The X can be any SQL expression, including aggregation functions.

Example sql expression

" sum(l_quantity) as sum_qty, sum(l_extendedprice) as sum_base_price, sum(l_extendedprice * (1 - l_discount)) as sum_disc_price, sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge"

You must supply an alias for each transformation. Any syntax that is supported by DuckDB can be used. You can optionally also specify groupby columns. This will apply the SQL expression select X from Y group by Z to each batch in the DataStream.

Parameters:

Name	Type	Description	Default
sql_expression	str	The SQL expression to apply to each batch in the DataStream.	required
groupby	list	The list of columns to group by.	[]
foldable	bool	Whether or not the transformation can be executed as part of the batch post-processing of the previous operation in the execution graph. This is set to True by default. Probably should be True.	True

Source code in pyquokka/datastream.py

def transform_sql(self, sql_expression, groupby = [], foldable = True):

    """

    This is a SQL version of the `transform` method. It allows you to write SQL expressions that can be applied to each batch in the DataStream.
    This is the X in `select X from Y`. The Y is the DataStream. The X can be any SQL expression, including aggregation functions.
    Example sql expression:
        "
        sum(l_quantity) as sum_qty,
        sum(l_extendedprice) as sum_base_price,
        sum(l_extendedprice * (1 - l_discount)) as sum_disc_price,
        sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge"

    You must supply an alias for each transformation. Any syntax that is supported by DuckDB can be used. 
    You can optionally also specify groupby columns. This will apply the SQL expression `select X from Y group by Z` to each batch in the DataStream.

    Args:
        sql_expression (str): The SQL expression to apply to each batch in the DataStream.
        groupby (list): The list of columns to group by.
        foldable (bool): Whether or not the transformation can be executed as part of the batch post-processing of the previous operation in the
            execution graph. This is set to True by default. Probably should be True.

    """

    assert type(groupby) == list

    enhanced_exp = "select " +",".join(groupby) + ", " + sql_expression + " from batch_arrow"
    if len(groupby) > 0:
        enhanced_exp = enhanced_exp + " group by " + ",".join(groupby)

    enhanced_exp = label_sample_table_names(sqlglot.parse_one(enhanced_exp), 'batch_arrow').sql()

    sqlglot_node = sqlglot.parse_one(enhanced_exp)
    required_columns = required_columns_from_exp(sqlglot_node)
    if len(required_columns) == 0:
        # TODO: this is to make sure count works by picking a random column to download. 
        required_columns = {self.schema[1]}

    for col in required_columns:
        assert col in self.schema, "required column %s not in schema" % col

    new_columns = [i.alias for i in sqlglot_node.selects if i.name not in groupby]
    assert '' not in new_columns, "must provide alias for each computation"

    assert type(required_columns) == set
    for column in new_columns:
        assert column not in self.schema, "For now new columns cannot have same names as existing columns"

    if self.materialized:
        batch_arrow = self._get_materialized_df().to_arrow()
        con = duckdb.connect().execute('PRAGMA threads=%d' % multiprocessing.cpu_count())
        df = polars.from_arrow(con.execute(enhanced_exp).arrow())
        return self.quokka_context.from_polars(df)

    def duckdb_func(func, batch):
        batch_arrow = batch.to_arrow()
        for i, (col_name, type_) in enumerate(zip(batch_arrow.schema.names, batch_arrow.schema.types)):
            if pa.types.is_boolean(type_):
                batch_arrow = batch_arrow.set_column(i, col_name, compute.cast(batch_arrow.column(col_name), pa.int32()))
        con = duckdb.connect().execute('PRAGMA threads=%d' % multiprocessing.cpu_count())
        return polars.from_arrow(con.execute(func).arrow())

    return self.quokka_context.new_stream(
        sources={0: self},
        partitioners={0: PassThroughPartitioner()},
        node=MapNode(
            schema=groupby + new_columns,
            schema_mapping={
                **{new_column: {-1: new_column} for new_column in new_columns}, **{col: {0: col} for col in self.schema}},
            required_columns={0: required_columns},
            function=partial(duckdb_func, enhanced_exp),
            foldable=foldable),
        schema= groupby + new_columns,
        sorted = self.sorted
        )