feat: Add support for window_by TVF window semantics

[zhenzhongxu]

What happened?

I have a simple example in this notebook showing that the window semantics are not supported by the compiler yet. Is this something we can add relatively easily?

https://github.com/zhenzhongxu/composable-data-architecture/blob/3c5324466c4355cda98d2443fc1f01b011a5241b/ibis-substrait-example.ipynb

import ibis
person_schema = ibis.schema({
        "id": "int64",
        "name": "string",
        "emailaddress": "string",
        "creditcard": "string",
        "city": "string",
        "state": "string"
    }
)

auction_schema = ibis.schema({
        "id": "int64",
        "itemname": "string",
        "description": "string",
        "initialbid": "float64",
        "reserve": "float64",
        "expires": "timestamp(3)",
        "seller": "int64",
    }
)

bid_schema = ibis.schema({
        "auction": "int64",
        "bidder": "int64",
        "price": "float64",
        "datetime": "timestamp(3)"
    }
)

import ibis
from ibis import _

# Create source tables
person_table = ibis.table(name="Person", schema=person_schema)
auction_table = ibis.table(name="Auction", schema=auction_schema)
bid_table = ibis.table(name="Bid", schema=bid_schema)

# Create SubstraitCompiler object
from ibis_substrait.compiler.core import SubstraitCompiler
compiler = SubstraitCompiler()

# over window
over_window = bid_table.filter(_ is not None)[_.price.mean().over(range=(-ibis.interval(seconds=10), 0), order_by=_.datetime).name("avg_price")]
compiler.compile(over_window)

# hop window
hop_window = bid_table.filter(_ is not None).window_by(time_col=_.datetime).hop(size=ibis.interval(seconds=10), slide=ibis.interval(seconds=2)).agg(_.price.mean().name("avg_price"))

compiler.compile(hop_window)

What version of ibis-substrait are you using?

v.4.0.1

What substrait consumer(s) are you using, if any?

N/A

Relevant log output

---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
Cell In[47], line 14
     10 # hop window
     11 hop_window = bid_table.filter(_ is not None).window_by(time_col=_.datetime).hop(size=ibis.interval(seconds=10), slide=ibis.interval(seconds=2)).agg(_.price.mean().name("avg_price"))
---> 14 compiler.compile(hop_window)

