Incorrect results / failure from query requiring multiple scans of arrow stream

[jonathanswenson]

What happens?

Undefined, undesirable behavior running query over arrow stream that requires multiple passes.

Perhaps this issue needs to be in the duckdb-java or duckdb-python repos, but I think the behavior is likely stemming from a problem in duckdb core.

If, after first registering an arrow stream with duckdb, you run a query that requires mutiple table scans, the query will fail (java) or produce incorrect results (python).

In duckdb-java, it appears as though the second scan of the results fails with an error indicating that the stream has been released: Invalid Input Error: This stream has been released

In duckdb-python, the query completes successfully, but returns seemingly incorrect data.

Ideally, the query would complete, successfully, with the correct result, having materialized the intermediate result as necessary (and apply any predicate / filter pushdown if posible).

To Reproduce

Java repro that throws an error

package com.acme;

import org.apache.arrow.c.ArrowArrayStream;
import org.apache.arrow.c.Data;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VarCharVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.ipc.ArrowStreamReader;
import org.apache.arrow.vector.ipc.ArrowStreamWriter;
import org.duckdb.DuckDBConnection;
import org.duckdb.DuckDBDriver;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.List;
import java.util.Properties;

public class DuckDBStreamIngestTest {

    private static byte[] createStream(BufferAllocator allocator) throws Exception {
        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();

        IntVector intVector = new IntVector("id", allocator);
        VarCharVector stringVector = new VarCharVector("value", allocator);

        try (
                VectorSchemaRoot vsr = new VectorSchemaRoot(List.of(intVector, stringVector));
                ArrowStreamWriter writer = new ArrowStreamWriter(vsr, null, outputStream)
        ) {
            vsr.setRowCount(5);
            for (int i = 0; i < 5; i++) {
                intVector.setSafe(i, i);
                stringVector.setSafe(i, ("v " + Integer.valueOf(i).toString()).getBytes(StandardCharsets.UTF_8));
            }
            writer.writeBatch();
        }

        return outputStream.toByteArray();
    }

    public static void main(final String[] args) throws Exception {
        BufferAllocator allocator = new RootAllocator();
        byte[] bytes = createStream(allocator);

        ByteArrayInputStream inputStream = new ByteArrayInputStream(bytes);

        ArrowStreamReader arrowReader = new ArrowStreamReader(inputStream, allocator);
        ArrowArrayStream arrowArrayStream = ArrowArrayStream.allocateNew(allocator);
        Data.exportArrayStream(allocator, arrowReader, arrowArrayStream);
        DuckDBDriver driver = new DuckDBDriver();
        try (Connection connection = driver.connect("jdbc:duckdb:", new Properties())) {
            DuckDBConnection conn = connection.unwrap(DuckDBConnection.class);
            conn.registerArrowStream("arrow_table", arrowArrayStream);

            try (
                    Statement statement = connection.createStatement();
                    ResultSet resultSet = statement.executeQuery("select id from arrow_table union all select id + 1 from arrow_table");
            ) {
                printResultSet(resultSet);
            }

        }
    }

    private static void printResultSet(ResultSet resultSet) throws SQLException {
        for (int i = 1; i <= resultSet.getMetaData().getColumnCount(); i++) {
            System.out.print(resultSet.getMetaData().getColumnLabel(i) + ", ");
        }
        System.out.println();
        while (resultSet.next()) {
            resultSet.getMetaData().getColumnCount();
            for (int i = 1; i <= resultSet.getMetaData().getColumnCount(); i++) {
                System.out.print(resultSet.getString(i) + ", ");
            }
            System.out.println();
        }
    }
}

Java output:

Exception in thread "main" java.sql.SQLException: Invalid Input Error: This stream has been released
	at org.duckdb.DuckDBNative.duckdb_jdbc_execute(Native Method)
	at org.duckdb.DuckDBPreparedStatement.execute(DuckDBPreparedStatement.java:193)
	at org.duckdb.DuckDBPreparedStatement.execute(DuckDBPreparedStatement.java:159)
	at org.duckdb.DuckDBPreparedStatement.executeQuery(DuckDBPreparedStatement.java:229)
	at org.duckdb.DuckDBPreparedStatement.executeQuery(DuckDBPreparedStatement.java:263)
	at com.acme.DuckDBStreamIngestTest.main(DuckDBStreamIngestTest.java:63)

Python repro

import pyarrow as pa
import duckdb
import io


def create_arrow_stream(table):
    buffer = io.BytesIO()

    with pa.ipc.new_stream(buffer, table.schema) as writer:
        writer.write(table)

    buffer.seek(0)
    return buffer


def main():
    data = {
        'id': [1, 2, 3, 4, 5],
        'value': ['one', 'two', 'three', 'four', 'five']
    }

    table = pa.table(data)

    stream_buffer1 = create_arrow_stream(table)

    with pa.ipc.open_stream(stream_buffer1) as stream1:
        duckdb.register("arrow_stream", stream1)

        sql = "SELECT id FROM arrow_stream union all select id + 1 from arrow_stream"

        print("Query Results:")
        print(duckdb.sql(sql).show())

if __name__ == "__main__":
    main()

Python output:

Query Results:
┌───────┐
│  id   │
│ int64 │
├───────┤
│     1 │
│     2 │
│     3 │
│     4 │
│     5 │
└───────┘

In both cases, you can materialize the table (issuing a CREATE TABLE <tablename> AS SELECT * FROM arrow_stream) and the query completes and produces the appropriate result.