File ~/miniconda3/envs/ibis-dev-arm64/lib/python3.12/site-packages/ibis_substrait/compiler/core.py:222, in SubstraitCompiler.compile(self, expr, **kwargs)
    217 from .translate import translate
    219 expr_schema = expr.schema()
    220 rel = stp.PlanRel(
    221     root=stalg.RelRoot(
--> 222         input=translate(expr.op(), compiler=self, **kwargs),
    223         names=translate(expr_schema).names,
    224     )
    225 )
    226 ver = vparse(__substrait_version__)
    227 return stp.Plan(
    228     version=stp.Version(
    229         major_number=ver.major,
   (...)
    256     relations=[rel],
    257 )

File ~/miniconda3/envs/ibis-dev-arm64/lib/python3.12/functools.py:907, in singledispatch.<locals>.wrapper(*args, **kw)
    903 if not args:
    904     raise TypeError(f'{funcname} requires at least '
    905                     '1 positional argument')
--> 907 return dispatch(args[0].__class__)(*args, **kw)

File ~/miniconda3/envs/ibis-dev-arm64/lib/python3.12/site-packages/ibis_substrait/compiler/translate.py:55, in translate(*args, **kwargs)
     53 @functools.singledispatch
     54 def translate(*args: Any, **kwargs: Any) -> Any:
---> 55     raise NotImplementedError(*args)

NotImplementedError: <ibis.expr.operations.temporal_windows.WindowAggregate object at 0x16d904dd0>

[gforsyth]

Hey @zhenzhongxu -- we can add a target for the window_by TVF stuff, but I don't know how to spell those style functions in Substrait -- do you have a reference for what the target should look like?

We already support batch window functions.

[zhenzhongxu]

@gforsyth does this help, or we can also try to look at the Calcite plan on Flink.

[zhenzhongxu]

Also add @richtia @EpsilonPrime for guidance on the substrait side.

[gforsyth]

I'm reasonably familiar with the flink-style TVF windowing, but I'm not clear on whether there's a way to express that in Substrait -- so yeah, if @EpsilonPrime or @richtia have any pointers.

If there isn't currently a way to express that in Substrait, then the first step will be adding that to the upstream. Then it should be straightforward to target from here.

[EpsilonPrime]

{
  "extensionUris": [
    {
      "extensionUriAnchor": 1,
      "uri": "/functions_arithmetic.yaml"
    }
  ],
  "extensions": [
    {
      "extensionFunction": {
        "extensionUriReference": 1,
        "functionAnchor": 1,
        "name": "row_number"
      }
    }
  ],
  "relations": [
    {
      "root": {
        "input": {
          "project": {
            "common": {
              "emit": {
                "outputMapping": [
                  3,
                  4,
                  5,
                  6
                ]
              }
            },
            "input": {
              "read": {
                "common": {
                  "direct": {}
                },
                "baseSchema": {
                  "names": [
                    "user_id",
                    "name",
                    "paid_for_service"
                  ],
                  "struct": {
                    "types": [
                      {
                        "string": {
                          "nullability": "NULLABILITY_REQUIRED"
                        }
                      },
                      {
                        "string": {
                          "nullability": "NULLABILITY_REQUIRED"
                        }
                      },
                      {
                        "bool": {
                          "nullability": "NULLABILITY_REQUIRED"
                        }
                      }
                    ],
                    "nullability": "NULLABILITY_REQUIRED"
                  }
                },
                "namedTable": {
                  "names": [
                    "users"
                  ]
                }
              }
            },
            "expressions": [
              {
                "selection": {
                  "directReference": {
                    "structField": {}
                  },
                  "rootReference": {}
                }
              },
              {
                "selection": {
                  "directReference": {
                    "structField": {
                      "field": 1
                    }
                  },
                  "rootReference": {}
                }
              },
              {
                "selection": {
                  "directReference": {
                    "structField": {
                      "field": 2
                    }
                  },
                  "rootReference": {}
                }
              },
              {
                "windowFunction": {
                  "functionReference": 1,
                  "sorts": [
                    {
                      "expr": {
                        "selection": {
                          "directReference": {
                            "structField": {
                              "field": 1
                            }
                          },
                          "rootReference": {}
                        }
                      },
                      "direction": "SORT_DIRECTION_ASC_NULLS_FIRST"
                    }
                  ],
                  "upperBound": {
                    "unbounded": {}
                  },
                  "lowerBound": {
                    "unbounded": {}
                  },
                  "outputType": {
                    "i64": {
                      "nullability": "NULLABILITY_REQUIRED"
                    }
                  },
                  "invocation": 3
                }
              }
            ]
          }
        },
        "names": [
          "user_id",
          "name",
          "paid_for_service",
          "row_number"
        ]
      }
    }
  ],
  "version": {
    "minorNumber": 52,
    "producer": "spark-substrait-gateway"
  }
}

[gforsyth]

Thanks, @EpsilonPrime -- I think that's a standard window function, which is a different beast than the table-valued functions used by some of the streaming engines.

I don't know that I saw any chatter on the substrait issue tracker about anyone implementing this.

We can probably desugar the hop and tumble operations on the Ibis side, but some of the other TVF windows are harder to write in batch form

[EpsilonPrime]

The only other way I know of to do window related behavior is with the ConsistentPartitionWindowRel -- https://github.com/substrait-io/substrait/blob/main/proto/substrait/algebra.proto#L272.

[zhenzhongxu]

Let me grab the Calcite plan for the window TVF so we can compare the standard window vs streaming window. Will circle back shortly.