id, 
0, 
1, 
2, 
3, 
4, 
1, 
2, 
3, 
4, 
5, 

OS:

linux x86_64, macOS arch64

DuckDB Version:

1.3.2

DuckDB Client:

duckdb-java,python

Hardware:

No response

Full Name:

Jonathan Swenson

Affiliation:

Omni

What is the latest build you tested with? If possible, we recommend testing with the latest nightly build.

I have tested with a stable release

Did you include all relevant data sets for reproducing the issue?

Yes

Did you include all code required to reproduce the issue?

• Yes, I have

Did you include all relevant configuration (e.g., CPU architecture, Python version, Linux distribution) to reproduce the issue?

• Yes, I have

[Tishj]

Streams really mess with the execution model of DuckDB, in this case the underlying object is "destroyed" while scanning.
Table scans expect to be able to scan a source without that scan modifying external state.

Which doesnt hold up here.

Perhaps we can try to utilize the same system that facilitates materialized CTEs, where one scan populates the cache and the other scans become a read from that cache instead.

That does require detecting that these are destructive scans though, which has proved difficult in the past, when I looked into this problem in the python client.

[Mytherin]

Given that arrow streams cannot be scanned multiple times, there are really only two possible options:

• Throw an error when a stream is scanned twice
• Materialize the stream

Given that materializing the stream up-front nullifies many of the benefits of it being a stream in the first place, I think this is not desirable to have on by default. The ideal situation here in my view is to throw a descriptive error, and have the user perform a materialize manually if that is the desired behavior.

[Tishj]

Agreed, I have attempted this before, but could not manage to correctly detect the cases where a stream was scanned twice in the same query.

And also, the problem stretches further than just a single query, if the same query is reran it'll produce an empty result, because the input stream is already destroyed.

[Mytherin]

When register is called we create a view over the stream, perhaps we can wrap that view in an object that detects when a scan is initiated over it? Then it can throw an error if a second scan is initiated?

[Tishj]

This issue should probably live in duckdb-java and duckdb-python, because this can definitely not be fixed in a client-agnostic way

[Mytherin]

Yes, I think Java actually has the desired behavior here - so we can move this to duckdb-python.

[jonathanswenson]

Perhaps a little bit of extra context here would be helpful. I agree that this is a "bug" in duckdb python that the incorrect results are displayed without error, however the error functionality is somewhat limiting to our usage of duckdb (in particular duckdb-java).

We're currently having to materialize EVERY table into duckdb java because it isn't trivially clear whether or not a particular query we're going to execute will require multiple passes over the result set. In the vast majority of cases we're doing a single pass and a stream would work (and it we'd love to be able to take advantage of functionality like predicate / filter pushdown over the stream as well -- however that appears not to work in duckdb-java either -- see: duckdb/duckdb-java#392.

Perhaps the need is better phrased as a feature request to add additional functionality to duckdb (core) to allow for what @Tishj describe and materialize the stream if necessary (perhaps on disk if a temp directory is required) -- similar to CTE materialization.

While we may be able to take advantage of our own relational layer to decide if we expect a multi-pass query is required, we may also try to issue an explain to determine if the duckdb planner expects a multi-pass table scan, and, if so, materialize it explicitly.

I expect it is also complex to do so, but in the event of a multi-pass query where the passes over the stream are different (say different projection / different applied filters), it would also be interesting to apply the appropriate predicate / filter pushdowns (perhaps the union of predicates / filters to the materialized table).

I also attempted to accomplish this by manually constructing a CTE:

with materialized_table as (select * from stream)

select ... from materialized_table t1 join materialized_table t2 on ...

However, this also fails with the same This stream has been released -- Is this expected (given the disccusion above about materialization of CTEs?) a limitation of CTEs over streams? or just so happens that the planner would rather do thepass over the results twice rather than materialize the CTE (which would then have to be passed over twice)?

[Mytherin]

Perhaps try explicitly materializing the CTE? e.g.:

WITH t(x) AS MATERIALIZED (complex_query)
SELECT *
FROM
    t AS t1,
    t AS t2,
    t AS t3;

By default CTEs are not materialized in DuckDB v1.3.2 and below. In DuckDB v1.4.0 they should be automatically materialized if used multiple times in the same query. See also the documentation for more info.

[paultiq]

Instead of eagerly creating the stream in python_replacement_scan, duckdb-python could store the PyCapsule interface object and defer __arrow_c_stream__() to the point of use. You can create multiple streams from the same object, you just can't reuse the same stream capsule.

The problem is that the stream capsule is created once in CreateArrowScan, even though the scan may be consumed multiple times:

duckdb-python/src/duckdb_py/python_replacement_scan.cpp

Lines 55 to 58 in bbd4389

	if (type == PyArrowObjectType::PyCapsuleInterface) {
	entry = entry.attr("__arrow_c_stream__")();
	type = PyArrowObjectType::PyCapsule;
	}

Deferring the __arrow_c_stream__() call to the points of consumption, namely PythonTableArrowArrayStreamFactory::Produce and PythonTableArrowArrayStreamFactory::GetSchema, solves this issue by creating "fresh" stream capsules for each access.

* I have a working branch, can open a PR if interested. It's significantly faster (25%) for scanning Polars dataframes
** Even with separate calls to __arrow_c_stream for each use in GetSchema and Product, it's still faster
*** There are a few TODOs, such as handling __arrow_c_schema__() for schema-only queries